"""Claude Agent SDK backend adapter.

One :class:`ClaudeSdkBackend` per :class:`ClaudeAgent`. A live session is
one ``ClaudeSDKClient`` (one claude subprocess) and runs one turn at a
time. Sessions are keyed by ``conversation_id`` when the frontend passes
one (markdown chats) or by a text-only fingerprint of ``messages[:-1]``
for stateless callers (``/v1/messages``). Without a live session the
adapter resumes ``session_id`` from the session store, or seeds the
incoming history into the store via ``core/transcript`` and resumes that.

Events on the wire are the Anthropic ``MessageStreamEvent`` family: one
``message_start``/``message_stop`` envelope per turn, block indices
rebased across the API calls claude makes inside the turn.

Process isolation: claude is spawned through a small exec wrapper that
drops every inherited environment variable outside a whitelist, and,
when ``RunnerConfig.user`` is set, under that uid.
"""

from __future__ import annotations

import asyncio
import contextlib
import fnmatch
import hashlib
import json
import logging
import os
import pwd
import shutil
import sys
import tempfile
import time
import uuid
from collections.abc import Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol, Self, cast

import claude_agent_sdk
from claude_agent_sdk import (
    AssistantMessage,
    ClaudeAgentOptions,
    ClaudeSDKClient,
    ResultMessage,
    StreamEvent,
    TextBlock,
    ThinkingBlock,
    ToolResultBlock,
    ToolUseBlock,
    UserMessage,
    project_key_for_directory,
)

from beaver_gateway.core import prompt as prompt_assembly
from beaver_gateway.core.events import (
    StopReason,
    build_content_block_stop,
    build_input_json_delta,
    build_message_delta,
    build_message_start,
    build_message_stop,
    build_signature_delta,
    build_text_block_start,
    build_text_delta,
    build_thinking_block_start,
    build_thinking_delta,
    build_tool_use_block_start,
)
from beaver_gateway.core.transcript import build_entries
from beaver_gateway.core.turn_capture import TurnCapture, TurnUsage

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence

    from anthropic.types import MessageParam
    from claude_agent_sdk import SessionStore

    from beaver_gateway.agents.base import BaseAgent
    from beaver_gateway.agents.claude import ClaudeAgent
    from beaver_gateway.core.events import MessageStreamEvent


_log = logging.getLogger("beaver_gateway.backends.claude_sdk")

__all__ = [
    "ClaudeSdkBackend",
    "RunnerConfig",
    "SessionClient",
    "UsageSink",
    "fingerprint",
]

ENV_KEEP: tuple[str, ...] = (
    "PATH",
    "HOME",
    "LANG",
    "LANGUAGE",
    "LC_ALL",
    "LC_CTYPE",
    "TZ",
    "TERM",
    "USER",
    "LOGNAME",
    "SHELL",
    "TMPDIR",
    "PWD",
    "SSL_CERT_FILE",
    "SSL_CERT_DIR",
    "NODE_EXTRA_CA_CERTS",
    "NODE_OPTIONS",
    "IS_SANDBOX",
    "HTTP_PROXY",
    "HTTPS_PROXY",
    "NO_PROXY",
    "http_proxy",
    "https_proxy",
    "no_proxy",
)
ENV_KEEP_PREFIXES: tuple[str, ...] = ("CLAUDE_", "ANTHROPIC_", "DISABLE_")

_REAP_INTERVAL = 60.0
_STOP_REASONS: dict[str, StopReason] = {
    "end_turn": "end_turn",
    "tool_use": "tool_use",
    "max_tokens": "max_tokens",
    "stop_sequence": "stop_sequence",
    "refusal": "refusal",
    "pause_turn": "pause_turn",
}


class SessionClient(Protocol):
    async def connect(self) -> None: ...
    async def query(self, prompt: str) -> None: ...
    def receive_response(self) -> AsyncIterator[Any]: ...
    async def disconnect(self) -> None: ...


ClientFactory = "Callable[[ClaudeAgentOptions], SessionClient]"
UsageSink = "Callable[[UsageEvent], Awaitable[None]]"


@dataclass(frozen=True, slots=True)
class RunnerConfig:
    user: str | None = None
    home: Path | None = None


