"""Anthropic backend — direct API calls with tool_choice schema enforcement.

Reused from swarm with import path adaptation.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
from typing import TypeVar

from pydantic import BaseModel, ValidationError

from pact.budget import BudgetExceeded, BudgetTracker

logger = logging.getLogger(__name__)

T = TypeVar("T", bound=BaseModel)

_MODEL_MAX_TOKENS: dict[str, int] = {
    "claude-opus-4-6": 32768,
    "claude-sonnet-4-5-20250929": 64000,
    "claude-haiku-4-5-20251001": 8192,
}
_DEFAULT_MAX_TOKENS_CAP = 32768


class AnthropicBackend:
    """Backend using the Anthropic API with tool_choice for structured extraction."""

    def __init__(self, budget: BudgetTracker, model: str = "claude-opus-4-6") -> None:
        try:
            import anthropic
        except ImportError as exc:
            raise ImportError(
                "The 'anthropic' package is required. Install with: pip install -e '.[cli]'"
            ) from exc

        api_key = os.environ.get("ANTHROPIC_API_KEY", "")
        if not api_key:
            raise ValueError("ANTHROPIC_API_KEY environment variable is required.")

        self._client = anthropic.AsyncAnthropic(
            api_key=api_key,
            max_retries=3,
            timeout=600.0,
        )
        self._model = model
        self._budget = budget

    def set_model(self, model: str) -> None:
        self._model = model

    def _max_tokens_cap(self) -> int:
        return _MODEL_MAX_TOKENS.get(self._model, _DEFAULT_MAX_TOKENS_CAP)

    async def assess(
        self,
        schema: type[T],
        prompt: str,
        system: str,
        max_tokens: int = 32768,
    ) -> tuple[T, int, int]:
        """Call LLM with schema enforcement via tool_choice."""
        total_in = 0
        total_out = 0
        cap = self._max_tokens_cap()
        current_max = min(max_tokens, cap)

        for attempt in range(3):
            raw_input, stop_reason, in_tok, out_tok = await self._call_llm(
                schema, prompt, system, current_max,
            )
            total_in += in_tok
            total_out += out_tok

            if raw_input is None:
                raise RuntimeError(
                    f"No tool_use block found for {schema.__name__}"
                )

            if stop_reason == "max_tokens" and attempt < 2:
                new_max = min(current_max * 2, cap)
                if new_max > current_max:
                    current_max = new_max
                    continue

            raw_input = self._coerce_fields(raw_input)

            try:
                parsed = schema.model_validate(raw_input)
                return parsed, total_in, total_out
            except ValidationError as e:
                if attempt < 2:
                    new_max = min(current_max * 2, cap)
                    if new_max > current_max:
                        current_max = new_max
                    continue
                raise

        raise RuntimeError(f"Failed to get valid {schema.__name__} after 3 attempts")

    @classmethod
    def _coerce_fields(cls, data):
        """Recursively parse string fields that look like JSON.

        LLMs sometimes return structured fields (lists, dicts) as JSON
        strings instead of native objects.  This walks the entire tree
        so nested structures are also fixed.
        """
        if isinstance(data, dict):
            return {k: cls._coerce_fields(v) for k, v in data.items()}
        if isinstance(data, list):
            return [cls._coerce_fields(item) for item in data]
        if isinstance(data, str):
            stripped = data.strip()
            if stripped and stripped[0] in ("[", "{"):
                try:
                    parsed = json.loads(stripped)
                    return cls._coerce_fields(parsed)
                except (json.JSONDecodeError, ValueError):
                    pass
        return data

    async def _call_llm(
        self,
        schema: type[T],
        prompt: str,
        system: str,
        max_tokens: int,
        stall_timeout: float = 300.0,
    ) -> tuple[dict | None, str, int, int]:
        """Call LLM with streaming progress detection.

        Uses streaming so we can distinguish a stalled connection
        (no events for stall_timeout seconds) from a legitimately
        long generation that's actively producing tokens.
        """
        tool_name = schema.__name__
        tool_schema = schema.model_json_schema()
        tool_schema.pop("title", None)

        try:
            message = await self._stream_with_stall_detection(
                tool_name, tool_schema, schema.__doc__ or f"Extract {tool_name}",
                prompt, system, max_tokens, stall_timeout,
            )
        except asyncio.TimeoutError:
            logger.error(
                "Anthropic API stalled (no progress for %.0fs) for %s",
                stall_timeout, tool_name,
            )
            raise RuntimeError(
                f"Anthropic API stalled (no progress for {stall_timeout:.0f}s) for {tool_name}"
            )

        in_tok = message.usage.input_tokens
        out_tok = message.usage.output_tokens
        stop_reason = message.stop_reason or ""

        if not self._budget.record_tokens(in_tok, out_tok):
            raise BudgetExceeded(f"Budget exceeded after {in_tok}+{out_tok} tokens")

        for block in message.content:
            if block.type == "tool_use" and block.name == tool_name:
                return block.input, stop_reason, in_tok, out_tok

        return None, stop_reason, in_tok, out_tok

    async def _stream_with_stall_detection(
        self,
        tool_name: str,
        tool_schema: dict,
        tool_description: str,
        prompt: str,
        system: str,
        max_tokens: int,
        stall_timeout: float,
    ):
        """Stream a response, raising TimeoutError if no event arrives within stall_timeout.

        Unlike a hard timeout on the full request, this only fires when the
        connection goes silent — a 10-minute generation that's actively
        streaming tokens will never trigger it.
        """
        async with self._client.messages.stream(
            model=self._model,
            max_tokens=max_tokens,
            system=system,
            messages=[{"role": "user", "content": prompt}],
            tools=[{
                "name": tool_name,
                "description": tool_description,
                "input_schema": tool_schema,
            }],
            tool_choice={"type": "tool", "name": tool_name},
        ) as stream:
            aiter = stream.__aiter__()
            while True:
                try:
                    await asyncio.wait_for(aiter.__anext__(), timeout=stall_timeout)
                except StopAsyncIteration:
                    break
                except asyncio.TimeoutError:
                    raise asyncio.TimeoutError()

        return await stream.get_final_message()

    # ── Prompt caching helpers ──────────────────────────────────────────

    _CACHE_MIN_CHARS = 300
    _CACHE_MIN_SYSTEM_CHARS = 4  # system prompts are almost always reused

    def _build_system_blocks(self, system: str, *, cache: bool = False) -> list[dict]:
        """Convert system string to content block list.

        If *cache* is True and the text is long enough, the block gets
        ``cache_control: {"type": "ephemeral"}``.  Very short system
        strings (< 4 chars) skip cache_control since the overhead is
        not worthwhile.
        """
        block: dict = {"type": "text", "text": system}
        if cache and len(system) >= self._CACHE_MIN_SYSTEM_CHARS:
            block["cache_control"] = {"type": "ephemeral"}
        return [block]

    def _build_user_blocks(
        self, cache_prefix: str, prompt: str
    ) -> str | list[dict]:
        """Build user content: prefix block (optionally cached) + dynamic prompt.

        Returns a plain string when there is no prefix, avoiding unnecessary
        overhead.
        """
        if not cache_prefix:
            return prompt

        prefix_block: dict = {"type": "text", "text": cache_prefix}
        if len(cache_prefix) >= self._CACHE_MIN_CHARS:
            prefix_block["cache_control"] = {"type": "ephemeral"}

        prompt_block: dict = {"type": "text", "text": prompt}
        return [prefix_block, prompt_block]

    # ── Cached LLM call ──────────────────────────────────────────────

    async def _call_llm_cached(
        self,
        schema: type[T],
        prompt: str,
        system: str,
        cache_prefix: str,
        max_tokens: int,
        stall_timeout: float = 300.0,
    ) -> tuple[dict | None, str, int, int]:
        """Like _call_llm but sends system + cache_prefix with cache_control."""
        tool_name = schema.__name__
        tool_schema = schema.model_json_schema()
        tool_schema.pop("title", None)

        system_blocks = self._build_system_blocks(system, cache=True)
        user_content = self._build_user_blocks(cache_prefix, prompt)

        try:
            async with self._client.messages.stream(
                model=self._model,
                max_tokens=max_tokens,
                system=system_blocks,
                messages=[{"role": "user", "content": user_content}],
                tools=[{
                    "name": tool_name,
                    "description": schema.__doc__ or f"Extract {tool_name}",
                    "input_schema": tool_schema,
                }],
                tool_choice={"type": "tool", "name": tool_name},
            ) as stream:
                aiter = stream.__aiter__()
                while True:
                    try:
                        await asyncio.wait_for(aiter.__anext__(), timeout=stall_timeout)
                    except StopAsyncIteration:
                        break
                    except asyncio.TimeoutError:
                        raise asyncio.TimeoutError()

            message = await stream.get_final_message()
        except asyncio.TimeoutError:
            logger.error(
                "Anthropic API stalled (no progress for %.0fs) for %s",
                stall_timeout, tool_name,
            )
            raise RuntimeError(
                f"Anthropic API stalled (no progress for {stall_timeout:.0f}s) for {tool_name}"
            )

        in_tok = message.usage.input_tokens
        out_tok = message.usage.output_tokens
        stop_reason = message.stop_reason or ""

        # Record cache metrics
        cache_creation = getattr(message.usage, 'cache_creation_input_tokens', 0) or 0
        cache_read = getattr(message.usage, 'cache_read_input_tokens', 0) or 0
        if cache_creation or cache_read:
            self._budget.record_cache_tokens(cache_creation, cache_read)

        if not self._budget.record_tokens(in_tok, out_tok):
            raise BudgetExceeded(f"Budget exceeded after {in_tok}+{out_tok} tokens")

        for block in message.content:
            if block.type == "tool_use" and block.name == tool_name:
                return block.input, stop_reason, in_tok, out_tok

        return None, stop_reason, in_tok, out_tok

    # ── Public cached assess ─────────────────────────────────────────

    async def assess_with_cache(
        self,
        schema: type[T],
        prompt: str,
        system: str,
        cache_prefix: str = "",
        max_tokens: int = 32768,
    ) -> tuple[T, int, int]:
        """Like assess() but marks system + cache_prefix for prompt caching.

        Args:
            system: System prompt -- always cached.
            cache_prefix: Static portion of user prompt (SOPs, contracts, etc.).
                          Sent as a separate content block with cache_control.
            prompt: Dynamic portion of user prompt (attempt-specific).
        """
        total_in = 0
        total_out = 0
        cap = self._max_tokens_cap()
        current_max = min(max_tokens, cap)

        for attempt in range(3):
            raw_input, stop_reason, in_tok, out_tok = await self._call_llm_cached(
                schema, prompt, system, cache_prefix, current_max,
            )
            total_in += in_tok
            total_out += out_tok

            if raw_input is None:
                raise RuntimeError(
                    f"No tool_use block found for {schema.__name__}"
                )

            if stop_reason == "max_tokens" and attempt < 2:
                new_max = min(current_max * 2, cap)
                if new_max > current_max:
                    current_max = new_max
                    continue

            raw_input = self._coerce_fields(raw_input)

            try:
                parsed = schema.model_validate(raw_input)
                return parsed, total_in, total_out
            except ValidationError as e:
                if attempt < 2:
                    new_max = min(current_max * 2, cap)
                    if new_max > current_max:
                        current_max = new_max
                    continue
                raise

        raise RuntimeError(f"Failed to get valid {schema.__name__} after 3 attempts")

    async def close(self) -> None:
        await self._client.close()
