"""Pi coding agent interface — v1's only coding agent.

Runs `pi -p --mode json` and tails its JSONL stdout line by line, forwarding
each event to a callback WHILE the agent works (the streaming crack, solved
by construction). `--session-id` creates-or-continues, so running and
continuing an agent are the same call: same session id = same context window.
"""

from __future__ import annotations

import json
import os
import queue
import signal
import subprocess
import threading
import time
from enum import Enum
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable, Optional

from .data_types import PiRequest, PiResult
from .utils import now_iso, operator_env

PI_PATH = os.environ.get("PI_PATH", "pi")
MODELS_JSON = os.environ.get("PI_MODELS_PATH",
                             str(Path.home() / ".pi" / "agent" / "models.json"))

RESULT_SNIPPET_CHARS = 20_000   # tool output rides along whole; clip only guards pathological cases
ARG_VALUE_CHARS = 20_000        # args too — the UI scrolls, it must not be handed cut-off data
LABEL_CHARS = 80                # "bash: <command>" shown as the event name
DEFAULT_AGENT_TIMEOUT_SECONDS = 1800
DEFAULT_AGENT_IDLE_TIMEOUT_SECONDS = 600
TERMINATE_GRACE_SECONDS = 2.0

# The arg that identifies a call at a glance, in the order tools tend to use.
PRIMARY_ARGS = ("command", "path", "file_path", "pattern", "query", "url")


class TimeoutLifecycle(str, Enum):
    STARTING = "starting"
    STREAMING = "streaming"
    TERMINATING = "terminating"
    TERMINAL = "terminal"


def _count(value: str) -> int:
    """Parse pi's compact model-list counts (`272K`, `1.0M`)."""
    suffixes = {"K": 1_000, "M": 1_000_000}
    suffix = value[-1:].upper()
    if suffix in suffixes:
        return int(float(value[:-1]) * suffixes[suffix])
    return int(value)


@lru_cache(maxsize=1)
def _pi_catalog() -> list[tuple[str, str, int]]:
    """Read pi's merged catalog, including built-in providers and custom models."""
    try:
        result = subprocess.run(
            [PI_PATH, "--list-models"], capture_output=True, text=True,
            timeout=30, env=operator_env(), check=False,
        )
    except (OSError, subprocess.TimeoutExpired):
        return []
    if result.returncode != 0:
        return []
    rows = []
    for line in result.stdout.splitlines()[1:]:
        columns = line.split()
        if len(columns) < 3:
            continue
        try:
            rows.append((columns[0], columns[1], _count(columns[2])))
        except ValueError:
            continue
    return rows


def resolve_model(pattern: str) -> tuple[str, str]:
    """Resolve a model pattern to an explicit ``(provider, model_id)`` pair.

    Pi's catalog merges built-in models with ``~/.pi/agent/models.json``. Using
    that same merged view lets SSSF target direct providers such as
    ``openai/gpt-5.6-terra`` without re-registering built-in models locally.
    """
    catalog = [(provider, model_id) for provider, model_id, _ in _pi_catalog()]
    if "/" in pattern:
        provider, model_id = pattern.split("/", 1)
        if (provider, model_id) in catalog:
            return provider, model_id
    matches = [(provider, model_id) for provider, model_id in catalog
               if pattern == model_id or pattern in model_id]
    exact = [match for match in matches
             if match[1] == pattern or match[1].endswith("/" + pattern)]
    if len(exact) == 1:
        return exact[0]
    if len(matches) == 1:
        return matches[0]
    if not matches:
        raise ValueError(f"model pattern {pattern!r} not found in pi --list-models — "
                         "authenticate/register it or fix the config")
    raise ValueError(f"model pattern {pattern!r} is ambiguous: {matches}")


