"""
GrokCLIAdapter — drives the `grok` CLI (Superagent `@vibe-kit/grok-cli`, variant A)
via per-turn subprocess.

⚠️ UNVERIFIED / needs-device-test: written from the superagent-ai/grok-cli README,
NOT observed live. Per-event FIELD NAMES are undocumented — this parser reads them
defensively (multiple candidate keys). Two incompatible `grok` binaries exist:
  A. Superagent `@vibe-kit/grok-cli`  — `grok -p … --format json`, GROK_API_KEY   (THIS adapter)
  B. xAI official "Grok Build"        — `grok -p/--single … --output-format streaming-json`
                                        or ACP (`grok agent stdio`), XAI_API_KEY   (TODO: separate branch)
Confirm variant + field names before flipping `beta=False`.

Each turn spawns (headless one-prompt-then-exit ⇒ per-turn/leaseless):
  grok -p <text> --format json [--session <id> | --session latest]

`--format json` emits newline-delimited JSON events with semantic types
`step_start` / `text` / `tool_use` / `step_finish` / `error`. Turn end = process
exit (grok -p exits after one prompt) — handled by the 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, ErrorBlock,
    TextBlock, ToolCallBlock, tool_result_or_large,
)
from .base import PerTurnCLIAdapter

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


class GrokCLIAdapter(PerTurnCLIAdapter):
    cli_name = "grok"
    _capture_stderr = True
    _close_stdin = True
    cap_images = True          # via OCR-text inline
    cap_documents = False
    cap_thinking = False
    cap_interrupt = True
    cap_input_during_turn = False
    # MCP via ~/.grok/user-settings.json (setup_integrations); unverified until a live run.
    cap_sentinel = False
    provides_models = False
    provides_usage = False
    agent_id = "grok"
    display_name = "Grok CLI"
    beta = True
    cli_binary = "grok"

    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("grok") or "grok"
        self._cwd = cwd
        self._model = model
        self._session_id: str | None = resume_session_id
        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._text_acc = ""
        self._proc: asyncio.subprocess.Process | None = None

    @classmethod
    def readiness(cls) -> tuple[bool, str]:
        if (os.environ.get("GROK_API_KEY") or os.environ.get("XAI_API_KEY")
                or (Path.home() / ".grok").exists()):
            return True, ""
        return False, "Set GROK_API_KEY (Superagent grok-cli) or run its login."

    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="Grok 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, "-p", prompt, "--format", "json"]
        if self._model:
            argv += ["--model", self._model]
        if self._cwd:
            argv += ["--directory", self._cwd]
        if self._session_id:
            argv += ["--session", self._session_id]
        elif not self._first_turn:
            argv += ["--session", "latest"]
        self._first_turn = False
        return argv

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

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

        # Defensive session-id capture (field name unconfirmed).
        sid = event.get("session_id") or event.get("sessionId") 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)

        if etype == "text":
            chunk = event.get("text") or event.get("content") or event.get("delta") or ""
            if chunk:
                self._text_acc += chunk
                yield TextBlock(text=chunk, done=False)
            return

        if etype == "tool_use":
            tcid = event.get("id") or event.get("toolCallId") or str(uuid.uuid4())
            name = event.get("name") or event.get("tool") or "tool"
            args = (event.get("input") or event.get("arguments")
                    or event.get("args") or {})
            if not isinstance(args, dict):
                args = {"value": args}
            # Some step-level formats fold the result into the same event.
            result = event.get("result") or event.get("output")
            yield AudioHintBlock(state="agent_silent")
            if result is None:
                self._pending_tools[tcid] = {"name": name, "args": args}
                yield ToolCallBlock(id=tcid, name=name, args=args, status="running")
            else:
                content = result if isinstance(result, str) else json.dumps(result)
                is_error = bool(event.get("is_error") or event.get("error"))
                yield ToolCallBlock(id=tcid, name=name, args=args,
                                    status="error" if is_error else "done")
                yield tool_result_or_large(tool_call_id=tcid, content=content,
                                           fetch_store=self._fetch_store, is_error=is_error)
                yield AudioHintBlock(state="agent_speaking")
            return

        if etype == "tool_result":  # if emitted as a distinct event
            tcid = event.get("id") or event.get("toolCallId") or ""
            meta = self._pending_tools.pop(tcid, {})
            raw = event.get("result") or event.get("output") or event.get("content") or ""
            content = raw if isinstance(raw, str) else json.dumps(raw)
            is_error = bool(event.get("is_error") or event.get("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=content,
                                       fetch_store=self._fetch_store, is_error=is_error)
            yield AudioHintBlock(state="agent_speaking")
            return

        if etype == "error":
            msg = event.get("message") or event.get("error") or "grok error"
            yield ErrorBlock(message=str(msg)[:500], recoverable=True)
            return
        # step_start / step_finish and unknown events: no-op. Turn end = process exit.

    async def _on_no_session_end(self) -> AsyncIterator[Block]:
        # Flush any streamed text into a finalized block before the end tail, so the
        # streaming TextBlock is closed (grok emits `text` deltas, no explicit finish).
        if self._text_acc:
            yield TextBlock(text=self._text_acc, done=True)
            self._text_acc = ""
        async for b in super()._on_no_session_end():
            yield b