@dataclass(frozen=True, slots=True)
class UsageEvent:
    agent_name: str
    model: str
    effort: str | None
    conversation_id: str | None
    session_id: str | None
    usage: TurnUsage


@dataclass
class _Live:
    client: SessionClient
    session_id: str | None
    resumed: bool
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    last_used: float = field(default_factory=time.monotonic)
    turns: int = 0


class _RunnerClient(ClaudeSDKClient):
    """``ClaudeSDKClient`` that hands the materialized resume dir to the runner uid."""

    def __init__(self, options: ClaudeAgentOptions, *, uid: int | None) -> None:
        super().__init__(options=options)
        self._runner_uid = uid

    async def _connect_inner(self, prompt: Any, actual_prompt: Any) -> None:
        materialized = self._materialized
        if materialized is not None and self._runner_uid is not None:
            _chown_tree(materialized.config_dir, self._runner_uid)
        await super()._connect_inner(prompt, actual_prompt)


class ClaudeSdkBackend:
    def __init__(
        self,
        *,
        agent: ClaudeAgent,
        mcp_internal_urls: Mapping[str, str],
        session_store: SessionStore,
        mcp_tool_names: Mapping[str, Sequence[str]] | None = None,
        runner: RunnerConfig | None = None,
        usage_sink: Callable[[UsageEvent], Awaitable[None]] | None = None,
        client_factory: Callable[[ClaudeAgentOptions], SessionClient] | None = None,
        work_dir: Path | None = None,
    ) -> None:
        self._agent = agent
        self._store = session_store
        self._runner = runner or RunnerConfig()
        self._usage_sink = usage_sink
        self._factory = client_factory or self._default_factory
        self._work_dir = work_dir or Path(tempfile.gettempdir()) / "beaver-claude"
        self._servers = _mcp_servers(agent, mcp_internal_urls)
        self._mcp_disallowed = _mcp_disallowed(agent, mcp_tool_names or {})
        self._sessions: dict[str, _Live] = {}
        self._reaper: asyncio.Task[None] | None = None
        self._uid = _resolve_uid(self._runner.user)
        self._wrapper: Path | None = None

    @property
    def agent(self) -> ClaudeAgent:
        return self._agent

    @property
    def sessions(self) -> dict[str, dict[str, Any]]:
        now = time.monotonic()
        return {
            key: {
                "session_id": live.session_id,
                "idle_seconds": now - live.last_used,
                "turns": live.turns,
                "busy": live.lock.locked(),
            }
            for key, live in self._sessions.items()
        }

    async def __aenter__(self) -> Self:
        self._reaper = asyncio.create_task(self._reap_loop())
        return self

    async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
        await self.aclose()

    async def aclose(self) -> None:
        if self._reaper is not None:
            self._reaper.cancel()
            with contextlib.suppress(BaseException):
                await self._reaper
            self._reaper = None
        for key in list(self._sessions):
            await self._close(key)

    async def complete(
        self,
        *,
        agent: BaseAgent,
        messages: Iterable[MessageParam],
        system: str | None = None,  # noqa: ARG002 - the agent owns its prompt
        conversation_id: str | None = None,
        session_id: str | None = None,
        capture: TurnCapture | None = None,
        **options: Any,  # noqa: ARG002 - per-request knobs are not supported
    ) -> AsyncIterator[MessageStreamEvent]:
        if agent.name != self._agent.name:
            msg = f"backend bound to {self._agent.name!r}, got {agent.name!r}"
            raise ValueError(msg)
        history = [dict(m) for m in messages]
        if not history or history[-1].get("role") != "user":
            msg = "the last message must be a user turn"
            raise ValueError(msg)
        prompt = _prompt_text(history[-1].get("content"))
        prior = history[:-1]
        key = conversation_id or fingerprint(prior)
        live = await self._acquire(key, session_id=session_id, history=prior)
        message_id = f"msg_{uuid.uuid4().hex}"
        yield build_message_start(message_id=message_id, model=self._agent.model)
        async with live.lock:
            live.last_used = time.monotonic()
            try:
                turn = await self._run_turn(live, prompt)
            except Exception:
                if live.resumed and live.turns == 0:
                    _log.exception(
                        "resume of %s failed, reseeding from history", live.session_id
                    )
                    await self._close(key)
                    live = await self._acquire(key, session_id=None, history=prior)
                    async with live.lock:
                        turn = await self._run_turn(live, prompt)
                else:
                    raise
            for event in turn.events:
                yield event
            live.turns += 1
            live.last_used = time.monotonic()
        if turn.result is not None and turn.result.session_id:
            live.session_id = turn.result.session_id
        if conversation_id is None:
            self._rekey(key, fingerprint([*history, *turn.synthesized]))
        usage = _usage_of(turn.result)
        if capture is not None:
            capture.synthesized_messages = turn.synthesized
            capture.session_id = live.session_id
            capture.usage = usage
        if self._usage_sink is not None:
            await self._usage_sink(
                UsageEvent(
                    agent_name=self._agent.name,
                    model=self._agent.model,
                    effort=self._agent.options.effort,
                    conversation_id=conversation_id,
                    session_id=live.session_id,
                    usage=usage,
                )
            )
        if turn.result is not None and turn.result.is_error:
            msg = f"claude: {turn.result.result or turn.result.subtype}"
            raise RuntimeError(msg)
        yield build_message_delta(
            stop_reason=turn.stop_reason, usage=_wire_usage(usage)
        )
        yield build_message_stop()

    async def _run_turn(self, live: _Live, prompt: str) -> _Turn:
        streaming = self._agent.options.include_partial_messages
        turn = _Turn()
        raw: list[Any] = []
        next_index = 0
        offset = 0
        await live.client.query(prompt)
        async for message in live.client.receive_response():
            if getattr(message, "parent_tool_use_id", None) is not None:
                continue
            if isinstance(message, StreamEvent):
                event = message.event
                if event.get("type") == "message_start":
                    offset = next_index
                    continue
                index = event.get("index")
                if isinstance(index, int):
                    next_index = max(next_index, offset + index + 1)
                    if streaming:
                        turn.events.extend(_emit_stream_event(event, offset + index))
            elif isinstance(message, AssistantMessage):
                raw.append(message)
                if not streaming:
                    for block in message.content:
                        turn.events.extend(_emit_block(block, next_index))
                        next_index += 1
            elif isinstance(message, UserMessage):
                raw.append(message)
            elif isinstance(message, ResultMessage):
                turn.result = message
                turn.stop_reason = _STOP_REASONS.get(
                    message.stop_reason or "", "end_turn"
                )
        turn.synthesized = synthesize_turn_messages(raw)
        _log.info(
            "turn: agent=%s session=%s events=%d synthesized=%d stop=%s",
            self._agent.name,
            live.session_id,
            len(turn.events),
            len(turn.synthesized),
            turn.stop_reason,
        )
        return turn

    async def _acquire(
        self, key: str, *, session_id: str | None, history: list[dict[str, Any]]
    ) -> _Live:
        live = self._sessions.get(key)
        if live is not None:
            return live
        resume = session_id
        if resume is None and history:
            resume = await self._seed(history)
        live = await self._spawn(resume)
        self._sessions[key] = live
        return live

    async def _seed(self, history: list[dict[str, Any]]) -> str:
        session_id = str(uuid.uuid4())
        entries = build_entries(
            history,
            session_id=session_id,
            cwd=str(self._agent.cwd),
            model=self._agent.model,
            permission_mode=self._agent.options.permission_mode,
        )
        key = {
            "project_key": project_key_for_directory(str(self._agent.cwd)),
            "session_id": session_id,
        }
        await self._store.append(cast("Any", key), cast("Any", entries))
        _log.info(
            "seeded session %s with %d entries from %d messages",
            session_id,
            len(entries),
            len(history),
        )
        return session_id

    async def _spawn(self, resume: str | None) -> _Live:
        options = self._build_options(resume)
        client = self._factory(options)
        await client.connect()
        _log.info(
            "spawned claude: agent=%s resume=%s user=%s",
            self._agent.name,
            resume,
            self._runner.user,
        )
        return _Live(client=client, session_id=resume, resumed=resume is not None)

    def _default_factory(self, options: ClaudeAgentOptions) -> SessionClient:
        return _RunnerClient(options, uid=self._uid)

    def _build_options(self, resume: str | None) -> ClaudeAgentOptions:
        agent = self._agent
        opt = agent.options
        env = dict(opt.env)
        if self._runner.home is not None:
            env["HOME"] = str(self._runner.home)
        plugins = self._plugins()
        system_prompt = (
            prompt_assembly.assemble(agent.prompt_sources)
            if agent.prompt_sources
            else agent.system_prompt
        )
        return ClaudeAgentOptions(
            model=agent.model or None,
            effort=cast("Any", opt.effort),
            system_prompt=system_prompt,
            setting_sources=[],
            strict_mcp_config=True,
            mcp_servers=cast("Any", self._servers),
            permission_mode=cast("Any", opt.permission_mode),
            tools=list(opt.tools) if opt.tools is not None else None,
            disallowed_tools=[*opt.disallowed_tools, *self._mcp_disallowed],
            cwd=str(agent.cwd),
            add_dirs=list(opt.add_dirs),
            env=env,
            user=self._runner.user,
            cli_path=str(self._exec_wrapper(extra_keep=tuple(env))),
            include_partial_messages=opt.include_partial_messages,
            session_store=self._store,
            session_store_flush=cast("Any", opt.session_store_flush),
            resume=resume,
            plugins=cast("Any", plugins),
            skills="all" if plugins else None,
            max_turns=opt.max_turns,
            stderr=lambda line: _log.warning(
                "claude[%s]: %s", agent.name, line.rstrip()
            ),
        )

    def _plugins(self) -> list[dict[str, str]]:
        plugins: list[dict[str, str]] = []
        root = self._work_dir / "plugins" / self._agent.name
        for raw in sorted(self._agent.skill_sets, key=str):
            source = Path(str(raw))
            name = source.name
            target = root / name
            if target.exists():
                shutil.rmtree(target)
            shutil.copytree(source, target / "skills")
            (target / ".claude-plugin").mkdir(parents=True, exist_ok=True)
            (target / ".claude-plugin" / "plugin.json").write_text(
                json.dumps({"name": name, "version": "0.0.0"}), encoding="utf-8"
            )
            _chmod_tree(target)
            plugins.append({"type": "local", "path": str(target)})
        return plugins

    def _exec_wrapper(self, *, extra_keep: tuple[str, ...]) -> Path:
        if self._wrapper is not None:
            return self._wrapper
        keep = sorted({*ENV_KEEP, *self._agent.options.env_keep, *extra_keep})
        target = _claude_binary()
        script = _WRAPPER.format(
            python=sys.executable,
            target=json.dumps(target),
            keep=json.dumps(keep),
            prefixes=json.dumps(list(ENV_KEEP_PREFIXES)),
        )
        digest = hashlib.sha256(script.encode("utf-8")).hexdigest()[:12]
        path = self._work_dir / f"claude-exec-{digest}.py"
        path.parent.mkdir(parents=True, exist_ok=True)
        if not path.exists() or path.read_text(encoding="utf-8") != script:
            path.write_text(script, encoding="utf-8")
        path.chmod(0o755)
        path.parent.chmod(0o755)
        self._wrapper = path
        return path

    def _rekey(self, old: str, new: str) -> None:
        live = self._sessions.pop(old, None)
        if live is None:
            return
        stale = self._sessions.pop(new, None)
        self._sessions[new] = live
        if stale is not None and stale is not live:
            asyncio.get_running_loop().create_task(_disconnect(stale))

    async def _close(self, key: str) -> None:
        live = self._sessions.pop(key, None)
        if live is not None:
            await _disconnect(live)

    async def _reap_loop(self) -> None:
        ttl = self._agent.options.idle_session_ttl
        while True:
            await asyncio.sleep(_REAP_INTERVAL)
            if ttl <= 0:
                continue
            now = time.monotonic()
            for key, live in list(self._sessions.items()):
                if not live.lock.locked() and now - live.last_used > ttl:
                    _log.info("closing idle session %s (%s)", live.session_id, key)
                    await self._close(key)


