"""Anthropic implementation of structured intent and streamed reranking."""

import json
import re
from collections.abc import AsyncGenerator
from typing import Literal, cast

from anthropic import AsyncAnthropic
from anthropic.lib._parse._transform import transform_schema
from anthropic.types import Message, TextBlockParam, Usage
from pydantic import BaseModel, ConfigDict, Field, ValidationError

from app.config import Settings
from app.domain.models import (
    ConversationTurn,
    Familiarity,
    Intent,
    PreviousRecommendation,
    RerankSelection,
    Track,
    TrackCandidate,
)
from app.observability.timing import record_llm_tokens
from app.ports.protocols import RecommenderOutputError
from app.prompts import INTENT_SYSTEM_PROMPT, RERANK_SYSTEM_PROMPT

Effort = Literal["low", "medium", "high", "xhigh", "max"]
_RECOMMENDATION_ARRAY = re.compile(r'"recommendations"\s*:\s*\[')


class CandidateOutput(BaseModel):
    """One bounded candidate in the intent response."""

    model_config = ConfigDict(extra="forbid")

    title: str = Field(min_length=1, max_length=200)
    artist: str = Field(min_length=1, max_length=200)


class IntentOutput(BaseModel):
    """Bounded structured output for the intent call."""

    model_config = ConfigDict(extra="forbid")

    mood: list[str] = Field(default_factory=list, max_length=6)
    activity: str | None = Field(default=None, max_length=100)
    era: list[str] = Field(default_factory=list, max_length=5)
    languages: list[str] = Field(default_factory=list, max_length=8)
    genres: list[str] = Field(default_factory=list, max_length=8)
    familiarity: Familiarity
    is_refinement: bool
    intent_summary: str = Field(min_length=1, max_length=300, pattern=r"^[^\r\n]+$")
    candidates: list[CandidateOutput] = Field(min_length=10, max_length=50)


class RerankSelectionOutput(BaseModel):
    """One validated object extracted from the rerank stream."""

    model_config = ConfigDict(extra="forbid")

    track_id: str = Field(min_length=1, max_length=200)
    justification: str = Field(min_length=1, max_length=300, pattern=r"^[^\r\n]+$")


class RerankOutput(BaseModel):
    """Complete bounded rerank output used for final validation."""

    model_config = ConfigDict(extra="forbid")

    recommendations: list[RerankSelectionOutput] = Field(min_length=1, max_length=50)


class AnthropicRecommender:
    """Use two Anthropic calls for intent generation and grounded ranking."""

    def __init__(self, client: AsyncAnthropic, settings: Settings) -> None:
        """Bind the shared asynchronous client and immutable settings."""
        self.client = client
        self.settings = settings

    async def create_intent(
        self,
        query: str,
        history: tuple[ConversationTurn, ...],
        previous_recommendations: tuple[PreviousRecommendation, ...],
        taste_summary: str,
        candidate_count: int,
    ) -> Intent:
        """Interpret a request through Anthropic structured output."""
        message = await self.client.messages.parse(
            model=self.settings.llm_model,
            max_tokens=self.settings.intent_max_tokens,
            output_config={"effort": _parse_effort(self.settings.intent_effort)},
            output_format=IntentOutput,
            system=[_system_block(INTENT_SYSTEM_PROMPT)],
            messages=[
                {
                    "role": "user",
                    "content": _render_intent_input(
                        query,
                        history,
                        previous_recommendations,
                        taste_summary,
                        candidate_count,
                    ),
                }
            ],
        )
        _validate_stop_reason(message)
        _record_usage(message.usage)
        parsed = message.parsed_output
        if parsed is None:
            raise RecommenderOutputError("Intent response contained no structured output")
        validated = IntentOutput.model_validate(parsed.model_dump())
        return _to_intent(validated)

    async def stream_rerank(
        self,
        intent: Intent,
        grounded_tracks: tuple[Track, ...],
        taste_summary: str,
        history: tuple[ConversationTurn, ...],
        selection_count: int,
        correction: str | None = None,
    ) -> AsyncGenerator[RerankSelection]:
        """Yield each complete valid selection while the JSON is streaming."""
        schema = transform_schema(RerankOutput.model_json_schema())
        parser = _RecommendationObjectParser()
        async with self.client.messages.stream(
            model=self.settings.llm_model,
            max_tokens=self.settings.rerank_max_tokens,
            output_config={
                "effort": _parse_effort(self.settings.rerank_effort),
                "format": {"type": "json_schema", "schema": schema},
            },
            system=[_system_block(RERANK_SYSTEM_PROMPT)],
            messages=[
                {
                    "role": "user",
                    "content": _render_rerank_input(
                        intent,
                        grounded_tracks,
                        taste_summary,
                        history,
                        selection_count,
                        correction,
                    ),
                }
            ],
        ) as stream:
            async for text_delta in stream.text_stream:
                for selection in parser.feed(text_delta):
                    yield RerankSelection(
                        track_id=selection.track_id,
                        justification=selection.justification,
                    )
            final_message = await stream.get_final_message()

        _validate_stop_reason(final_message)
        _record_usage(final_message.usage)
        try:
            validated = RerankOutput.model_validate_json(parser.complete_text)
        except ValidationError as error:
            raise RecommenderOutputError("Rerank response failed final validation") from error
        if len(validated.recommendations) > selection_count:
            raise RecommenderOutputError("Rerank response returned too many selections")


