"""
Module for parsing and normalizing LLM responses across different providers.
"""

import json
from typing import TYPE_CHECKING, Any, ClassVar, Literal, Type, cast, overload
from uuid import uuid4

from httpx._types import ResponseContent
from pydantic import ConfigDict, Field

from moxn.base_models.blocks.text import TextContent
from moxn.base_models.blocks.tool import ToolCall
from moxn_types.content import Provider
from moxn_types.request_config import RequestConfig
from moxn_types.response import (
    ParsedResponseCandidateModelBase,
    ParsedResponseModelBase,
    ResponseMetadata,
    ResponseParserProtocol,
    StopReason,
    TokenUsage,
)
from moxn_types.telemetry import LLMEventModelBase, ResponseType
from moxn_types.type_aliases.anthropic import (
    AnthropicContentBlock,
    AnthropicMessage,
    AnthropicToolUseBlock,
)
from moxn_types.type_aliases.google import (
    GoogleGenerateContentResponse,
    GoogleGenerateContentResponseCandidate,
)
from moxn_types.type_aliases.openai_chat import (
    OpenAIChatCompletion,
)
from moxn_types.type_aliases.openai_responses import OpenAIResponse

if TYPE_CHECKING:
    from moxn.models.message import Message
else:
    Message = Any


class ParsedResponseCandidate(ParsedResponseCandidateModelBase[TextContent, ToolCall]):
    content_blocks: list[TextContent | ToolCall]
    metadata: ResponseMetadata


class ParsedResponse(ParsedResponseModelBase[ParsedResponseCandidate]):
    """
    Normalized response content from any LLM provider.

    Contains parsed content blocks, metadata, and original response for reference.
    """

    candidates: list[ParsedResponseCandidate]

    model_config = ConfigDict(arbitrary_types_allowed=True)


class ResponseParserAnthropic(ResponseParserProtocol[AnthropicMessage, ParsedResponse]):
    """Parser for Anthropic Claude responses."""

    @staticmethod
    def parse_candidate(
        response_content: list[AnthropicContentBlock | AnthropicToolUseBlock],
        stop_reason: StopReason,
        raw_stop_reason: str | None,
    ) -> ParsedResponseCandidate:
        """Parse a single Anthropic response candidate."""
        content_blocks: list[TextContent | ToolCall] = []

        # Parse content blocks
        if ResponseContent:
            for block in response_content:
                if block.type == "text" and block.text:
                    text_block = TextContent(text=block.text)
                    content_blocks.append(text_block)
                elif (
                    block.type == "tool_use"
                    and isinstance(block, AnthropicToolUseBlock)
                    and block.id
                    and block.name
                    and block.input
                ):
                    # Anthropic arguments are already a dict
                    tool_call = ToolCall(
                        id=block.id,
                        name=block.name,
                        arguments=cast(str | dict[str, Any] | None, block.input),
                    )
                    content_blocks.append(tool_call)

        return ParsedResponseCandidate(
            content_blocks=content_blocks,
            metadata=ResponseMetadata(
                normalized_finish_reason=stop_reason,
                raw_finish_reason=raw_stop_reason or "",
            ),
        )

    @classmethod
    def parse_response(
        cls, response: AnthropicMessage, provider: Provider = Provider.ANTHROPIC
    ) -> ParsedResponse:
        """Parse an Anthropic response into a normalized format."""
        stop_reason = StopReason.END_TURN
        raw_stop_reason = response.stop_reason

        # Anthropic stop reasons: ["end_turn", "max_tokens", "stop_sequence", "tool_use"]]
        if response.stop_reason == "end_turn":
            stop_reason = StopReason.END_TURN
        elif response.stop_reason == "tool_use":
            stop_reason = StopReason.TOOL_CALL
        elif response.stop_reason == "max_tokens":
            stop_reason = StopReason.MAX_TOKENS
        elif response.stop_reason == "stop_sequence":
            stop_reason = StopReason.END_TURN
        else:
            stop_reason = StopReason.OTHER

        # Parse candidate (Anthropic has only one candidate)
        candidates = [
            cls.parse_candidate(response.content, stop_reason, raw_stop_reason)
        ]

        # Extract token usage
        usage = TokenUsage(
            input_tokens=response.usage.input_tokens if response.usage else 0,
            completion_tokens=response.usage.output_tokens if response.usage else 0,
        )

        return ParsedResponse(
            provider=provider,
            candidates=candidates,
            stop_reason=stop_reason,
            usage=usage,
            model=response.model,
            raw_response=response.model_dump(by_alias=True, mode="json"),
        )

    @classmethod
    def extract_metadata(cls, response: AnthropicMessage) -> dict[str, Any]:
        """Extract additional metadata from Anthropic response."""
        metadata = {}

        # Add any Anthropic-specific metadata
        if hasattr(response, "id"):
            metadata["anthropic_message_id"] = response.id

        return metadata