@dataclass
class _Turn:
    events: list[Any] = field(default_factory=list)
    synthesized: list[dict[str, Any]] = field(default_factory=list)
    result: ResultMessage | None = None
    stop_reason: StopReason = "end_turn"


_WRAPPER = """#!{python}
import os
import sys

TARGET = {target}
KEEP = set({keep})
PREFIXES = tuple({prefixes})
env = {{
    k: v for k, v in os.environ.items() if k in KEEP or k.startswith(PREFIXES)
}}
os.execve(TARGET, [TARGET, *sys.argv[1:]], env)
"""


async def _disconnect(live: _Live) -> None:
    try:
        await live.client.disconnect()
    except Exception:  # noqa: BLE001
        _log.exception("disconnect failed for session %s", live.session_id)


def _claude_binary() -> str:
    bundled = Path(claude_agent_sdk.__file__).parent / "_bundled" / "claude"
    if bundled.is_file():
        return str(bundled)
    found = shutil.which("claude")
    if found is None:
        msg = "claude CLI not found: neither bundled in claude_agent_sdk nor on PATH"
        raise FileNotFoundError(msg)
    return found


def _resolve_uid(user: str | None) -> int | None:
    if user is None:
        return None
    if user.isdigit():
        return int(user)
    return pwd.getpwnam(user).pw_uid


