import math
from datetime import datetime
from typing import List, Optional, Union, Any, Dict
from uuid import UUID
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from .common_types import ModelClass, ModelCore, criteria_starts

# Maximum length for evaluation criteria across all endpoints
MAX_EVALUATION_CRITERIA_LENGTH = 4096


def is_binary_evaluation(criteria: str) -> bool:
    """Check if evaluation criteria is binary (Passes if/Fails if) vs reward-based."""
    return any(criteria.startswith(cs) for cs in criteria_starts["binary"])


def validate_criteria_length(criteria: str | None) -> str | None:
    """Validate that criteria length does not exceed the maximum allowed length."""
    if criteria is not None and len(criteria) > MAX_EVALUATION_CRITERIA_LENGTH:
        raise ValueError(
            f"Criteria length cannot exceed {MAX_EVALUATION_CRITERIA_LENGTH} characters"
        )
    return criteria


class LLMToolResponse(BaseModel):
    id: str = Field(..., description="Unique identifier for the tool call.")
    response: str = Field(..., description="Tool response")

    name: Optional[str] = Field(default=None, description="Tool name")

    model_config = ConfigDict(
        extra="allow",
    )


class LLMToolCall(BaseModel):
    id: str = Field(..., description="Unique identifier for the tool call.")
    name: str = Field(..., description="Tool name")
    parameters: str = Field(..., description="Tool parameters.")

    model_config = ConfigDict(
        extra="allow",
    )


class LLMToolDefinition(BaseModel):
    name: str = Field(..., description="Tool name")
    description: str = Field(..., description="Tool description")
    parameters: str = Field(..., description="Tool parameters.")

    model_config = ConfigDict(
        extra="allow",
    )


class LLMInput(BaseModel):
    role: str = Field(
        ...,
        description="Role of the message sender (e.g., user, system, assistant, tool)",
    )
    type: Optional[str] = Field(
        default=None,
        description="Type of message (e.g., text, thinking_tokens, tool_call). This can be arbitrary but should be consistent within a trace.",
    )
    content: Union[str, LLMToolCall, LLMToolResponse] = Field(
        ..., description="Content of the LLM Input."
    )


class LLMOutput(BaseModel):
    type: Optional[str] = Field(
        default=None,
        description="Type of message (e.g., text, thinking_tokens, tool_call). This can be arbitrary but should be consistent within a trace.",
    )
    content: Union[str, LLMToolCall] = Field(..., description="Content of the message")


class LLMInteraction(BaseModel):
    input_messages: List[LLMInput] = Field(
        ...,
        description="List of input messages to the LLM invocation.",
    )
    output_messages: List[LLMOutput] = Field(
        ...,
        description="List of output messages from the LLM invocation.",
    )
    tools_available: Optional[List[LLMToolDefinition]] = Field(
        default=None,
        description="List of tools available to the agent in this interaction.",
    )