class ResponseParserOpenAI(
    ResponseParserProtocol[OpenAIChatCompletion, ParsedResponse]
):
    """Parser for OpenAI responses."""

    @staticmethod
    def parse_candidate(message, finish_reason) -> ParsedResponseCandidate:
        """Parse a single OpenAI response candidate."""
        content_blocks: list[TextContent | ToolCall] = []

        # Map OpenAI stop reasons to our normalized format
        if finish_reason == "stop":
            stop_reason = StopReason.END_TURN
        elif finish_reason == "length":
            stop_reason = StopReason.MAX_TOKENS
        elif finish_reason in ("tool_calls", "function_call"):
            stop_reason = StopReason.TOOL_CALL
        elif finish_reason == "content_filter":
            stop_reason = StopReason.CONTENT_FILTER
        else:
            stop_reason = StopReason.OTHER

        # Parse content
        if message.content:
            content_blocks.append(TextContent(text=message.content))

        # Parse tool calls
        if message.tool_calls:
            for tool_call in message.tool_calls:
                # OpenAI tool calls have arguments as a string that needs to be parsed
                args = tool_call.function.arguments
                try:
                    # Convert string arguments to dict
                    args_dict = json.loads(args) if isinstance(args, str) else args
                except json.JSONDecodeError:
                    args_dict = {"raw_arguments": args}

                tool_call_obj = ToolCall(
                    id=tool_call.id, name=tool_call.function.name, arguments=args_dict
                )
                content_blocks.append(tool_call_obj)

        return ParsedResponseCandidate(
            content_blocks=content_blocks,
            metadata=ResponseMetadata(
                normalized_finish_reason=stop_reason, raw_finish_reason=finish_reason
            ),
        )

    @classmethod
    def parse_response(
        cls, response: OpenAIChatCompletion, provider: Provider = Provider.OPENAI_CHAT
    ) -> ParsedResponse:
        """Parse an OpenAI response into a normalized format."""
        candidates: list[ParsedResponseCandidate] = []

        # Get the choices from the response
        if not response.choices:
            return ParsedResponse(
                provider=provider,
                candidates=[],
                stop_reason=StopReason.ERROR,
                raw_response=response.model_dump(by_alias=True, mode="json"),
            )

        # Parse each choice as a candidate
        for choice in response.choices:
            candidates.append(cls.parse_candidate(choice.message, choice.finish_reason))

        # Extract token usage
        usage = TokenUsage(
            input_tokens=response.usage.prompt_tokens if response.usage else 0,
            completion_tokens=response.usage.completion_tokens if response.usage else 0,
        )

        # Use the first candidate's stop reason for the overall response
        stop_reason = (
            candidates[0].metadata.normalized_finish_reason
            if candidates
            else StopReason.ERROR
        )

        return ParsedResponse(
            provider=provider,
            candidates=candidates,
            stop_reason=stop_reason,
            usage=usage,
            model=response.model,
            raw_response=response.model_dump(by_alias=True, mode="json"),
        )

    @classmethod
    def extract_metadata(cls, response: OpenAIChatCompletion) -> dict[str, Any]:
        """Extract additional metadata from OpenAI response."""
        metadata = {}

        # Add any OpenAI-specific metadata
        if response.id:
            metadata["openai_completion_id"] = response.id

        # Add system fingerprint if available
        if response.system_fingerprint:
            metadata["system_fingerprint"] = response.system_fingerprint

        return metadata


