import datetime
from typing import Any, Literal, Optional

import nanoid  # type: ignore
from openai.types.chat import ChatCompletion
from openai.types.chat.chat_completion_reasoning_effort import (
    ChatCompletionReasoningEffort,
)
from pydantic import BaseModel, Field, computed_field, model_validator

from retab.types.chat import ChatCompletionRetabMessage
from retab.types.documents.extractions import RetabParsedChatCompletion

from ..utils.usage.usage import (
    CostBreakdown,
    compute_cost_from_model,
    compute_cost_from_model_with_breakdown,
)
from .ai_models import Amount
from .modalities import Modality

ValidationsState = Literal["pending", "validated", "invalid"]


class ExtractionSource(BaseModel):
    type: Literal[
        "api",
        "annotation",
        "processor",
        "automation",
        "automation.link",
        "automation.mailbox",
        "automation.cron",
        "automation.outlook",
        "automation.endpoint",
        "schema.extract",
    ] = Field(description="Type of extraction")
    id: str | None = Field(default=None, description="ID the trigger of the extraction")


ExtractionSteps = str | Literal["initialization", "prepare_messages", "yield_first_token", "completion"]  # Steps are meant to not overlap
BrowserCanvas = Literal["A3", "A4", "A5"]


class ExtractionTimingStep(BaseModel):
    name: ExtractionSteps
    duration: float  # in seconds
    notes: str | None = None


class Extraction(BaseModel):
    id: str = Field(
        default_factory=lambda: "extr_" + nanoid.generate(),
        description="Unique identifier of the analysis",
    )
    messages: list[ChatCompletionRetabMessage] = Field(default_factory=list)
    messages_gcs: str = Field(..., description="GCS path to the messages")
    file_gcs_paths: list[str] = Field(..., description="GCS paths to the files")
    file_ids: list[str] = Field(..., description="IDs of the files")
    # Legacy fields for backward compatibility
    file_gcs: str = Field(default="", description="GCS path to the first file (deprecated)")
    file_id: str = Field(default="", description="ID of the first file (deprecated)")

    status: Literal["success", "failed"] = Field(..., description="Whether the analysis was successful")
    completion: RetabParsedChatCompletion | ChatCompletion = Field(..., description="Response generated by the analysis")
    json_schema: Any = Field(..., description="Response format (JSON Schema or pydantic_v2.BaseModel)")
    model: str = Field(..., description="Model used for the analysis")
    temperature: float = Field(default=0.0, description="Temperature used for the analysis")
    source: ExtractionSource = Field(..., description="Source of the extraction")
    image_resolution_dpi: int = Field(default=96, description="Resolution of the image sent to the LLM")
    browser_canvas: BrowserCanvas = Field(
        default="A4",
        description="Sets the size of the browser canvas for rendering documents in browser-based processing. Choose a size that matches the document type.",
    )
    modality: Modality = Field(default="native", description="Modality of the extraction")
    reasoning_effort: Optional[ChatCompletionReasoningEffort] = Field(
        default=None,
        description="The effort level for the model to reason about the input data.",
    )
    n_consensus: int = Field(default=1, description="Number of consensus models used for the extraction")
    timings: list[ExtractionTimingStep] = Field(default_factory=list, description="Timings of the extraction")

    # Infered from the schema
    schema_id: str = Field(..., description="Version of the schema used for the analysis")
    schema_data_id: str = Field(..., description="Version of the schema data used for the analysis")
    created_at: datetime.datetime = Field(
        default_factory=lambda: datetime.datetime.now(datetime.timezone.utc),
        description="Timestamp of the creation of the extraction object",
    )
    request_at: datetime.datetime | None = Field(default=None, description="Timestamp of the extraction request if provided.")
    organization_id: str = Field(..., description="Organization ID of the user or application")
    validation_state: Optional[ValidationsState] = Field(default=None, description="Validation state of the extraction")
    billed: bool = Field(default=False, description="Whether the extraction has been billed or not")

    @model_validator(mode="before")
    def handle_legacy_fields(cls, data):
        """Handle backward compatibility for legacy file_gcs and file_id fields."""
        if isinstance(data, dict):
            # If only legacy fields are provided, convert to new format
            if "file_gcs" in data and "file_gcs_paths" not in data:
                data["file_gcs_paths"] = [data["file_gcs"]]
            if "file_id" in data and "file_ids" not in data:
                data["file_ids"] = [data["file_id"]]

            # Set legacy fields from new format for backward compatibility
            if "file_gcs_paths" in data and data["file_gcs_paths"]:
                data["file_gcs"] = data["file_gcs_paths"][0]
            if "file_ids" in data and data["file_ids"]:
                data["file_id"] = data["file_ids"][0]
        return data

    @computed_field  # type: ignore
    @property
    def api_cost(self) -> Optional[Amount]:
        if self.completion and self.completion.usage:
            try:
                cost = compute_cost_from_model(self.completion.model, self.completion.usage)
                return cost
            except Exception as e:
                print(f"Error computing cost: {e}")
                return None
        return None

    @computed_field  # type: ignore
    @property
    def cost_breakdown(self) -> Optional[CostBreakdown]:
        if self.completion and self.completion.usage:
            try:
                cost = compute_cost_from_model_with_breakdown(self.completion.model, self.completion.usage)
                return cost
            except Exception as e:
                print(f"Error computing cost: {e}")
                return None
        return None
