# -*- coding: utf-8 -*-
"""Astraea — the world's best Washington divorce-law attorney.
Multi-agent RAG + no-KYC profiles (About Me) + document prep + comms missions + TTS.
"""
import json
import re
import time
import urllib.request

from flask import Flask, jsonify, render_template, request, send_file, Response

import rag
import store
import comms
import documents
import tts
from agents import AGENTS, AGENT_BY_ID, OLLAMA_URL, RAG_MODEL, GENERAL_MODEL

app = Flask(__name__)

# ── Professional avatars (gradient + glyph, rendered client-side) ──
AVATARS = [
    {"id": 0, "glyph": "\u2696", "c1": "#6ea8ff", "c2": "#a78bfa", "label": "Scales"},
    {"id": 1, "glyph": "\U0001F3DB", "c1": "#7dd3fc", "c2": "#0ea5e9", "label": "Court"},
    {"id": 2, "glyph": "\U0001F6E1", "c1": "#4ade80", "c2": "#059669", "label": "Shield"},
    {"id": 3, "glyph": "\U0001F4DC", "c1": "#f5d06f", "c2": "#d97706", "label": "Scroll"},
    {"id": 4, "glyph": "\U0001F54A", "c1": "#f0abfc", "c2": "#c026d3", "label": "Dove"},
    {"id": 5, "glyph": "\U0001F4BC", "c1": "#fdba74", "c2": "#ea580c", "label": "Briefcase"},
    {"id": 6, "glyph": "\U0001F512", "c1": "#67e8f9", "c2": "#0891b2", "label": "Key"},
    {"id": 7, "glyph": "\u2731", "c1": "#fda4af", "c2": "#e11d48", "label": "Star"},
    {"id": 8, "glyph": "\U0001F3AF", "c1": "#a3e635", "c2": "#65a30d", "label": "Target"},
    {"id": 9, "glyph": "\U0001F525", "c1": "#f87171", "c2": "#b91c1c", "label": "Flame"},
    {"id": 10, "glyph": "\U0001F31F", "c1": "#fbbf24", "c2": "#b45309", "label": "Glow"},
    {"id": 11, "glyph": "\U0001F9ED", "c1": "#5eead4", "c2": "#0d9488", "label": "Compass"},
]