class ResponseParserGoogle(
    ResponseParserProtocol[GoogleGenerateContentResponse, ParsedResponse]
):
    """Parser for Google Gemini and Vertex responses."""

    @staticmethod
    def parse_candidate(
        candidate: GoogleGenerateContentResponseCandidate,
    ) -> ParsedResponseCandidate:
        content_blocks: list[TextContent | ToolCall] = []
        finish_reason = (
            candidate.finish_reason.value.upper() if candidate.finish_reason else ""
        )
        if "STOP" in finish_reason:
            stop_reason = StopReason.END_TURN
        elif "MAX_TOKENS" in finish_reason:
            stop_reason = StopReason.MAX_TOKENS
        elif "SAFETY" in finish_reason or "BLOCK" in finish_reason:
            stop_reason = StopReason.CONTENT_FILTER
        elif "RECITATION" in finish_reason:
            stop_reason = StopReason.ERROR
        else:
            stop_reason = StopReason.OTHER

        # Parse content
        content = candidate.content
        if content:
            if content.parts:
                # Extract text from parts
                for part in content.parts:
                    if part.text:
                        text_content = TextContent(text=part.text)
                        content_blocks.append(text_content)
                    if part.function_call:
                        name = part.function_call.name
                        if not name:
                            raise ValueError("Tool call name is required")
                        # Google's args are already a dict
                        args = part.function_call.args

                        tool_call = ToolCall(
                            id=str(
                                uuid4()
                            ),  # Generate a unique ID as Google doesn't provide one
                            name=name,
                            arguments=args,
                        )
                        content_blocks.append(tool_call)

        return ParsedResponseCandidate(
            content_blocks=content_blocks,
            metadata=ResponseMetadata(
                normalized_finish_reason=stop_reason, raw_finish_reason=finish_reason
            ),
        )

    @classmethod
    def parse_response(
        cls, response: GoogleGenerateContentResponse, provider: Provider
    ) -> ParsedResponse:
        """Parse a Google response into a normalized format."""
        # Ensure we have candidates
        if not response.candidates:
            return ParsedResponse(
                provider=provider,
                candidates=[],
                stop_reason=StopReason.ERROR,
                raw_response=response.model_dump(by_alias=True, mode="json"),
            )

        # Parse each candidate
        parsed_candidates = [
            cls.parse_candidate(candidate) for candidate in response.candidates
        ]

        # Extract token usage
        usage_metadata = response.usage_metadata
        if usage_metadata:
            prompt_tokens = usage_metadata.prompt_token_count or 0
            completion_tokens = usage_metadata.candidates_token_count or 0

            usage = TokenUsage(
                input_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )
        else:
            usage = TokenUsage()
        model = response.model_version

        return ParsedResponse(
            provider=provider,
            candidates=parsed_candidates,
            stop_reason=parsed_candidates[0].metadata.normalized_finish_reason,
            usage=usage,
            model=model,
            raw_response=response.model_dump(by_alias=True, mode="json"),
        )

    @classmethod
    def extract_metadata(
        cls, response: GoogleGenerateContentResponse
    ) -> dict[str, Any]:
        """Extract additional metadata from Google response."""
        metadata = {}

        # Add safety ratings if available
        if response.candidates:
            candidate = response.candidates[0]
            if candidate.safety_ratings:
                metadata["safety_ratings"] = candidate.safety_ratings

        return metadata