class AgentInstance(BaseModel):
    id: str = Field(..., description="Unique identifier for the agent instance.")
    name: str = Field(..., description="User defined name of the agent invoked.")
    interactions: List[Union[str, LLMInteraction]] = Field(
        ...,
        description="List of ordered LLM interactions or agent instance ids for this agent instance.",
    )

    @classmethod
    def from_chat_session(cls, chat_session: Dict[str, Any]) -> "AgentInstance":
        """
        Convert a chat session dictionary to an AgentInstance.

        Args:
            chat_session: A dictionary containing chat session data with 'messages'
                         and optionally 'tools' keys. Can also accept a ChatSession
                         object (for backward compatibility).

        Returns:
            An AgentInstance with interactions representing each LLM invocation
        """
        # Handle both dict and ChatSession object (for backward compatibility)
        if hasattr(chat_session, "raw_data"):
            raw_data = chat_session.raw_data
        else:
            raw_data = chat_session

        messages = raw_data.get("messages", [])
        tools = raw_data.get("tools", [])

        # Convert tools to LLMToolDefinition
        tools_available = None
        if tools:
            tools_available = [
                LLMToolDefinition(
                    name=tool.get("function", {}).get("name", ""),
                    description=tool.get("function", {}).get("description", ""),
                    parameters=str(tool.get("function", {}).get("parameters", {})),
                )
                for tool in tools
            ]

        # Build interactions: each assistant message represents one LLM invocation
        # Each interaction's input includes the full conversation history up to that point
        interactions = []
        all_messages = (
            []
        )  # Track all messages (inputs and outputs) in chronological order
        tool_call_map = {}  # Map tool_call_id to tool_name for resolving tool responses

        for msg in messages:
            role = msg.get("role")
            content = msg.get("content")

            if role in ["user", "system", "tool"]:
                # These are inputs to the next LLM invocation
                if role == "tool":
                    tool_call_id = msg.get("tool_call_id")
                    # Look up the tool name from the tool_call_map
                    tool_name = tool_call_map.get(tool_call_id)
                    input_msg = LLMInput(
                        role=role,
                        type="tool_response",
                        content=LLMToolResponse(
                            id=tool_call_id,
                            name=tool_name,
                            response=(
                                content if isinstance(content, str) else str(content)
                            ),
                        ),
                    )
                else:
                    input_msg = LLMInput(
                        role=role,
                        content=content if isinstance(content, str) else str(content),
                    )
                all_messages.append(input_msg)

            elif role == "assistant":
                # This is an LLM output - create an interaction
                output_messages = []

                tool_calls = msg.get("tool_calls")
                if tool_calls:
                    for tool_call in tool_calls:
                        tc_id = tool_call.get("id")
                        tc_func = tool_call.get("function", {})
                        tc_name = tc_func.get("name")
                        tc_args = tc_func.get("arguments", "")

                        # Store tool call ID -> name mapping for later tool responses
                        tool_call_map[tc_id] = tc_name

                        output_msg = LLMOutput(
                            type="tool_call",
                            content=LLMToolCall(
                                id=tc_id,
                                name=tc_name,
                                parameters=tc_args,
                            ),
                        )
                        output_messages.append(output_msg)

                if content:
                    output_msg = LLMOutput(
                        type="text",
                        content=content if isinstance(content, str) else str(content),
                    )
                    output_messages.append(output_msg)

                # Create the interaction for this LLM invocation
                # Input includes all messages up to this point (full conversation history)
                interaction = LLMInteraction(
                    input_messages=all_messages.copy(),
                    output_messages=output_messages,
                    tools_available=tools_available,
                )
                interactions.append(interaction)

                # Add the outputs to the running history
                # Convert LLMOutput to LLMInput for the next interaction's input
                for output_msg in output_messages:
                    if output_msg.type == "tool_call":
                        # Tool calls should be added to all_messages for next interaction
                        all_messages.append(
                            LLMInput(
                                role="assistant",
                                type="tool_call",
                                content=output_msg.content,
                            )
                        )
                    elif output_msg.type == "text":
                        # Assistant text responses become part of conversation history
                        # Convert to LLMInput for next interaction
                        all_messages.append(
                            LLMInput(
                                role="assistant",
                                type="text",
                                content=output_msg.content,
                            )
                        )

        return cls(
            id="agent_instance_0",
            name="default_agent",
            interactions=interactions,
        )


class MultiAgentTrace(BaseModel):
    root_agent_instance_id: str = Field(
        ...,
        description="User invoked Agent Instance ID. This is the root of the trace.",
    )
    agent_instance_by_id: dict[str, AgentInstance] = Field(
        ..., description="Mapping of ID to Agent Instance objects."
    )


