from datetime import date, datetime
from enum import Enum
from typing import Annotated, Any, List

from openai.types.chat import (
    ChatCompletionMessage,
    ChatCompletionMessageParam,
    ChatCompletionToolMessageParam,
    ChatCompletionUserMessageParam,
)
from pydantic import AfterValidator, BaseModel, BeforeValidator, model_validator

from elluminate.schemas.base import BatchCreateStatus
from elluminate.schemas.generation_metadata import GenerationMetadata
from elluminate.schemas.prompt import Prompt
from elluminate.schemas.rating import Rating, RatingValue


def fix_message_tool_calls(messages: List[Any]) -> List[Any]:
    """Fix tool_calls field in assistant messages.

    The backend may return assistant messages with tool_calls=None and
    we expect tool_calls to be an empty list when not present.
    """
    for msg in messages:
        if isinstance(msg, dict) and msg.get("role") == "assistant":
            if "tool_calls" in msg and msg["tool_calls"] is None:
                msg["tool_calls"] = []

    return messages


def ensure_tool_calls_list(messages: List[ChatCompletionMessageParam]) -> List[ChatCompletionMessageParam]:
    """Ensure tool_calls is a proper list for assistant messages with tool_calls.

    The `ChatCompletionMessageParam` parses the tool calls as an Iterator. Convert this to a list for ease of use.
    """
    for msg in messages:
        if isinstance(msg, dict) and msg.get("role") == "assistant" and "tool_calls" in msg:
            msg["tool_calls"] = list(msg["tool_calls"])

    return messages


class PromptResponseFilter(BaseModel):
    """Filter for prompt responses."""

    experiment_id: int | None = None
    prompt_id: int | None = None
    prompt_template_id: int | None = None
    template_variables_id: int | None = None
    collection_id: int | None = None
    llm_model_name: str | None = None
    epoch: int | None = None
    response_ids: list[int] | None = None
    duration_seconds_min: float | None = None
    duration_seconds_max: float | None = None
    input_tokens_min: int | None = None
    input_tokens_max: int | None = None
    output_tokens_min: int | None = None
    output_tokens_max: int | None = None


class PromptResponse(BaseModel):
    """Prompt response model."""

    id: int
    prompt: Prompt
    messages: Annotated[
        List[ChatCompletionMessage | ChatCompletionToolMessageParam | ChatCompletionUserMessageParam],
        BeforeValidator(fix_message_tool_calls),
        AfterValidator(ensure_tool_calls_list),
    ]
    generation_metadata: GenerationMetadata | None
    error: str | None
    epoch: int
    ratings: list[Rating] = []
    trajectory: dict[str, Any] | None = None
    annotation: str = ""
    created_at: datetime

    @property
    def response_str(self) -> str:
        """Extract the text content from the assistant's response.

        Returns the content of the last assistant message in the conversation.
        Returns an empty string if no assistant message is found.
        """
        for msg in reversed(self.messages):
            if isinstance(msg, dict) and msg.get("role") == "assistant":
                return msg.get("content") or ""
            elif hasattr(msg, "role") and msg.role == "assistant":
                return msg.content or ""
        return ""


class CreatePromptResponseRequest(BaseModel):
    """Request to create a new prompt response."""

    template_variables_id: int
    experiment_id: int
    epoch: int = 1
    llm_config_id: int | None = None
    messages: Annotated[
        List[ChatCompletionMessage | ChatCompletionToolMessageParam | ChatCompletionUserMessageParam],
        BeforeValidator(fix_message_tool_calls),
        AfterValidator(ensure_tool_calls_list),
    ] = []
    metadata: GenerationMetadata | None = None

    @model_validator(mode="after")
    def validate_messages_have_response_content(self) -> "CreatePromptResponseRequest":
        if self.messages and not any(
            (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) in ("assistant", "tool")
            for m in self.messages
        ):
            raise ValueError("messages must contain at least one assistant or tool message to be ratable")
        return self


class BatchCreatePromptResponseRequest(BaseModel):
    prompt_response_ins: list[CreatePromptResponseRequest]


class BatchCreatePromptResponseStatus(BatchCreateStatus[PromptResponse]):
    # The result is a tuple with epoch and response
    pass


class ResponsesSampleFilter(BaseModel):
    """Filter for responses samples."""

    experiment_id: int
    criteria_filter_by: dict[int, RatingValue] | None = None
    duration_seconds_min: float | None = None
    duration_seconds_max: float | None = None
    input_tokens_min: int | None = None
    input_tokens_max: int | None = None
    output_tokens_min: int | None = None
    output_tokens_max: int | None = None
    show_only_annotated_responses: bool | None = None


# TODO Switch to StrEnum when we drop support for Python 3.10
class ResponsesSampleSortBy(str, Enum):
    """Sort by for responses samples."""

    RATING_ASC = "rating-asc"
    RATING_DESC = "rating-desc"
    INPUT_TOKENS_ASC = "input-tokens-asc"
    INPUT_TOKENS_DESC = "input-tokens-desc"
    OUTPUT_TOKENS_ASC = "output-tokens-asc"
    OUTPUT_TOKENS_DESC = "output-tokens-desc"
    DURATION_ASC = "duration-asc"
    DURATION_DESC = "duration-desc"


class ResponsesSample(BaseModel):
    """Responses sample model."""

    experiment_id: int
    template_variables_id: int
    prompt_response_ids: list[int]


class DailyUsageStats(BaseModel):
    """Daily usage stats model."""

    date: date
    requests: int = 0
    input_tokens: int = 0
    output_tokens: int = 0


class ResponsesStats(BaseModel):
    """Responses stats model."""

    total_responses: int
    average_duration_seconds: float | None = None
    min_duration_seconds: float | None = None
    max_duration_seconds: float | None = None
    total_input_tokens: int
    total_output_tokens: int
    average_input_tokens: float
    average_output_tokens: float
    average_rating: float | None = None
    daily_usage: list[DailyUsageStats]