# ── Ollama helpers ──
def _ollama_chat(model, messages, tools=None, num_ctx=16384, num_predict=1400,
                 temperature=0.2, timeout=180):
    payload = {"model": model, "messages": messages, "stream": False, "think": False,
               "options": {"temperature": temperature, "num_predict": num_predict, "num_ctx": num_ctx}}
    if tools:
        payload["tools"] = tools
    req = urllib.request.Request(f"{OLLAMA_URL}/api/chat", data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    op = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    with op.open(req, timeout=timeout) as r:
        return json.loads(r.read().decode("utf-8"))


def _strip_sources(text):
    text = re.split(r"\n\s*(?:Sources|References|Citations)\s*:\s*\n", text, flags=re.I)[0]
    return text.strip()


# ── About Me injection (designated plainly, used silently) ──
def about_me_block(about_me):
    if not about_me or not about_me.strip():
        return ""
    return (
        "\n\n=== ABOUT THE USER (PRIVATE CONTEXT) ===\n"
        f"{about_me.strip()}\n"
        "=== END ABOUT THE USER ===\n"
        "This is private background about the user's situation. Use it silently to tailor your "
        "answer to their facts. NEVER reference it, quote it, summarize it back, or say anything "
        "like 'based on what you shared'. Answer ONLY what the user asks."
    )


# ── Comms tools (agent can invoke to send SMS/email/call) ──
TOOLS = [
    {"type": "function", "function": {
        "name": "send_sms", "description": "Send a professionally-worded text message on the user's behalf.",
        "parameters": {"type": "object", "properties": {
            "recipient": {"type": "string", "description": "recipient phone number in E.164 (e.g. +14255550123)"},
            "message": {"type": "string", "description": "the message text"}},
            "required": ["recipient", "message"]}}},
    {"type": "function", "function": {
        "name": "send_email", "description": "Send a formal email on the user's behalf.",
        "parameters": {"type": "object", "properties": {
            "recipient": {"type": "string", "description": "recipient email address"},
            "subject": {"type": "string"}, "body": {"type": "string"}},
            "required": ["recipient", "subject", "body"]}}},
    {"type": "function", "function": {
        "name": "make_call", "description": "Place a phone call that speaks a message.",
        "parameters": {"type": "object", "properties": {
            "recipient": {"type": "string", "description": "recipient phone number in E.164"},
            "message": {"type": "string", "description": "the message to speak"}},
            "required": ["recipient", "message"]}}},
]


def _execute_tool(name, args, settings, uid, display_name):
    recipient = args.get("recipient", "")
    if comms.no_contact_blocked(settings, recipient):
        return ("BLOCKED: a no-contact order is in effect for this person. No message was sent. "
                "Inform the user you cannot contact this person due to the no-contact order.")
    try:
        if name == "send_sms":
            comms.send_sms(settings, recipient, comms.sms_template(args.get("message", "")))
            return f"SMS sent to {recipient}."
        if name == "send_email":
            body = comms.email_template(display_name, recipient, args.get("subject", ""), args.get("body", ""))
            comms.send_email(settings, recipient, args.get("subject", ""), body)
            return f"Email sent to {recipient}."
        if name == "make_call":
            twiml_url = request.host_url.rstrip("/") + "/twilio/voice"
            comms.make_call(settings, recipient, twiml_url)
            return f"Call initiated to {recipient}."
    except Exception as e:
        return f"ERROR sending {name}: {e}"
    return "Unknown tool."


# ── RAG chat with tool-calling ──
def safety_block(settings):
    if not (settings.get("domestic_violence") or settings.get("protection_order")):
        return ""
    lines = ["\n\n=== SAFETY CONTEXT (PRIVATE — CRITICAL) ==="]
    if settings.get("domestic_violence"):
        lines.append("DOMESTIC VIOLENCE IS A FACTOR IN THIS CASE.")
    if settings.get("protection_order"):
        lines.append("An active protection order is in effect (RCW 7.105).")
    notes = (settings.get("safety_notes") or "").strip()
    if notes:
        lines.append(f"Safety notes: {notes}")
    lines.append("=== END SAFETY CONTEXT ===")
    lines.append(
        "Prioritize the user's safety above all else. Apply RCW 26.09.191 restrictions on "
        "residential time and decision-making where domestic violence or abuse is a factor. "
        "Recommend protection orders, supervised visitation, and safety planning where "
        "appropriate. NEVER advise or facilitate direct contact with the abusive or restrained party."
    )
    return "\n".join(lines)


def _build_answer(agent, message, history, about_me, settings, uid, display_name):
    top = rag.retrieve(message, agent["id"], top_k=5)
    context_blocks = []
    citations = []
    seen = set()
    for i, (score, c) in enumerate(top, 1):
        key = (c["source"], c["title"])
        if key in seen:
            continue
        seen.add(key)
        context_blocks.append(f"[{i}] ({c['source']} — {c['title']})\n{c['text']}")
        citations.append({"source": c["source"], "title": c["title"],
                          "score": round(score, 3), "snippet": c["text"][:260]})

    system = (
        f"{agent['system']}\n\n"
        "Rules:\n"
        "- Give REAL, practical legal advice: state the rule, apply it to the user's facts, "
        "tell them what to do next.\n"
        "- Answer DIRECTLY. Do not restate the question or narrate your reasoning.\n"
        "- Ground every legal claim in the reference documents. Cite the RCW section AND the "
        "controlling case law by name.\n"
        "- Be plain-English, specific to Washington State.\n"
        "- You may use the available tools (send_sms / send_email / make_call) ONLY when the "
        "user explicitly asks you to contact someone on their behalf.\n"
        + about_me_block(about_me)
        + safety_block(settings)
    )

    user = ("REFERENCE DOCUMENTS:\n\n" + "\n\n".join(context_blocks)
            + f"\n\nUSER QUESTION: {message}")

    messages = [{"role": "system", "content": system}]
    if history:
        for turn in history[-6:]:
            if turn.get("role") in ("user", "assistant"):
                messages.append({"role": turn["role"], "content": turn["content"]})
    messages.append({"role": "user", "content": user})

    try:
        resp = _ollama_chat(RAG_MODEL, messages, tools=TOOLS)
    except Exception:
        try:
            resp = _ollama_chat(GENERAL_MODEL, messages)
        except Exception:
            return {"answer": "The legal engine is temporarily unavailable. Try again shortly.",
                    "citations": [], "grounded": False}

    # tool-calling loop
    for _ in range(3):
        msg = resp.get("message", {})
        tool_calls = msg.get("tool_calls") or []
        if not tool_calls:
            break
        messages.append(msg)
        for tc in tool_calls:
            fn = tc.get("function", {})
            name = fn.get("name", "")
            args = fn.get("arguments", {})
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except Exception:
                    args = {}
            result = _execute_tool(name, args, settings, uid, display_name)
            messages.append({"role": "tool", "content": result})
        resp = _ollama_chat(RAG_MODEL, messages, tools=TOOLS)

    final = resp.get("message", {})
    answer = (final.get("content") or final.get("thinking") or "").strip()
    answer = _strip_sources(answer)
    return {"answer": answer, "citations": citations, "grounded": bool(context_blocks)}


# ── Auth helpers ──
def _uid():
    token = request.headers.get("Authorization", "").replace("Bearer ", "")
    return store.user_from_token(token)


def _profile_or_none():
    uid = _uid()
    return uid, store.get_profile(uid) if uid else None, store.get_settings(uid) if uid else {}


# ── Pages ──
@app.route("/")
def landing():
    return render_template("landing.html", agents=AGENTS)


@app.route("/app")
def app_page():
    return render_template("index.html", agents=AGENTS)


# ── Auth API ──
@app.route("/api/register", methods=["POST"])
def api_register():
    d = request.get_json(force=True, silent=True) or {}
    u = (d.get("username") or "").strip().lower()
    p = (d.get("password") or "").strip()
    if len(u) < 2 or len(u) > 30 or not u.isalnum():
        return jsonify({"error": "Username must be 2-30 alphanumeric chars"}), 400
    if len(p) < 4:
        return jsonify({"error": "Password must be 4+ chars"}), 400
    c = store._conn()
    if c.execute("SELECT id FROM users WHERE username=?", (u,)).fetchone():
        c.close()
        return jsonify({"error": "Username taken"}), 409
    c.close()
    uid = store.create_user(u, p)
    return jsonify({"ok": True, "token": store.issue_token(uid), "username": u})


@app.route("/api/login", methods=["POST"])
def api_login():
    d = request.get_json(force=True, silent=True) or {}
    u = (d.get("username") or "").strip().lower()
    p = (d.get("password") or "").strip()
    uid = store.authenticate(u, p)
    if not uid:
        return jsonify({"error": "Invalid username or password"}), 401
    return jsonify({"ok": True, "token": store.issue_token(uid), "username": u})


@app.route("/api/me")
def api_me():
    uid = _uid()
    if not uid:
        return jsonify({"authenticated": False}), 401
    return jsonify({"authenticated": True, "profile": store.get_profile(uid),
                    "settings": store.get_settings(uid)})


@app.route("/api/profile", methods=["GET", "POST"])
def api_profile():
    uid = _uid()
    if not uid:
        return jsonify({"error": "auth required"}), 401
    if request.method == "GET":
        return jsonify(store.get_profile(uid))
    d = request.get_json(force=True, silent=True) or {}
    prof = store.set_profile(uid, display_name=d.get("display_name"),
                             avatar=d.get("avatar"), about_me=d.get("about_me"))
    return jsonify(prof)


@app.route("/api/settings", methods=["GET", "POST"])
def api_settings():
    uid = _uid()
    if not uid:
        return jsonify({"error": "auth required"}), 401
    if request.method == "GET":
        return jsonify(store.get_settings(uid))
    d = request.get_json(force=True, silent=True) or {}
    s = store.set_settings(uid, **d)
    # never echo the auth token back in full
    if "twilio_auth_token" in s:
        s["twilio_auth_token"] = "••••" if s["twilio_auth_token"] else ""
    if "smtp_pass" in s:
        s["smtp_pass"] = "••••" if s["smtp_pass"] else ""
    return jsonify(s)


@app.route("/api/avatars")
def api_avatars():
    return jsonify(AVATARS)


# ── Chat ──
@app.route("/api/chat", methods=["POST"])
def api_chat():
    uid = _uid()
    d = request.get_json(force=True, silent=True) or {}
    agent_id = (d.get("agent_id") or "navigator").strip()
    message = (d.get("message") or "").strip()
    client_history = d.get("history") or []
    agent = AGENT_BY_ID.get(agent_id)
    if not agent:
        return jsonify({"error": "Unknown agent"}), 400
    if not message:
        return jsonify({"error": "Message required"}), 400
    message = message[:4000]
    profile = store.get_profile(uid) if uid else {"about_me": "", "display_name": ""}
    settings = store.get_settings(uid) if uid else {}
    # persisted history for authed users; client history as fallback for anonymous
    history = store.get_conversation(uid, agent_id) if uid else client_history
    t0 = time.time()
    result = _build_answer(agent, message, history, profile.get("about_me", ""),
                           settings, uid, profile.get("display_name", ""))
    if uid and result.get("answer"):
        store.save_message(uid, agent_id, "user", message)
        store.save_message(uid, agent_id, "assistant", result["answer"], result.get("citations"))
    result["latency_ms"] = round((time.time() - t0) * 1000)
    result["agent"] = agent_id
    return jsonify(result)


@app.route("/api/conversation/<agent_id>", methods=["GET", "DELETE"])
def api_conversation(agent_id):
    uid = _uid()
    if not uid:
        return jsonify({"error": "auth required"}), 401
    if agent_id not in AGENT_BY_ID:
        return jsonify({"error": "unknown agent"}), 400
    if request.method == "DELETE":
        store.clear_conversation(uid, agent_id)
        return jsonify({"ok": True})
    return jsonify({"history": store.get_conversation(uid, agent_id)})


# ── Documents ──
@app.route("/api/documents/types")
def api_doc_types():
    return jsonify([{"id": k, "name": v["name"], "form": v["form"], "description": v["description"]}
                    for k, v in documents.DOC_TYPES.items()])


@app.route("/api/documents", methods=["POST"])
def api_documents():
    uid = _uid()
    d = request.get_json(force=True, silent=True) or {}
    doc_type = (d.get("doc_type") or "").strip()
    instructions = (d.get("instructions") or "").strip()
    if doc_type not in documents.DOC_TYPES:
        return jsonify({"error": "Unknown document type"}), 400
    profile = store.get_profile(uid) if uid else {"about_me": ""}
    # retrieve relevant context for grounding
    top = rag.retrieve(documents.DOC_TYPES[doc_type]["name"], "divorce", top_k=4)
    ctx = "\n\n".join(f"{c['text']}" for s, c in top)
    try:
        text = documents.generate(doc_type, profile.get("about_me", ""), instructions, ctx)
    except Exception as e:
        return jsonify({"error": f"generation failed: {e}"}), 500
    return jsonify({"ok": True, "document": text,
                    "name": documents.DOC_TYPES[doc_type]["name"]})


# ── TTS ──
@app.route("/api/voices")
def api_voices():
    return jsonify({"voices": tts.list_voices(), "default": "aria"})


@app.route("/api/speak", methods=["POST"])
def api_speak():
    d = request.get_json(force=True, silent=True) or {}
    text = (d.get("text") or "").strip()
    voice = d.get("voice") or "aria"
    if not text:
        return jsonify({"error": "no text"}), 400
    try:
        path = tts.synthesize(text, voice)
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    return send_file(path, mimetype="audio/mpeg")


# ── Missions (manual comms) ──
@app.route("/api/mission", methods=["POST"])
def api_mission():
    uid = _uid()
    if not uid:
        return jsonify({"error": "auth required"}), 401
    d = request.get_json(force=True, silent=True) or {}
    mtype = d.get("type")
    recipient = d.get("recipient") or ""
    settings = store.get_settings(uid)
    profile = store.get_profile(uid)
    if comms.no_contact_blocked(settings, recipient):
        return jsonify({"ok": False, "error": "Blocked: a no-contact order is in effect for this person."}), 403
    try:
        if mtype == "sms":
            comms.send_sms(settings, recipient, comms.sms_template(d.get("message", "")))
        elif mtype == "email":
            body = comms.email_template(profile.get("display_name", ""), recipient,
                                        d.get("subject", ""), d.get("body", ""))
            comms.send_email(settings, recipient, d.get("subject", ""), body)
        elif mtype == "call":
            comms.make_call(settings, recipient, request.host_url.rstrip("/") + "/twilio/voice")
        else:
            return jsonify({"ok": False, "error": "unknown type"}), 400
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)}), 500
    return jsonify({"ok": True})


# ── Twilio voice webhook (TwiML <Say>) ──
@app.route("/twilio/voice", methods=["GET", "POST"])
def twilio_voice():
    msg = request.values.get("message", "This is an automated message from a legal assistant.")
    twiml = (f'<?xml version="1.0" encoding="UTF-8"?>'
             f'<Response><Say voice="Polly.Joanna">{msg}</Say></Response>')
    return Response(twiml, mimetype="text/xml")


@app.route("/api/health")
def health():
    return jsonify({"ok": True, "agent_count": len(AGENTS), "index": rag.index_status()})


@app.route("/api/index/status")
def index_status():
    return jsonify(rag.index_status())


@app.route("/api/reindex", methods=["POST"])
def reindex():
    try:
        idx = rag.build_index()
        return jsonify({"ok": True, "chunks": len(idx)})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)}), 500


store.init_db()
rag.start_background_index()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, threaded=True)