class RequestBase(BaseModel):
    evaluation_criteria: Optional[str] = Field(
        None,
        min_length=1,
        description="Criteria used for evaluation. Begins with one of the following: "
        + ", ".join(criteria_starts["reward"] + criteria_starts["binary"])
        + "Note, align-lightning-20251127 only supports Passes/Fails criteria. "
        + "Mutually exclusive with evaluator_name.",
    )
    evaluator_name: Optional[str] = Field(
        None,
        min_length=1,
        description="Name of an existing evaluator (group_id) to use for evaluation. "
        "Mutually exclusive with evaluation_criteria. ",
        max_length=255,
    )
    model_core: ModelCore = Field(
        default=ModelCore.ALIGN_20251111,
        description="The model core for reward evaluation. Defaults to align-20251111 if not specified.",
    )

    tags: Optional[dict[str, str]] = Field(
        default={},
        description="Optional key-value pairs to tag and categorize the request. All tag values must be strings.",
    )

    trace_id: Optional[UUID] = Field(
        None,
        description="Optional trace ID to associate this evaluation with an existing trace.",
    )

    @field_validator("tags")
    @classmethod
    def validate_tags(cls, value):
        if value is None:
            return {}
        if not isinstance(value, dict):
            raise ValueError("Tags must be a dictionary of key-value pairs.")

        for k, v in value.items():
            if not isinstance(k, str):
                raise ValueError("Tag keys must be strings.")
            if len(k) > 64:
                raise ValueError("Tag keys must not exceed 64 characters.")

            # Enforce that tag values must be strings
            if not isinstance(v, str):
                raise ValueError("Tag values must be strings.")
            if len(v) > 64:
                raise ValueError("Tag values must not exceed 64 characters.")

        return value

    _validate_evaluation_criteria_length = field_validator("evaluation_criteria")(
        validate_criteria_length
    )

    @field_validator("evaluation_criteria")
    @classmethod
    def evaluation_criteria_must_start_with_correct_prefix(cls, value):
        # Only validate if criteria is provided
        if value is None:
            return None

        # Guard against empty or whitespace-only strings
        if not value or not value.strip():
            raise ValueError("Evaluation criteria cannot be empty or whitespace-only")

        value = value[0].upper() + value[1:]
        if not any(
            value.startswith(start)
            for start in criteria_starts["reward"] + criteria_starts["binary"]
        ):
            raise ValueError(
                "Evaluation criteria must start with one of the following: "
                + ", ".join(criteria_starts["reward"] + criteria_starts["binary"])
            )
        return value

    @model_validator(mode="after")
    def validate_exactly_one_evaluation_input(self):
        """Ensure exactly one of evaluation_criteria or evaluator_name is provided."""
        has_criteria = self.evaluation_criteria is not None
        has_name = self.evaluator_name is not None

        if has_criteria and has_name:
            raise ValueError(
                "Provide either 'evaluation_criteria' or 'evaluator_name', not both"
            )

        if not has_criteria and not has_name:
            raise ValueError(
                "Either 'evaluation_criteria' or 'evaluator_name' is required"
            )

        return self


class RewardRequest(RequestBase):
    """
    Request model for reward score evaluation of LLM responses against specified criteria.
    """

    messages: List[dict] = Field(
        ...,
        description="A list of chat messages in [OpenAI](https://platform.openai.com/docs/api-reference/chat/create#chat-create-messages) or [Anthropic](https://platform.claude.com/docs/en/api/messages#message_param) messages format",
        min_length=2,
    )
    system: Optional[str] = Field(
        None,
        description="System message is separate for Anthropic-style LLM calls. Optional.",
    )
    tools: Optional[List[dict]] = Field(
        None,
        description="List of tools available for the assistant to call in [OpenAI](https://platform.openai.com/docs/api-reference/chat/create#chat_create-tools) or [Anthropic](https://platform.claude.com/docs/en/api/messages#tool_union) format. Optional.",
    )
    evaluate_latest: bool = Field(
        default=True,
        description="Whether to evaluate only the latest response or all responses. Only True is supported for `align-lightning` model cores.",
    )

    @model_validator(mode="after")
    def lightning_not_allowed_for_all_responses_evaluation(self):
        if (
            not self.evaluate_latest
            and self.model_core.model_class() == ModelClass.ALIGN_LIGHTNING
        ):
            raise ValueError(
                "Align Lightning model core not supported when evaluating all responses."
            )
        return self