def _chown_tree(root: Path, uid: int) -> None:
    for path in [root, *root.rglob("*")]:
        with contextlib.suppress(OSError):
            os.chown(path, uid, -1)


def _chmod_tree(root: Path) -> None:
    for path in [root, *root.rglob("*")]:
        with contextlib.suppress(OSError):
            path.chmod(0o755 if path.is_dir() else 0o644)


def _mcp_servers(
    agent: ClaudeAgent, urls: Mapping[str, str]
) -> dict[str, dict[str, Any]]:
    servers: dict[str, dict[str, Any]] = {}
    for exposed in agent.expose_mcps:
        url = urls.get(exposed.name)
        if url is None:
            msg = f"agent {agent.name!r} exposes MCP {exposed.name!r} without a URL"
            raise ValueError(msg)
        servers[exposed.name] = {"type": "http", "url": url}
    return servers


def _mcp_disallowed(
    agent: ClaudeAgent, catalog: Mapping[str, Sequence[str]]
) -> list[str]:
    out: list[str] = []
    for exposed in agent.expose_mcps:
        if exposed.tools is None and not exposed.deny:
            continue
        names = list(catalog.get(exposed.name, ()))
        if not names:
            _log.warning(
                "MCP %r has no tool catalog; disallowing the whole server for %s",
                exposed.name,
                agent.name,
            )
            out.append(f"mcp__{exposed.name}")
            continue
        for name in names:
            allowed = exposed.tools is None or name in exposed.tools
            denied = any(fnmatch.fnmatchcase(name, pat) for pat in exposed.deny)
            if not allowed or denied:
                out.append(f"mcp__{exposed.name}__{name}")
    return out


