"""
AmpCLIAdapter — drives Sourcegraph's `amp` CLI via per-turn subprocess.

⚠️ UNVERIFIED / needs-device-test: written from docs (ampcode.com/manual +
ampcode.com/news/streaming-json), not observed against a live `amp`. The
stream-json format is documented as "Claude Code compatible", so this parser
mirrors the Anthropic-messages shape. Confirm tool_use/tool_result field names
and whether assistant text streams as deltas before flipping `beta=False`.

Each turn spawns (fresh process; resume via the `threads continue` subcommand):
  amp -x <text> --stream-json --dangerously-allow-all                 (first turn)
  amp threads continue <T-id> --execute <text> --stream-json --dangerously-allow-all

`--stream-json` emits JSON Lines. Event `type` ∈ {system, user, assistant, result}:
  {type:"system", subtype:"init", session_id:"T-…", cwd, tools[], mcp_servers[]}
  {type:"user",   message:{role, content:[…]}, session_id}            (tool_result here)
  {type:"assistant", message:{role, content:[{type:text|thinking|tool_use,…}],
                              stop_reason, usage:{…}}, session_id}
  {type:"result", subtype:"success"|…, is_error, num_turns, result, session_id}
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import shutil
import uuid
from pathlib import Path
from typing import AsyncIterator, Callable

from ..attachments import Attachment, build_text_with_inline
from ..blocks import (
    Block, AgentInfoBlock, AudioHintBlock, ErrorBlock, SessionEndBlock,
    TextBlock, ThinkingBlock, ToolCallBlock, tool_result_or_large,
)
from .base import PerTurnCLIAdapter

log = logging.getLogger("beacon.amp")


class AmpCLIAdapter(PerTurnCLIAdapter):
    cli_name = "amp"
    _capture_stderr = True
    _close_stdin = True
    cap_images = True          # via OCR-text inline (native path TODO)
    cap_documents = False
    cap_thinking = True        # --stream-json-thinking emits thinking blocks
    cap_interrupt = True
    cap_input_during_turn = False
    # MCP via ~/.config/amp/settings.json (setup_integrations); unverified until a live run.
    cap_sentinel = False
    provides_models = False
    provides_usage = False
    agent_id = "amp"
    display_name = "Amp"
    beta = True                # UNVERIFIED wire format — keep the Beta badge
    cli_binary = "amp"

    def __init__(
        self,
        binary: str | None = None,
        cwd: str | None = None,
        model: str | None = None,   # accepted for parity; Amp picks its own model
        resume_session_id: str | None = None,
        on_session_id: "Callable[[str], None] | None" = None,
    ) -> None:
        self._binary = binary or shutil.which("amp") or "amp"
        self._cwd = cwd
        self._model = model               # accepted for parity; Amp picks its own
        self._thread_id: str | None = resume_session_id   # Amp thread id "T-…"
        self._on_session_id = on_session_id
        self._send_queue: asyncio.Queue[tuple[str, list[Attachment]]] = None  # type: ignore
        self._pending_tools: dict[str, dict] = {}
        self._fetch_store: dict[str, str] = {}
        self._proc: asyncio.subprocess.Process | None = None

    @classmethod
    def readiness(cls) -> tuple[bool, str]:
        if os.environ.get("AMP_API_KEY") or (Path.home() / ".config" / "amp" / "settings.json").exists():
            return True, ""
        return False, "Run `amp login` (or set AMP_API_KEY)."

    async def start(self, session_id: str) -> AgentInfoBlock:
        self._beacon_session_id = session_id
        self._send_queue = asyncio.Queue()
        return self.make_agent_info(name="Amp",
                                    working_dirs=[self._cwd] if self._cwd else [])

    def clear_history(self) -> None:
        self._thread_id = None
        if self._on_session_id:
            self._on_session_id("")

    async def _build_argv(self, text: str, attachments: list[Attachment]) -> list[str]:
        prompt = await asyncio.to_thread(build_text_with_inline, text, attachments)
        common = ["--stream-json", "--dangerously-allow-all"]
        if self._thread_id:
            # Resume a specific thread by id.
            return [self._binary, "threads", "continue", self._thread_id,
                    "--execute", prompt, *common]
        return [self._binary, "-x", prompt, *common]

    def _reset_turn_state(self) -> None:
        self._pending_tools.clear()

    def _capture_sid(self, sid: str | None) -> None:
        if sid and sid != self._thread_id:
            self._thread_id = sid
            if self._on_session_id:
                self._on_session_id(sid)

    async def _parse_event(self, event: dict) -> AsyncIterator[Block]:
        etype = event.get("type", "")

        if etype == "system":
            if event.get("subtype") == "init":
                self._capture_sid(event.get("session_id"))
            return

        if etype == "assistant":
            self._capture_sid(event.get("session_id"))
            msg = event.get("message") or {}
            for block in (msg.get("content") or []):
                if not isinstance(block, dict):
                    continue
                btype = block.get("type")
                if btype == "text":
                    txt = block.get("text", "")
                    if txt:
                        yield TextBlock(text=txt, done=True)
                elif btype == "thinking":
                    th = block.get("thinking") or block.get("text", "")
                    if th:
                        yield ThinkingBlock(thinking=th)
                elif btype == "tool_use":
                    tcid = block.get("id") or str(uuid.uuid4())
                    name = block.get("name", "tool")
                    args = block.get("input") or {}
                    self._pending_tools[tcid] = {"name": name, "args": args}
                    yield AudioHintBlock(state="agent_silent")
                    yield ToolCallBlock(id=tcid, name=name, args=args, status="running")
            return

        if etype == "user":
            # Tool results arrive as tool_result blocks in a user message.
            msg = event.get("message") or {}
            content = msg.get("content")
            if not isinstance(content, list):
                return
            for block in content:
                if not isinstance(block, dict) or block.get("type") != "tool_result":
                    continue
                tcid = block.get("tool_use_id") or ""
                meta = self._pending_tools.pop(tcid, {})
                raw = block.get("content", "")
                text = raw if isinstance(raw, str) else json.dumps(raw)
                is_error = bool(block.get("is_error"))
                if meta:
                    yield ToolCallBlock(id=tcid, name=meta.get("name", "tool"),
                                        args=meta.get("args", {}),
                                        status="error" if is_error else "done")
                yield tool_result_or_large(tool_call_id=tcid, content=text,
                                           fetch_store=self._fetch_store, is_error=is_error)
                yield AudioHintBlock(state="agent_speaking")
            return

        if etype == "result":
            self._capture_sid(event.get("session_id"))
            yield AudioHintBlock(state="agent_silent")
            if event.get("is_error") or event.get("subtype") not in (None, "success"):
                msg = event.get("result") or event.get("error") or "amp turn failed"
                yield ErrorBlock(message=str(msg)[:500], recoverable=True)
                yield SessionEndBlock(reason="error")
            else:
                yield SessionEndBlock(reason="completed")
            return