class TraceRewardRequest(RequestBase):
    """
    Request model for a trace based evaluation of LLM responses against specified criteria.
    """

    trace: MultiAgentTrace = Field(
        ...,
        description="A Multi Agent Trace object representing the full trace of agent interactions.",
    )
    evaluate_latest: bool = Field(
        default=False,
        description="Whether to evaluate only the latest response or all responses.",
    )

    @field_validator("model_core")
    @classmethod
    def lightning_not_allowed_for_trace_evaluation(cls, value):
        if value.model_class() == ModelClass.ALIGN_LIGHTNING:
            raise ValueError(
                "Align Lightning model core not supported for trace evaluations."
            )
        return value

    @field_validator("evaluate_latest")
    @classmethod
    def evaluate_latest_must_be_false_for_trace_evaluation(cls, value):
        if value:
            raise ValueError("evaluate_latest must be false for trace evaluations.")
        return value

    @field_validator("trace")
    @classmethod
    def trace_must_have_llm_interactions(cls, body):
        ### Check if trace has any LLM interactions
        has_llm_interaction = False
        for agent_instance in body.agent_instance_by_id.values():
            for interaction in agent_instance.interactions:
                if isinstance(interaction, LLMInteraction):
                    has_llm_interaction = True
                    break
            if has_llm_interaction:
                break

        if not has_llm_interaction:
            raise ValueError(
                "Trace contains no LLM interactions. At least one agent instance must have an LLM interaction."
            )
        return body


class RewardGPURequest(RewardRequest):
    """
    Request model for reward score evaluation of LLM responses against specified criteria,
    specifically for GPU-based evaluation.
    """

    explanation: str = Field(
        description="Explanation of the evaluation criteria. Optional.",
    )


class Source(BaseModel):
    id: str = Field(description="Identifier of the source.")
    type: str = Field(description="Type of the source.")


class ScoreResponse(BaseModel):
    """
    Response model for evaluation scores.
    """

    score: Optional[float] = Field(
        None,
        description="Evaluation score between 0 and 1. If null, the criteria was deemed not applicable.",
    )
    explanation: str = Field(description="Explanation of the evaluation score")
    sources: Optional[List[Source]] = Field(
        default=None,
        description="Sources referenced in the evaluation, e.g. memories, annotations, or documents.",
    )

    @field_validator("score")
    def validate_score(cls, value):
        if value is None:
            return value
        if math.isnan(value):
            return None
        if not 0 <= value <= 1:
            raise ValueError("Score must be between 0 and 1.")
        return value

    def get_rounded(self) -> "ScoreResponse":
        """
        Round the score to either 0 or 1 based on a threshold of 0.5.
        """
        if self.score is not None:
            self.score = 1.0 if self.score >= 0.5 else 0.0
        return self


class SummaryStatistics(BaseModel):
    median: Optional[float] = Field(default=None, description="Median score.")
    min: Optional[float] = Field(default=None, description="Minimum score.")
    max: Optional[float] = Field(default=None, description="Maximum score.")
    std: Optional[float] = Field(
        default=None, description="Standard deviation of scores."
    )


class AgentTraceResult(BaseModel):
    agent_name: str = Field(..., description="Name of the agent evaluated.")
    results_by_agent_instance_id: dict[str, Union[ScoreResponse, None]] = Field(
        ...,
        description="Mapping of Agent Instance ID to their respective Score Response, or None depending on criteria.",
    )
    summary_statistics: Optional[SummaryStatistics] = Field(
        default=None,
        description="Summary statistics for the agent's evaluations. Only applicable for reward evaluations.",
    )