class ResponseParserOpenAIResponses(
    ResponseParserProtocol[OpenAIResponse, ParsedResponse]
):
    """Parser for OpenAI Responses API responses."""

    @staticmethod
    def parse_output_items(output: list) -> ParsedResponseCandidate:
        """Parse output items into a single candidate.

        Note: Responses API returns a single response with multiple output items,
        not multiple choices/candidates like Chat API.
        """
        content_blocks: list[TextContent | ToolCall] = []

        # Process each output item
        for item in output:
            item_type = getattr(item, "type", None)

            if item_type == "message":
                # ResponseOutputMessage - extract text content
                if hasattr(item, "content"):
                    for content_item in item.content:
                        content_type = getattr(content_item, "type", None)
                        if content_type == "output_text":
                            text = getattr(content_item, "text", "")
                            content_blocks.append(TextContent(text=text))

            elif item_type == "function_call":
                # ResponseFunctionToolCall
                call_id = getattr(item, "call_id", "")
                name = getattr(item, "name", "")
                arguments_str = getattr(item, "arguments", "{}")

                # Parse arguments from JSON string
                try:
                    args_dict = (
                        json.loads(arguments_str)
                        if isinstance(arguments_str, str)
                        else arguments_str
                    )
                except json.JSONDecodeError:
                    args_dict = {"raw_arguments": arguments_str}

                content_blocks.append(
                    ToolCall(id=call_id, name=name, arguments=args_dict)
                )

        # Default stop reason - will be updated from Response.status
        stop_reason = StopReason.END_TURN

        return ParsedResponseCandidate(
            content_blocks=content_blocks,
            metadata=ResponseMetadata(
                normalized_finish_reason=stop_reason,
                raw_finish_reason="completed",
            ),
        )

    @classmethod
    def parse_response(
        cls, response: OpenAIResponse, provider: Provider
    ) -> ParsedResponse:
        """Parse an OpenAI Responses API response into normalized format."""
        # Parse the single output array (not multiple choices)
        candidate = cls.parse_output_items(response.output)

        # Map status to stop reason
        status = getattr(response, "status", "completed")
        if status == "completed":
            stop_reason = StopReason.END_TURN
        elif status == "incomplete":
            # Check incomplete details for more specific reason
            incomplete_details = getattr(response, "incomplete_details", None)
            if incomplete_details:
                reason = getattr(incomplete_details, "reason", None)
                if reason == "max_output_tokens":
                    stop_reason = StopReason.MAX_TOKENS
                elif reason == "content_filter":
                    stop_reason = StopReason.CONTENT_FILTER
                else:
                    stop_reason = StopReason.OTHER
            else:
                stop_reason = StopReason.OTHER
        elif status == "failed":
            stop_reason = StopReason.ERROR
        else:
            stop_reason = StopReason.OTHER

        # Update candidate metadata with actual status
        candidate.metadata.normalized_finish_reason = stop_reason
        candidate.metadata.raw_finish_reason = status

        # Extract token usage
        usage_obj = getattr(response, "usage", None)
        usage = TokenUsage(
            input_tokens=getattr(usage_obj, "input_tokens", None)
            if usage_obj
            else None,
            completion_tokens=getattr(usage_obj, "output_tokens", None)
            if usage_obj
            else None,
        )

        return ParsedResponse(
            provider=provider,
            candidates=[candidate],  # Single candidate (no n parameter)
            stop_reason=stop_reason,
            usage=usage,
            model=response.model,
            raw_response=response.model_dump(by_alias=True, mode="json")
            if hasattr(response, "model_dump")
            else {},
        )

    @classmethod
    def extract_metadata(cls, response: OpenAIResponse) -> dict[str, Any]:
        """Extract OpenAI Responses API specific metadata."""
        metadata = {}
        if hasattr(response, "id"):
            metadata["openai_response_id"] = response.id
        return metadata


class ResponseParser:
    """
    Factory class to parse LLM responses based on provider type.

    Handles normalization of responses from different providers into
    a consistent ParsedResponse format.
    """

    _parsers: ClassVar[dict[Provider, Type[ResponseParserProtocol]]] = {
        Provider.ANTHROPIC: ResponseParserAnthropic,
        Provider.OPENAI_CHAT: ResponseParserOpenAI,
        Provider.OPENAI_RESPONSES: ResponseParserOpenAIResponses,
        Provider.GOOGLE_GEMINI: ResponseParserGoogle,
        Provider.GOOGLE_VERTEX: ResponseParserGoogle,
    }

    @classmethod
    def get_parser(cls, provider: Provider) -> ResponseParserProtocol:
        """Get the parser for a given provider."""
        if provider == Provider.ANTHROPIC:
            return cls._parsers[Provider.ANTHROPIC]
        elif provider == Provider.OPENAI_CHAT:
            return cls._parsers[Provider.OPENAI_CHAT]
        elif provider == Provider.OPENAI_RESPONSES:
            return cls._parsers[Provider.OPENAI_RESPONSES]
        elif provider == Provider.GOOGLE_GEMINI:
            return cls._parsers[Provider.GOOGLE_GEMINI]
        elif provider == Provider.GOOGLE_VERTEX:
            return cls._parsers[Provider.GOOGLE_VERTEX]
        else:
            raise ValueError(f"Unsupported provider: {provider}")

    @classmethod
    @overload
    def parse(
        cls, response: AnthropicMessage, provider: Literal[Provider.ANTHROPIC]
    ) -> ParsedResponse: ...

    @classmethod
    @overload
    def parse(
        cls, response: OpenAIChatCompletion, provider: Literal[Provider.OPENAI_CHAT]
    ) -> ParsedResponse: ...

    @classmethod
    @overload
    def parse(
        cls, response: OpenAIResponse, provider: Literal[Provider.OPENAI_RESPONSES]
    ) -> ParsedResponse: ...

    @classmethod
    @overload
    def parse(
        cls,
        response: GoogleGenerateContentResponse,
        provider: Literal[Provider.GOOGLE_GEMINI],
    ) -> ParsedResponse: ...

    @classmethod
    @overload
    def parse(
        cls,
        response: GoogleGenerateContentResponse,
        provider: Literal[Provider.GOOGLE_VERTEX],
    ) -> ParsedResponse: ...

    @classmethod
    def parse(
        cls,
        response: AnthropicMessage
        | OpenAIChatCompletion
        | OpenAIResponse
        | GoogleGenerateContentResponse,
        provider: Provider,
    ) -> ParsedResponse:
        """
        Parse an LLM response based on provider type.

        Args:
            response: The raw response from the provider
            provider: The provider type

        Returns:
            A normalized ParsedResponse object
        """
        if provider not in cls._parsers:
            raise ValueError(f"Unsupported provider: {provider}")

        parser = cls._parsers[provider]
        return parser.parse_response(response, provider)


