"""
KimiCLIAdapter — drives MoonshotAI's `kimi` (kimi-cli) via per-turn subprocess.

⚠️ UNVERIFIED / needs-device-test: written from docs
(moonshotai.github.io/kimi-cli print-mode + reference), not observed against a
live `kimi`. Confirm whether `content` streams as deltas or whole messages, and
whether any session-id line is emitted, before flipping `beta=False`.

Each turn spawns (print mode auto-exits ⇒ per-turn/leaseless):
  kimi --print -p <text> --output-format stream-json [--resume <id> | --continue]

`--output-format stream-json` emits JSON Lines in OpenAI chat-message shape,
keyed by `role`:
  {"role":"user","content":"…"}                                   (echo — ignored)
  {"role":"assistant","content":"…","tool_calls":[{"id","type":"function",
        "function":{"name","arguments":"<json-string>"}}]}
  {"role":"tool","tool_call_id":"…","content":"…"}                (tool result)
Turn end = process exit (no terminal event). Exit codes: 0 ok · 1 non-retryable
· 75 retryable — handled by PerTurnCLIAdapter's exit-code tail.
"""
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,
    TextBlock, ToolCallBlock, tool_result_or_large,
)
from .base import PerTurnCLIAdapter

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


class KimiCLIAdapter(PerTurnCLIAdapter):
    cli_name = "kimi"
    _capture_stderr = True
    _close_stdin = True
    cap_images = True          # via OCR-text inline (no native image path found)
    cap_documents = False
    cap_thinking = False       # no thinking block documented in stream-json
    cap_interrupt = True
    cap_input_during_turn = False
    # MCP via ~/.config/kimi/mcp.json (setup_integrations); unverified until a live run.
    cap_sentinel = False
    provides_models = False
    provides_usage = False
    agent_id = "kimi"
    display_name = "Kimi CLI"
    beta = True
    cli_binary = "kimi"

    def __init__(
        self,
        binary: str | None = None,
        cwd: str | None = None,
        model: str | None = None,
        resume_session_id: str | None = None,
        on_session_id: "Callable[[str], None] | None" = None,
    ) -> None:
        self._binary = binary or shutil.which("kimi") or "kimi"
        self._cwd = cwd
        self._model = model
        self._session_id: str | None = resume_session_id   # ULID under ~/.kimi/sessions
        self._on_session_id = on_session_id
        self._first_turn = True
        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("MOONSHOT_API_KEY") or os.environ.get("KIMI_API_KEY")
                or (Path.home() / ".kimi" / "credentials").exists()
                or (Path.home() / ".kimi" / "config.toml").exists()):
            return True, ""
        return False, "Run `kimi` then `/login` (or set MOONSHOT_API_KEY)."

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

    def clear_history(self) -> None:
        self._session_id = None
        self._first_turn = True
        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)
        argv = [self._binary, "--print", "-p", prompt,
                "--output-format", "stream-json"]
        if self._model:
            argv += ["-m", self._model]
        # Prefer an explicit id (adopted external session); else continue the most
        # recent session in this cwd on turns ≥2. stream-json may not emit an id, so
        # --continue is the reliable within-cwd continuity path.
        if self._session_id:
            argv += ["--resume", self._session_id]
        elif not self._first_turn:
            argv += ["--continue"]
        self._first_turn = False
        return argv

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

    async def _parse_event(self, event: dict) -> AsyncIterator[Block]:
        # Defensive: if any build emits a session id, capture it for --resume.
        sid = event.get("session_id") or event.get("sessionId")
        if isinstance(sid, str) and sid and sid != self._session_id:
            self._session_id = sid
            if self._on_session_id:
                self._on_session_id(sid)

        role = event.get("role", "")

        if role == "assistant":
            content = event.get("content")
            if isinstance(content, str) and content:
                yield TextBlock(text=content, done=True)
            for call in (event.get("tool_calls") or []):
                if not isinstance(call, dict):
                    continue
                fn = call.get("function") or {}
                tcid = call.get("id") or str(uuid.uuid4())
                name = fn.get("name", "tool")
                raw_args = fn.get("arguments")
                args = _parse_args(raw_args)
                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 role == "tool":
            tcid = event.get("tool_call_id") or ""
            meta = self._pending_tools.pop(tcid, {})
            raw = event.get("content", "")
            text = raw if isinstance(raw, str) else json.dumps(raw)
            is_error = bool(event.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
        # role == "user" (echo) and unknown roles: ignored. Turn end = process exit.


def _parse_args(raw) -> dict:
    """OpenAI tool-call arguments come as a JSON string; be forgiving."""
    if isinstance(raw, dict):
        return raw
    if isinstance(raw, str) and raw.strip():
        try:
            v = json.loads(raw)
            return v if isinstance(v, dict) else {"value": v}
        except Exception:
            return {"raw": raw}
    return {}