def fingerprint(messages: Iterable[Mapping[str, Any]]) -> str:
    turns: list[tuple[str, str]] = []
    for message in messages:
        text = _text_of(message.get("content"))
        if not text:
            continue
        role = str(message.get("role", ""))
        if turns and turns[-1][0] == role:
            turns[-1] = (role, turns[-1][1] + "\n" + text)
        else:
            turns.append((role, text))
    digest = hashlib.sha1(usedforsecurity=False)
    for role, text in turns:
        digest.update(role.encode("utf-8"))
        digest.update(b"\x00")
        digest.update(text.strip().encode("utf-8"))
        digest.update(b"\x01")
    return digest.hexdigest()


def _text_of(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = [
            str(b.get("text", ""))
            for b in content
            if isinstance(b, Mapping) and b.get("type") == "text"
        ]
        return "\n".join(p for p in parts if p)
    return ""


def _prompt_text(content: Any) -> str:
    text = _text_of(content)
    if not text:
        msg = "user message has no text content"
        raise ValueError(msg)
    return text


def synthesize_turn_messages(raw: Iterable[Any]) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    for message in raw:
        if isinstance(message, AssistantMessage):
            out.append(
                {
                    "role": "assistant",
                    "content": [_block_to_dict(b) for b in message.content],
                }
            )
        elif isinstance(message, UserMessage):
            content = message.content
            if isinstance(content, list) and content:
                out.append(
                    {"role": "user", "content": [_block_to_dict(b) for b in content]}
                )
    return out


def _block_to_dict(block: Any) -> dict[str, Any]:
    if isinstance(block, TextBlock):
        return {"type": "text", "text": block.text}
    if isinstance(block, ToolUseBlock):
        return {
            "type": "tool_use",
            "id": block.id,
            "name": block.name,
            "input": block.input,
        }
    if isinstance(block, ToolResultBlock):
        result: dict[str, Any] = {
            "type": "tool_result",
            "tool_use_id": block.tool_use_id,
            "content": block.content,
        }
        if block.is_error is not None:
            result["is_error"] = block.is_error
        return result
    if isinstance(block, ThinkingBlock):
        return {
            "type": "thinking",
            "thinking": block.thinking,
            "signature": block.signature,
        }
    msg = f"unknown content block type: {type(block).__name__}"
    raise TypeError(msg)


def _usage_of(result: ResultMessage | None) -> TurnUsage:
    if result is None:
        return TurnUsage()
    usage = result.usage or {}
    return TurnUsage(
        input_tokens=_int(usage.get("input_tokens")),
        output_tokens=_int(usage.get("output_tokens")),
        cache_read_tokens=_int(usage.get("cache_read_input_tokens")),
        cache_creation_tokens=_int(usage.get("cache_creation_input_tokens")),
        cost_usd=result.total_cost_usd,
        duration_ms=result.duration_ms,
        num_turns=result.num_turns,
    )


def _wire_usage(usage: TurnUsage) -> dict[str, int]:
    return {
        "input_tokens": usage.input_tokens,
        "output_tokens": usage.output_tokens,
        "cache_read_input_tokens": usage.cache_read_tokens,
        "cache_creation_input_tokens": usage.cache_creation_tokens,
    }


def _int(value: Any) -> int:
    return value if isinstance(value, int) else 0


def _emit_stream_event(
    event: Mapping[str, Any], index: int
) -> Iterable[MessageStreamEvent]:
    etype = event.get("type")
    if etype == "content_block_start":
        block = event.get("content_block")
        if not isinstance(block, Mapping):
            return ()
        btype = block.get("type")
        if btype == "text":
            return (build_text_block_start(index),)
        if btype == "thinking":
            return (build_thinking_block_start(index),)
        if btype == "tool_use":
            return (
                build_tool_use_block_start(
                    index,
                    tool_use_id=str(block.get("id", "")),
                    name=str(block.get("name", "")),
                ),
            )
        return ()
    if etype == "content_block_delta":
        delta = event.get("delta")
        if not isinstance(delta, Mapping):
            return ()
        dtype = delta.get("type")
        if dtype == "text_delta":
            return (build_text_delta(index, str(delta.get("text", ""))),)
        if dtype == "thinking_delta":
            return (build_thinking_delta(index, str(delta.get("thinking", ""))),)
        if dtype == "signature_delta":
            return (build_signature_delta(index, str(delta.get("signature", ""))),)
        if dtype == "input_json_delta":
            return (build_input_json_delta(index, str(delta.get("partial_json", ""))),)
        return ()
    if etype == "content_block_stop":
        return (build_content_block_stop(index),)
    return ()


def _emit_block(block: Any, index: int) -> Iterable[MessageStreamEvent]:
    if isinstance(block, TextBlock):
        return (
            build_text_block_start(index),
            build_text_delta(index, block.text),
            build_content_block_stop(index),
        )
    if isinstance(block, ThinkingBlock):
        return (
            build_thinking_block_start(index),
            build_thinking_delta(index, block.thinking),
            build_signature_delta(index, block.signature),
            build_content_block_stop(index),
        )
    if isinstance(block, ToolUseBlock):
        partial = json.dumps(block.input, separators=(",", ":"), ensure_ascii=False)
        return (
            build_tool_use_block_start(index, tool_use_id=block.id, name=block.name),
            build_input_json_delta(index, partial),
            build_content_block_stop(index),
        )
    return ()