import typing as t

from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
from pydantic.json_schema import SkipJsonSchema

from common_py.model import enums
from common_py.workflow.nodes.node import ErrorHandling
from common_py.workflow.nodes.workflow_types import InferenceMode
from common_py.workflow.utils.pydantic_model_serializer import (
    deserialize_model,
    serialize_model,
)

from . import ToolOutput


class LLMUsageData(BaseModel):
    prompt_tokens: int
    prompt_tokens_cost_usd: float | None
    completion_tokens: int
    completion_tokens_cost_usd: float | None
    total_tokens: int
    total_tokens_cost_usd: float | None


class ChatInput(BaseModel):
    MODEL_PREFIX: t.ClassVar[str | None] = None

    schema: dict[str, t.Any] | type[BaseModel] | None = None  # type: ignore
    model: str = enums.ModelName.GPT_4O_MINI
    prompt_str: str | None = None
    temperature: float = 0.1
    stop: list[str] | None = None
    error_handling: ErrorHandling = ErrorHandling(timeout_seconds=600)
    top_p: float | None = None
    seed: int | None = None
    presence_penalty: float | None = None
    frequency_penalty: float | None = None
    tool_choice: str | None = None
    inference_mode: InferenceMode = InferenceMode.DEFAULT
    output_schema: dict[str, t.Any] | type[BaseModel] | None = Field(
        default=None, description="Schema for structured output or function calling."
    )
    stream: bool = Field(default=False, description="Enable streaming response")

    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")

    @field_serializer("output_schema")
    def serialize_schema(self, value: t.Any | None, _info) -> dict[str, t.Any] | None:  # type: ignore
        if value is None:
            return None  # type: ignore
        elif isinstance(value, type) and issubclass(value, BaseModel):
            return serialize_model(value)
        else:
            return value  # type: ignore

    @field_validator("model")
    @classmethod
    def set_model(cls, value: str | None) -> str:
        """Set the model with the appropriate prefix.

        Args:
            value (str | None): The model value.

        Returns:
            str: The model value with the prefix.
        """
        if cls.MODEL_PREFIX is not None and not value.startswith(cls.MODEL_PREFIX):  # type: ignore
            value = f"{cls.MODEL_PREFIX}{value}"
        return value  # type: ignore

    @field_validator("output_schema")
    @classmethod
    def parse_inference_schema(
        cls, value: dict[str, t.Any] | type[BaseModel] | None
    ) -> dict[str, t.Any] | type[BaseModel] | None:
        if value is None:
            return None
        elif isinstance(value, dict):
            try:
                return deserialize_model(value)
            except Exception:
                return value
        else:
            return value


class ChatOutput(ToolOutput):
    llm_call: SkipJsonSchema[dict[str, t.Any]] = Field(default_factory=dict)
    content: str | None = Field(default="")
    schema_result_json: dict[str, t.Any] | None = Field(default=None)
    tool_calls: SkipJsonSchema[list[dict[str, t.Any]] | None] = Field(default=None)