def _context_tokens(usage: dict) -> int:
    """Tokens occupying the window after a turn.

    Mirrors pi's own `calculateContextTokens` (coding-agent
    `core/compaction/compaction.ts`), which is what pi compacts against and
    shows in its footer: prefer the provider's `totalTokens`, else sum the
    parts. Cache reads count — cached prompt is still prompt.
    """
    total = usage.get("totalTokens") or 0
    if total:
        return int(total)
    return int(sum(usage.get(part) or 0
                   for part in ("input", "output", "cacheRead", "cacheWrite")))


def context_window(provider: str, model_id: str) -> int:
    """The model's context ceiling from pi's merged model catalog."""
    registry = json.loads(Path(MODELS_JSON).read_text())
    for model in registry.get("providers", {}).get(provider, {}).get("models", []):
        if model.get("id") == model_id:
            return int(model.get("contextWindow") or 0)
    for listed_provider, listed_model, window in _pi_catalog():
        if listed_provider == provider and listed_model == model_id:
            return window
    return 0


def _text_of(container: dict) -> str:
    """Join the text blocks of anything pi shapes as {content: [...]} — a
    message or a tool result."""
    return "".join(part.get("text", "") for part in container.get("content", []) or []
                   if isinstance(part, dict) and part.get("type") == "text")


def _clip(text: str, limit: int) -> str:
    return text if len(text) <= limit else text[:limit].rstrip() + "…"


def _label(tool: str, args: dict) -> str:
    """One-line human name for a tool call: `bash: ls -la src`."""
    value = next((args[key] for key in PRIMARY_ARGS
                  if isinstance(args.get(key), str) and args[key].strip()), "")
    if not value:
        value = next((v for v in args.values() if isinstance(v, str) and v.strip()), "")
    value = " ".join(str(value).split())
    return f"{tool}: {_clip(value, LABEL_CHARS)}" if value else tool


class ToolCallTracker:
    """Folds pi's tool stream into ONE normalized record per completed call.

    pi announces a call as a `toolCall` content block, then emits
    tool_execution_start / _update / _end for it. Only the end carries the
    result, so that is where a record is emitted — one trace event per real
    tool call, the moment it returns, instead of three shapeless ones.

    The record carries the call's real span (`started_at`/`ended_at`), which the
    tracer writes to columns so the UI can lay tool calls on a time axis without
    parsing every payload.
    """

    def __init__(self) -> None:
        self._open: dict[str, dict] = {}

    def observe(self, event: dict) -> Optional[dict]:
        """Returns the record for a finished tool call, else None."""
        etype = event.get("type", "")
        if etype == "message_end":
            for block in event.get("message", {}).get("content", []) or []:
                if isinstance(block, dict) and block.get("type") == "toolCall":
                    self._announce(block.get("id"), block.get("name"),
                                   block.get("arguments"))
            return None
        if etype == "tool_execution_start":
            self._announce(event.get("toolCallId"), event.get("toolName"),
                           event.get("args"))
            return None
        if etype != "tool_execution_end":
            return None

        call_id = str(event.get("toolCallId") or "")
        opened = self._open.pop(call_id, {})
        tool = str(event.get("toolName") or opened.get("tool") or "tool")
        args = event.get("args") or opened.get("args") or {}
        record = {
            "tool": tool,
            "tool_call_id": call_id,
            "args": {key: _clip(value, ARG_VALUE_CHARS) if isinstance(value, str) else value
                     for key, value in args.items()},
            "ok": not event.get("isError", False),
            "label": _label(tool, args),
        }
        result_text = _text_of(event.get("result") or {})
        if result_text:
            record["result_snippet"] = _clip(result_text, RESULT_SNIPPET_CHARS)
        record["ended_at"] = now_iso()
        if opened.get("clock"):
            record["duration_ms"] = int((time.monotonic() - opened["clock"]) * 1000)
        if opened.get("started_at"):
            record["started_at"] = opened["started_at"]
        return record

    def _announce(self, call_id, tool, args) -> None:
        """First sighting starts the clock; a later sighting only fills gaps."""
        if not call_id:
            return
        known = self._open.get(str(call_id), {})
        self._open[str(call_id)] = {
            "tool": tool or known.get("tool", ""),
            "args": args or known.get("args", {}),
            "started_at": known.get("started_at") or now_iso(),   # wall clock, for the row
            "clock": known.get("clock") or time.monotonic(),      # monotonic, for duration
        }