class ResponseMetadataExtractor:
    """Utility class to extract and format response metadata for telemetry."""

    @staticmethod
    def extract_for_telemetry(parsed_response: ParsedResponse) -> dict[str, Any]:
        """
        Extract metadata from a parsed response for telemetry purposes.

        Args:
            parsed_response: The normalized parsed response

        Returns:
            A dictionary of metadata suitable for telemetry
        """
        metadata = {
            "provider": parsed_response.provider.value,
            "model": parsed_response.model,
            "stop_reason": parsed_response.stop_reason.value,
            "usage": parsed_response.usage.model_dump(),
        }

        return metadata


class ResponseTypeDetector:
    """Utility class for detecting and classifying LLM response types.

    Analyzes parsed responses and request configurations to determine
    the type of response for observability and UI rendering purposes.
    """

    @staticmethod
    def detect(
        parsed_response: ParsedResponse, request_config: RequestConfig | None = None
    ) -> ResponseType:
        """Detect the response type based on content and configuration.

        Args:
            parsed_response: The parsed LLM response
            request_config: Optional request configuration used for the call

        Returns:
            ResponseType enum value indicating the type of response
        """
        # Check if structured generation was requested
        is_structured = ResponseTypeDetector._is_structured_generation(request_config)

        # Analyze content blocks to determine presence of text and tool calls
        has_text = False
        has_tool_calls = False

        for candidate in parsed_response.candidates:
            for block in candidate.content_blocks:
                if isinstance(block, TextContent):
                    has_text = True
                elif isinstance(block, ToolCall):
                    has_tool_calls = True

        # Classify based on content and configuration
        if is_structured:
            if has_tool_calls:
                return ResponseType.STRUCTURED_WITH_TOOLS
            else:
                return ResponseType.STRUCTURED
        else:
            if has_text and has_tool_calls:
                return ResponseType.TEXT_WITH_TOOLS
            elif has_tool_calls:
                return ResponseType.TOOL_CALLS
            else:
                return ResponseType.TEXT

    @staticmethod
    def _is_structured_generation(request_config: RequestConfig | None) -> bool:
        """Check if the request was configured for structured generation.

        Args:
            request_config: The request configuration to check

        Returns:
            True if structured generation was configured, False otherwise
        """
        if request_config is None:
            return False

        # Check provider-specific structured generation indicators
        if hasattr(request_config, "response_format"):
            # OpenAI Chat API structured outputs
            response_format = request_config.response_format
            if isinstance(response_format, dict):
                format_type = response_format.get("type")
                if format_type in ("json_object", "json_schema"):
                    return True

        if hasattr(request_config, "response_schema"):
            # Google structured generation
            if request_config.response_schema is not None:
                return True

        if hasattr(request_config, "response_mime_type"):
            # Google JSON mode
            if request_config.response_mime_type == "application/json":
                return True

        return False

    @staticmethod
    def count_tool_calls(parsed_response: ParsedResponse) -> int:
        """Count the total number of tool calls across all candidates.

        Args:
            parsed_response: The parsed LLM response

        Returns:
            Total count of tool calls in the response
        """
        count = 0
        for candidate in parsed_response.candidates:
            for block in candidate.content_blocks:
                if isinstance(block, ToolCall):
                    count += 1
        return count


class LLMEvent(LLMEventModelBase[ParsedResponse, Message]):
    """Domain model for LLM interactions"""

    messages: list[Message] = Field(..., alias="messages")
    parsed_response: ParsedResponse = Field(..., alias="parsedResponse")
