import warnings
from enum import Enum
from typing import Any, Optional

from pydantic import model_validator

from galtea.utils.case_insensitive_enum import CaseInsensitiveEnumMixin
from galtea.utils.from_camel_case_base_model import FromCamelCaseBaseModel
from galtea.utils.input_normalization import coerce_to_string, unwrap_input

__all__ = ["InferenceResult", "InferenceResultBase", "InferenceResultStatus", "InferenceResultUpdate"]


class InferenceResultStatus(CaseInsensitiveEnumMixin, str, Enum):
    PENDING = "PENDING"
    GENERATED = "GENERATED"
    FAILED = "FAILED"


class CostInfoProperties(FromCamelCaseBaseModel):
    cost_per_input_token: Optional[float] = None
    cost_per_output_token: Optional[float] = None
    cost_per_cache_read_input_token: Optional[float] = None
    cost: Optional[float] = None


class UsageInfoProperties(FromCamelCaseBaseModel):
    input_tokens: Optional[int] = None
    output_tokens: Optional[int] = None
    cache_read_input_tokens: Optional[int] = None
    tokens: Optional[int] = None


class InferenceResultBase(CostInfoProperties, UsageInfoProperties):
    """Model for creating inference results with mandatory fields (POST operation).

    The API stores input as a JSON object (e.g. ``{"user_message": "..."}``).
    The ``input`` field holds the raw string sent to the API.
    The ``input_data`` field holds the raw JSON object returned by the API.
    The deprecated ``actual_input`` property returns the unwrapped string for
    backward compatibility.
    """

    session_id: str
    status: Optional[InferenceResultStatus] = InferenceResultStatus.GENERATED
    input: Optional[str] = None
    actual_output: Optional[str] = None
    retrieval_context: Optional[str] = None
    latency: Optional[float] = None
    conversation_simulator_version: Optional[str] = None

    # Raw JSON object from the API (for structured input access)
    input_data: Optional[dict[str, Any]] = None

    @property
    def actual_input(self) -> Optional[str]:
        """Deprecated: use ``input`` directly. Returns the same string value.

        .. deprecated::
            ``actual_input`` is kept for backward compatibility and will be removed
            in a future version. The API no longer returns ``actualInput``; use
            ``input`` for the string value or ``input_data`` for the raw JSON object.
        """
        warnings.warn(
            "InferenceResult.actual_input is deprecated. Use InferenceResult.input directly.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.input

    @model_validator(mode="before")
    @classmethod
    def _normalize_input_field(cls, data: Any) -> Any:
        """Accept ``actualInput``/``actual_input`` and unwrap JSON input objects."""
        if not isinstance(data, dict):
            return data

        # Migrate legacy actualInput/actual_input keys to input
        for legacy_key in ("actualInput", "actual_input"):
            if legacy_key in data and "input" not in data:
                data["input"] = data.pop(legacy_key)
            elif legacy_key in data:
                data.pop(legacy_key)

        # Unwrap JSON input objects to strings
        raw_input = data.get("input")
        if isinstance(raw_input, dict):
            data["input_data"] = raw_input
            data["input"] = coerce_to_string(unwrap_input(raw_input))

        return data


class InferenceResult(InferenceResultBase):
    id: str
    index: int


class InferenceResultUpdate(FromCamelCaseBaseModel):
    """Model for updating inference results with all optional fields (PATCH operation).

    Fields are Optional (defaulting to None) so they can be excluded from serialization
    if not set (via exclude_unset=True in the service).
    """

    # Overridden from CostInfoProperties
    cost_per_input_token: Optional[float] = None
    cost_per_output_token: Optional[float] = None
    cost_per_cache_read_input_token: Optional[float] = None
    cost: Optional[float] = None

    # Overridden from UsageInfoProperties
    input_tokens: Optional[int] = None
    output_tokens: Optional[int] = None
    cache_read_input_tokens: Optional[int] = None
    tokens: Optional[int] = None

    # Fields specific to InferenceResultUpdate
    session_id: Optional[str] = None
    status: Optional[InferenceResultStatus] = None
    input: Optional[str] = None
    actual_output: Optional[str] = None
    retrieval_context: Optional[str] = None
    latency: Optional[float] = None
    conversation_simulator_version: Optional[str] = None
    index: Optional[int] = None
    created_at: Optional[str] = None
    deleted_at: Optional[str] = None

    @model_validator(mode="before")
    @classmethod
    def _normalize_input_field(cls, data: Any) -> Any:
        """Accept ``actualInput``/``actual_input`` from older callers."""
        if not isinstance(data, dict):
            return data

        for legacy_key in ("actualInput", "actual_input"):
            if legacy_key in data and "input" not in data:
                data["input"] = data.pop(legacy_key)
            elif legacy_key in data:
                data.pop(legacy_key)

        # Unwrap JSON input objects to strings
        raw_input = data.get("input")
        if isinstance(raw_input, dict):
            data["input"] = coerce_to_string(unwrap_input(raw_input))

        return data
