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

from uiform.types.chat import ChatCompletionUiformMessage
from uiform.types.documents.extractions import UiParsedChatCompletion

from .._utils.usage.usage import compute_cost_from_model, compute_cost_from_model_with_breakdown, CostBreakdown
from .ai_models import Amount
from .modalities import Modality

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


class ExtractionSource(BaseModel):
    type: Literal["api", "annotation","processor", "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


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[ChatCompletionUiformMessage] = Field(default_factory=list)
    messages_gcs: str = Field(..., description="GCS path to the messages")
    file_gcs: str = Field(..., description="GCS path to the file")
    file_id: str = Field(..., description="ID of the file")

    status: Literal["success", "failed"] = Field(..., description="Whether the analysis was successful")
    completion: UiParsedChatCompletion | 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: Literal['A3', 'A4', 'A5'] = 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.")
    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")

    @computed_field
    @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
