"""Discover and drive a live mugi_backend over loopback HTTP (Cursor MCP steering).

Used by ``mcp_server.py`` steering tools and the existing probe scripts pattern
(``scripts/tool_call_probe_http.py``). Only talks to ``127.0.0.1``; auth token
is read from the running backend process environment (same as the Swift host).
"""
from __future__ import annotations

import json
import os
import subprocess
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Any


@dataclass(frozen=True)
class BackendConnection:
    pid: int
    port: int
    base_url: str
    auth_token: str
    repo_root: str | None = None


def _pgrep_backend_pids() -> list[int]:
    try:
        out = subprocess.check_output(
            ["pgrep", "-f", "mugi_backend"],
            text=True,
            stderr=subprocess.DEVNULL,
        ).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return []
    return [int(x) for x in out.splitlines() if x.strip().isdigit()]


def _port_from_pid(pid: int) -> int | None:
    try:
        out = subprocess.check_output(
            ["lsof", "-nP", "-p", str(pid)],
            text=True,
            stderr=subprocess.DEVNULL,
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        return None
    for line in out.splitlines():
        if "TCP" not in line or "LISTEN" not in line:
            continue
        part = line.split()[-2]
        if ":" in part:
            try:
                return int(part.rsplit(":", 1)[-1])
            except ValueError:
                continue
    return None


def _auth_from_pid(pid: int) -> str | None:
    try:
        out = subprocess.check_output(["ps", "eww", "-p", str(pid)], text=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        return None
    for tok in out.split():
        if tok.startswith("MUGI_AUTH_TOKEN="):
            return tok.split("=", 1)[1]
    # Headless dev (`make backend` without env): mugi_backend writes the token to a
    # mode-0600 tempfile and prints MUGI_READY:<port>:file:<path>. Swift sets
    # MUGI_AUTH_TOKEN in the child env; loopback probes fall back to the newest
    # mugi_token_*.txt in the system temp dir (single-backend dev assumption).
    import glob
    import tempfile

    temp = tempfile.gettempdir()
    candidates = sorted(
        glob.glob(os.path.join(temp, "mugi_token_*.txt")),
        key=os.path.getmtime,
        reverse=True,
    )
    for path in candidates[:5]:
        try:
            with open(path, encoding="utf-8") as f:
                tok = f.read().strip()
            if tok:
                return tok
        except OSError:
            continue
    return None


def _repo_root_from_pid_file() -> str | None:
    path = os.path.expanduser("~/.mugi/.swift_backend_pid.json")
    try:
        with open(path, encoding="utf-8") as f:
            obj = json.load(f)
        root = str(obj.get("repo_root") or "").strip()
        return root or None
    except Exception:
        return None


def discover_backend() -> BackendConnection | dict[str, Any]:
    """Return a connection to the listening mugi_backend, or ``{error: ...}``."""
    pid = None
    for candidate in _pgrep_backend_pids():
        if _port_from_pid(candidate) is not None:
            pid = candidate
            break
    if pid is None:
        return {
            "error": "No listening mugi_backend process found. Start the Mugi app first.",
        }
    port = _port_from_pid(pid)
    if port is None:
        return {"error": f"Backend pid {pid} has no TCP LISTEN port."}
    token = _auth_from_pid(pid)
    if not token:
        return {"error": f"Could not read MUGI_AUTH_TOKEN from backend pid {pid}."}
    return BackendConnection(
        pid=pid,
        port=port,
        base_url=f"http://127.0.0.1:{port}",
        auth_token=token,
        repo_root=_repo_root_from_pid_file(),
    )


def http_json(
    method: str,
    url: str,
    headers: dict[str, str],
    body: dict[str, Any] | None = None,
    *,
    timeout: float = 120,
) -> tuple[int, Any]:
    data = None
    hdrs = dict(headers)
    if body is not None:
        data = json.dumps(body).encode("utf-8")
        hdrs.setdefault("Content-Type", "application/json")
    req = urllib.request.Request(url, data=data, headers=hdrs, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
            code = resp.getcode()
    except urllib.error.HTTPError as e:
        raw = e.read().decode("utf-8", errors="replace")
        code = e.code
    try:
        return code, json.loads(raw) if raw.strip() else {}
    except json.JSONDecodeError:
        return code, raw


def _headers(conn: BackendConnection) -> dict[str, str]:
    return {"X-Auth-Token": conn.auth_token}


def backend_health(conn: BackendConnection) -> dict[str, Any]:
    code, body = http_json("GET", f"{conn.base_url}/health", {}, timeout=10)
    ok = code == 200 and isinstance(body, dict) and body.get("ok") is True
    return {"http_status": code, "ok": ok, "body": body}


def new_conversation(conn: BackendConnection) -> dict[str, Any]:
    code, body = http_json("POST", f"{conn.base_url}/conversations/new", _headers(conn))
    if code != 200 or not isinstance(body, dict) or not body.get("thread_id"):
        return {"error": f"POST /conversations/new failed ({code})", "body": body}
    return {"thread_id": str(body["thread_id"]), "body": body}


def send_message(
    conn: BackendConnection,
    *,
    text: str,
    thread_id: str,
    research_mode: bool = False,
) -> dict[str, Any]:
    msg = (text or "").strip()
    if not msg:
        return {"error": "text is required"}
    tid = (thread_id or "").strip()
    if not tid:
        return {"error": "thread_id is required"}
    code, body = http_json(
        "POST",
        f"{conn.base_url}/send",
        _headers(conn),
        {"text": msg, "thread_id": tid, "research_mode": bool(research_mode)},
        timeout=30,
    )
    if code != 200:
        return {"error": f"POST /send failed ({code})", "body": body}
    return {"ok": True, "thread_id": tid, "body": body}


def get_thread_state(conn: BackendConnection, thread_id: str) -> dict[str, Any]:
    tid = (thread_id or "").strip()
    code, body = http_json(
        "GET",
        f"{conn.base_url}/chat/thread_state?thread_id={tid}",
        _headers(conn),
        timeout=15,
    )
    if code != 200:
        return {"error": f"GET /chat/thread_state failed ({code})", "body": body}
    return body if isinstance(body, dict) else {"raw": body}


def get_messages(conn: BackendConnection, thread_id: str) -> dict[str, Any]:
    tid = (thread_id or "").strip()
    code, body = http_json(
        "GET",
        f"{conn.base_url}/chat/messages?thread_id={tid}",
        _headers(conn),
        timeout=60,
    )
    if code != 200:
        return {"error": f"GET /chat/messages failed ({code})", "body": body}
    if not isinstance(body, list):
        return {"error": "unexpected /chat/messages shape", "body": body}
    return {"messages": body, "count": len(body)}


def _summarize_messages(messages: list[dict[str, Any]], *, tail: int = 8) -> list[dict[str, str]]:
    out: list[dict[str, str]] = []
    for m in messages[-tail:]:
        role = str(m.get("role") or "")
        content = str(m.get("content") or "")
        if len(content) > 800:
            content = content[:797] + "..."
        out.append({"role": role, "content": content, "timestamp": str(m.get("timestamp") or "")})
    return out


def _parse_trace_tail(path: str, *, max_events: int = 12) -> list[dict[str, Any]]:
    if not os.path.isfile(path):
        return []
    try:
        with open(path, encoding="utf-8", errors="replace") as f:
            lines = f.readlines()
    except OSError:
        return []
    events: list[dict[str, Any]] = []
    for line in lines[-max(1, max_events * 3) :]:
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(obj, dict):
            events.append(obj)
    return events[-max_events:]


def wait_for_run(
    conn: BackendConnection,
    thread_id: str,
    *,
    timeout_sec: float = 180,
    poll_sec: float = 2.0,
    trace_path: str | None = None,
    baseline_message_count: int | None = None,
) -> dict[str, Any]:
    """Poll until the agent stops running and the last message is assistant (or timeout).

    *baseline_message_count* is the thread length **before** the user message that
    triggered this wait. Completion requires strictly more messages than that
    baseline and an assistant as the final turn — so a pre-existing init/welcome
    pair cannot satisfy a wait for the case prompt's reply.
    """
    tid = (thread_id or "").strip()
    deadline = time.monotonic() + max(5.0, float(timeout_sec))
    poll = max(0.5, float(poll_sec))
    last_running: bool | None = None
    last_msg_count = 0
    saw_assistant_after_send = False
    completed = False
    if baseline_message_count is None:
        msgs_payload = get_messages(conn, tid)
        messages = msgs_payload.get("messages") if isinstance(msgs_payload, dict) else None
        baseline_message_count = (
            len(messages) if isinstance(messages, list) else 0
        )

    while time.monotonic() < deadline:
        state = get_thread_state(conn, tid)
        running = bool(state.get("running")) if isinstance(state, dict) and "error" not in state else None
        if running is not None:
            last_running = running

        msgs_payload = get_messages(conn, tid)
        messages = msgs_payload.get("messages") if isinstance(msgs_payload, dict) else None
        if isinstance(messages, list):
            last_msg_count = len(messages)
            if messages and str(messages[-1].get("role") or "") == "assistant":
                saw_assistant_after_send = True
                if (
                    running is False
                    and last_msg_count > baseline_message_count
                ):
                    completed = True
                    break
        time.sleep(poll)

    timed_out = not completed

    trace_events: list[dict[str, Any]] = []
    if trace_path:
        trace_events = _parse_trace_tail(trace_path)

    final_messages: list[dict[str, str]] = []
    msgs_payload = get_messages(conn, tid)
    if isinstance(msgs_payload, dict) and isinstance(msgs_payload.get("messages"), list):
        final_messages = _summarize_messages(msgs_payload["messages"])

    return {
        "thread_id": tid,
        "timed_out": timed_out,
        "running": last_running,
        "message_count": last_msg_count,
        "saw_assistant_reply": saw_assistant_after_send,
        "messages_tail": final_messages,
        "trace_tail": [
            {
                k: ev.get(k)
                for k in ("event", "round", "reason", "tool_name", "tool_call_count", "result_preview")
                if k in ev
            }
            for ev in trace_events
        ],
    }


def send_and_wait(
    conn: BackendConnection,
    *,
    text: str,
    thread_id: str,
    research_mode: bool = False,
    timeout_sec: float = 180,
    poll_sec: float = 2.0,
    trace_path: str | None = None,
) -> dict[str, Any]:
    pre_payload = get_messages(conn, thread_id)
    pre_messages = pre_payload.get("messages") if isinstance(pre_payload, dict) else None
    pre_count = len(pre_messages) if isinstance(pre_messages, list) else 0
    sent = send_message(conn, text=text, thread_id=thread_id, research_mode=research_mode)
    if sent.get("error"):
        return sent
    waited = wait_for_run(
        conn,
        thread_id,
        timeout_sec=timeout_sec,
        poll_sec=poll_sec,
        trace_path=trace_path,
        baseline_message_count=pre_count,
    )
    return {"send": sent, "wait": waited}