def run(request: PiRequest, on_event: Optional[Callable[[dict], None]] = None,
        on_spawn: Optional[Callable[[int], None]] = None,
        on_exit: Optional[Callable[[int], None]] = None,
        timeout_seconds: float | None = None,
        idle_timeout_seconds: float | None = None, *,
        clock: Callable[[], float] | None = None,
        sleeper: Callable[[float], None] | None = None,
        popen_factory: Callable[..., Any] | None = None,
        kill_process_group: Callable[[int, int], None] | None = None,
        thread_factory: Callable[..., Any] | None = None) -> PiResult:
    """Run one non-interactive pi turn."""
    clock = clock or time.monotonic
    sleeper = sleeper or time.sleep
    popen_factory = popen_factory or subprocess.Popen
    kill_process_group = kill_process_group or os.killpg
    thread_factory = thread_factory or threading.Thread
    provider, model_id = resolve_model(request.model)
    cmd = [
        PI_PATH, "-p", "--mode", "json",
        "--provider", provider, "--model", model_id,
        "--thinking", request.thinking,
        "--session-id", request.session_id,
        "--session-dir", request.session_dir,
        "--system-prompt", request.system_prompt,
    ]
    if request.tools:
        cmd += ["--tools", ",".join(request.tools)]
    for extension in request.extensions:
        cmd += ["-e", extension]
    cmd.append(request.prompt)

    raw_path = Path(request.raw_output_path)
    raw_path.parent.mkdir(parents=True, exist_ok=True)
    stderr_path = raw_path.with_name("stderr.log")
    timeout = (timeout_seconds if timeout_seconds is not None else float(
        os.environ.get("FACTORY_AGENT_TIMEOUT_SECONDS", DEFAULT_AGENT_TIMEOUT_SECONDS)))
    idle_timeout = (idle_timeout_seconds if idle_timeout_seconds is not None else float(
        os.environ.get("FACTORY_AGENT_IDLE_TIMEOUT_SECONDS", DEFAULT_AGENT_IDLE_TIMEOUT_SECONDS)))
    result = PiResult(session_id=request.session_id,
                      context_window=context_window(provider, model_id))
    process: Any = None
    timed_out = False
    timeout_kind: str | None = None
    error_text: str | None = None
    lifecycle = TimeoutLifecycle.STARTING
    terminal_receipt_emitted = False
    try:
        process = popen_factory(cmd, stdin=subprocess.DEVNULL,
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                text=True, bufsize=1, cwd=request.cwd,
                                env=operator_env(), start_new_session=True)
        if on_spawn:
            on_spawn(process.pid)

        stream_records: queue.Queue[dict] = queue.Queue()
        last_output_clock: float | None = None

        def consume_stdout() -> None:
            nonlocal last_output_clock, lifecycle
            with raw_path.open("a") as raw:
                assert process is not None and process.stdout is not None
                for line in process.stdout:
                    last_output_clock = clock()
                    lifecycle = TimeoutLifecycle.STREAMING
                    raw.write(line)
                    raw.flush()
                    try:
                        event = json.loads(line.strip())
                    except json.JSONDecodeError:
                        continue
                    if event.get("type") == "message_end":
                        message = event.get("message", {})
                        if message.get("role") == "assistant":
                            text = _text_of(message)
                            if text:
                                result.text = text
                            usage = message.get("usage", {}) or {}
                            turn = _context_tokens(usage)
                            result.tokens += turn
                            result.usage.add_turn(usage, turn)
                            if turn and message.get("stopReason") not in ("aborted", "error"):
                                result.context_tokens = turn
                            result.cost += (usage.get("cost", {}) or {}).get("total", 0.0) or 0.0
                    stream_records.put(event)

        stderr_chunks: list[str] = []

        def consume_stderr() -> None:
            assert process is not None and process.stderr is not None
            for line in process.stderr:
                stderr_chunks.append(line)

        reader = thread_factory(target=consume_stdout, daemon=True)
        errors = thread_factory(target=consume_stderr, daemon=True)
        reader.start()
        errors.start()
        started_clock = clock()

        def forward_stream_records() -> None:
            while True:
                try:
                    event = stream_records.get_nowait()
                except queue.Empty:
                    return
                if on_event:
                    on_event(event)

        while process.poll() is None:
            forward_stream_records()
            now = clock()
            if now - started_clock >= timeout:
                timeout_kind = "wall"
                break
            if (lifecycle is TimeoutLifecycle.STREAMING and last_output_clock is not None
                    and now - last_output_clock >= idle_timeout):
                timeout_kind = "idle"
                break
            sleep_seconds = timeout - (now - started_clock)
            if lifecycle is TimeoutLifecycle.STREAMING and last_output_clock is not None:
                sleep_seconds = min(sleep_seconds, idle_timeout - (now - last_output_clock))
            sleeper(min(0.05, max(0.0, sleep_seconds)))
        if timeout_kind is not None:
            timed_out = True
            lifecycle = TimeoutLifecycle.TERMINATING
            grace_started_clock = clock()
            kill_process_group(process.pid, signal.SIGTERM)
            try:
                process.wait(timeout=TERMINATE_GRACE_SECONDS)
            except subprocess.TimeoutExpired:
                pass
            grace_remaining = TERMINATE_GRACE_SECONDS - (clock() - grace_started_clock)
            if grace_remaining > 0:
                sleeper(grace_remaining)
            try:
                kill_process_group(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            try:
                result.returncode = process.wait(timeout=TERMINATE_GRACE_SECONDS)
            except subprocess.TimeoutExpired as exc:
                raise TimeoutError("pi process group did not exit after SIGKILL") from exc
        else:
            result.returncode = process.wait()
        lifecycle = TimeoutLifecycle.TERMINAL
        reader.join(timeout=TERMINATE_GRACE_SECONDS)
        errors.join(timeout=TERMINATE_GRACE_SECONDS)
        if reader.is_alive() or errors.is_alive():
            raise TimeoutError("pi stream drain timed out")
        forward_stream_records()
        stderr = "".join(stderr_chunks)
        stderr_path.write_text(stderr)
        if timed_out:
            limit = timeout if timeout_kind == "wall" else idle_timeout
            error_text = f"pi {timeout_kind} timeout: timed out after {limit:g}s"
            raise TimeoutError(error_text)
        if result.returncode != 0 and not result.text:
            error_text = f"pi exited {result.returncode}: {stderr.strip()[-800:]}"
            raise RuntimeError(error_text)
        return result
    except BaseException as exc:
        error_text = error_text or str(exc)
        raise
    finally:
        lifecycle = TimeoutLifecycle.TERMINAL
        if process is not None and on_exit:
            on_exit(process.pid)
        if on_event and not terminal_receipt_emitted:
            terminal_receipt_emitted = True
            returncode = result.returncode if process is not None else None
            sig = -returncode if returncode is not None and returncode < 0 else None
            usage = result.usage.model_dump() if result.tokens else None
            on_event({"type": "agent_attempt_end", "returncode": returncode,
                      "signal": sig, "timedOut": timed_out,
                      "timeoutKind": timeout_kind, "stderrPath": str(stderr_path),
                      "tokens": result.tokens or None, "usage": usage,
                      "error": error_text, "lifecycle": lifecycle.value})