class MultiAgentTraceResponse(BaseModel):
    """
    Response model for multi-agent trace evaluations.
    """

    request_id: str = Field(
        ..., description="Unique identifier for the evaluation request."
    )
    results_by_agent_name: dict[str, AgentTraceResult] = Field(
        ..., description="Mapping of Agent Name to their respective trace results."
    )

    def __str__(self) -> str:
        return self.model_dump_json(indent=2)

    def get_rounded(self) -> "MultiAgentTraceResponse":
        """
        Round all scores in the response to either 0 or 1 based on a threshold of 0.5.
        """
        for agent_trace_result in self.results_by_agent_name.values():
            for (
                agent_instance_id,
                score_response,
            ) in agent_trace_result.results_by_agent_instance_id.items():
                if score_response is not None and score_response.score is not None:
                    score_response.score = 1.0 if score_response.score >= 0.5 else 0.0

        return self


class Evaluator(BaseModel):
    """Evaluator model for creating and retrieving evaluators."""

    id: str = Field(..., description="Unique identifier for the evaluator")
    group_id: str = Field(..., description="Group ID of the evaluator")
    criteria: str = Field(..., description="Evaluation criteria")
    created_at: str = Field(..., description="Timestamp when the evaluator was created")
    deleted: bool = Field(..., description="Whether this version has been superseded")


class CreateEvaluatorRequest(BaseModel):
    """Request model for creating a new evaluator."""

    group_id: Optional[str] = Field(
        None,
        description="Group ID of the evaluator. If not provided, one will be auto-generated.",
    )
    criteria: str = Field(..., description="Evaluation criteria")

    _validate_criteria = field_validator("criteria")(validate_criteria_length)


class UpdateEvaluatorRequest(BaseModel):
    """Request model for updating an evaluator's criteria."""

    criteria: str = Field(..., min_length=1, description="New evaluation criteria")

    _validate_criteria = field_validator("criteria")(validate_criteria_length)


class EvaluationContext(BaseModel):
    """
    Business context for evaluation requests.

    Handlers create this object and attach it to request.state.evaluation_context.
    Middleware reads it when creating usage logs.
    """

    model_config = ConfigDict(frozen=True)  # Immutable after creation for safety

    model_core: str = Field(
        ..., description="Model core identifier (e.g., 'align-20251111')"
    )
    evaluator_id: Optional[UUID] = Field(None, description="ID of the evaluator")
    is_binary: bool = Field(
        False, description="Whether this is a binary (Passes/Fails) evaluation"
    )

    @property
    def model_class(self) -> str:
        """Extract model class from model_core (e.g., 'align-20251111' -> 'align')."""
        if self.model_core and "-" in self.model_core:
            return self.model_core.rsplit("-", 1)[0]
        return "unknown"


class AttachCommentRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")
    trace_id: UUID = Field(..., description="The trace (request) ID to annotate")
    comment: str = Field(..., description="Comment to attach to this trace")


class AttachRatingRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")
    trace_id: UUID = Field(..., description="The trace (request) ID to annotate")
    rating: bool = Field(..., description="Rating: true = positive, false = negative")


class AttachCommentResponse(BaseModel):
    id: UUID
    trace_id: UUID
    comment: str
    created_at: datetime


class AttachRatingResponse(BaseModel):
    trace_id: UUID
    rating: bool


class TraceEvaluation(BaseModel):
    id: UUID
    evaluator_name: Optional[str]
    score: Optional[float]
    explanation: Optional[str]
    sub_trace: Optional[Dict[str, Any]]
    sources: Optional[List[Source]] = None


class TraceResponse(BaseModel):
    id: UUID
    created_at: datetime
    tags: Optional[Dict[str, str]]
    rating: Optional[bool]
    trace: Optional[Dict[str, Any]]
    comments: List[AttachCommentResponse]
    evaluations: List[TraceEvaluation]


class ListTracesResponse(BaseModel):
    traces: List[TraceResponse]
    next_page_token: Optional[str]