class _RecommendationObjectParser:
    def __init__(self) -> None:
        self.complete_text = ""
        self._scan_index = 0
        self._object_start: int | None = None
        self._object_depth = 0
        self._is_in_string = False
        self._is_escaped = False
        self._has_found_array = False

    def feed(self, text_delta: str) -> list[RerankSelectionOutput]:
        self.complete_text += text_delta
        if not self._has_found_array:
            match = _RECOMMENDATION_ARRAY.search(self.complete_text)
            if match is None:
                return []
            self._has_found_array = True
            self._scan_index = match.end()

        selections: list[RerankSelectionOutput] = []
        while self._scan_index < len(self.complete_text):
            character = self.complete_text[self._scan_index]
            completed = self._scan_character(character)
            self._scan_index += 1
            if completed is not None:
                selections.append(completed)
        return selections

    def _scan_character(self, character: str) -> RerankSelectionOutput | None:
        if self._object_start is None:
            if character == "{":
                self._object_start = self._scan_index
                self._object_depth = 1
            return None

        if self._is_in_string:
            if self._is_escaped:
                self._is_escaped = False
            elif character == "\\":
                self._is_escaped = True
            elif character == '"':
                self._is_in_string = False
            return None

        if character == '"':
            self._is_in_string = True
        elif character == "{":
            self._object_depth += 1
        elif character == "}":
            self._object_depth -= 1
            if self._object_depth == 0:
                return self._finish_object()
        return None

    def _finish_object(self) -> RerankSelectionOutput:
        assert self._object_start is not None
        object_text = self.complete_text[self._object_start : self._scan_index + 1]
        self._object_start = None
        try:
            return RerankSelectionOutput.model_validate_json(object_text)
        except ValidationError as error:
            raise RecommenderOutputError("Rerank item failed validation") from error


def _to_intent(output: IntentOutput) -> Intent:
    return Intent(
        mood=tuple(output.mood),
        activity=output.activity,
        era=tuple(output.era),
        languages=tuple(output.languages),
        genres=tuple(output.genres),
        familiarity=output.familiarity,
        is_refinement=output.is_refinement,
        intent_summary=output.intent_summary,
        candidates=tuple(
            TrackCandidate(title=candidate.title, artist=candidate.artist)
            for candidate in output.candidates
        ),
    )


def _render_intent_input(
    query: str,
    history: tuple[ConversationTurn, ...],
    previous_recommendations: tuple[PreviousRecommendation, ...],
    taste_summary: str,
    candidate_count: int,
) -> str:
    payload = {
        "query": query,
        "history": [turn.__dict__ for turn in history],
        "prior_recommendations": [
            recommendation.__dict__ for recommendation in previous_recommendations
        ],
        "taste_profile": taste_summary,
        "required_candidate_count": candidate_count,
    }
    return json.dumps(payload, ensure_ascii=True, separators=(",", ":"))


def _render_rerank_input(
    intent: Intent,
    grounded_tracks: tuple[Track, ...],
    taste_summary: str,
    history: tuple[ConversationTurn, ...],
    selection_count: int,
    correction: str | None,
) -> str:
    payload = {
        "intent": {
            "mood": intent.mood,
            "activity": intent.activity,
            "era": intent.era,
            "languages": intent.languages,
            "genres": intent.genres,
            "familiarity": intent.familiarity,
            "intent_summary": intent.intent_summary,
        },
        "grounded_pool": [
            {"track_id": track.id, "title": track.title, "artists": track.artists}
            for track in grounded_tracks
        ],
        "taste_profile": taste_summary,
        "history": [turn.__dict__ for turn in history],
        "requested_selection_count": selection_count,
        "correction": correction,
    }
    return json.dumps(payload, ensure_ascii=True, separators=(",", ":"))


def _system_block(prompt: str) -> TextBlockParam:
    return {
        "type": "text",
        "text": prompt,
        "cache_control": {"type": "ephemeral"},
    }


def _parse_effort(value: str) -> Effort:
    allowed = {"low", "medium", "high", "xhigh", "max"}
    if value not in allowed:
        raise ValueError(f"Unsupported Anthropic effort: {value}")
    return cast(Effort, value)


def _validate_stop_reason(message: Message) -> None:
    stop_reason_messages = {
        None: "Anthropic response had no stop reason",
        "max_tokens": "Anthropic response reached its token limit",
        "stop_sequence": "Anthropic response hit an unexpected stop sequence",
        "tool_use": "Anthropic response attempted tool use",
        "pause_turn": "Anthropic response paused before completion",
        "refusal": "Anthropic response was refused",
        "model_context_window_exceeded": "Anthropic context window was exceeded",
    }
    if message.stop_reason == "end_turn":
        return
    raise RecommenderOutputError(
        stop_reason_messages.get(message.stop_reason, "Anthropic response stopped unexpectedly")
    )


def _record_usage(usage: Usage) -> None:
    cache_creation_tokens = usage.cache_creation_input_tokens or 0
    cache_read_tokens = usage.cache_read_input_tokens or 0
    record_llm_tokens(
        usage.input_tokens + cache_creation_tokens + cache_read_tokens,
        usage.output_tokens,
    )