from typing import Annotated, Any, Literal
from uuid import UUID, uuid4

import pandas as pd
import polars as pl
from diffprivlib.validation import DiffprivlibMixin
from pydantic import (
    BaseModel,
    ConfigDict,
    Discriminator,
    Field,
    PlainSerializer,
    PlainValidator,
    ValidationInfo,
    field_validator,
)

from lomas_core.constants import DPLibraries
from lomas_core.models.exceptions import LomasServerExceptionType
from lomas_core.models.utils import (
    dataframe_from_dict,
    dataframe_to_dict,
    deserialize_model,
    polars_df_from_str,
    polars_df_to_str,
    serialize_model,
)


class ResponseModel(BaseModel):
    """Base model for any response from the server."""


class InitialBudgetResponse(ResponseModel):
    """Model for responses to initial budget queries."""

    initial_epsilon: float
    """The initial epsilon privacy loss budget."""
    initial_delta: float
    """The initial delta privacy loss budget."""


class SpentBudgetResponse(ResponseModel):
    """Model for responses to spent budget queries."""

    total_spent_epsilon: float
    """The total spent epsilon privacy loss budget."""
    total_spent_delta: float
    """The total spent delta privacy loss budget."""


class RemainingBudgetResponse(ResponseModel):
    """Model for responses to remaining budget queries."""

    remaining_epsilon: float
    """The remaining epsilon privacy loss budget."""
    remaining_delta: float
    """The remaining delta privacy loss budget."""


class DummyDsResponse(ResponseModel):
    """Model for responses to dummy dataset requests."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    dtypes: Any
    """The dummy_df column data types."""
    dummy_df: Annotated[pd.DataFrame, PlainSerializer(dataframe_to_dict), PlainValidator(dataframe_from_dict)]

    """The dummy dataframe."""

    @field_validator("dummy_df", mode="before")
    @classmethod
    def deserialize_dummy_df(cls, v: pd.DataFrame | dict, info: ValidationInfo) -> pd.DataFrame:
        """Decodes the dict representation of the dummy df with correct types.

        Only does so if the input value is not already a dataframe.
        Args:
            v (pd.DataFrame | dict): The dataframe to decode.
            info (ValidationInfo): Validation info to access other model fields.

        Returns:
            pd.DataFrame: The decoded dataframe.
        """
        if isinstance(v, pd.DataFrame):
            return v

        dtypes = info.data["dtypes"]
        dummy_df = dataframe_from_dict(v)
        dummy_df = dummy_df.astype(dtypes)
        return dummy_df


class CostResponse(ResponseModel):
    """Model for responses to cost estimation requests or queries."""

    model_config = ConfigDict(use_attribute_docstrings=True)

    epsilon: float
    """The epsilon cost of the query."""
    delta: float
    """The delta cost of the query."""


# Query Responses
# -----------------------------------------------------------------------------


# DiffPrivLib
class DiffPrivLibQueryResult(BaseModel):
    """Model for diffprivlib query result."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    type: Literal[DPLibraries.DIFFPRIVLIB] = DPLibraries.DIFFPRIVLIB
    """Result type description."""
    score: float
    """The trained model score."""
    model: Annotated[
        DiffprivlibMixin,
        PlainSerializer(serialize_model),
        PlainValidator(deserialize_model),
    ]
    """The trained model."""


# SmartnoiseSQL
class SmartnoiseSQLQueryResult(BaseModel):
    """Type for smartnoise_sql result type."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    type: Literal[DPLibraries.SMARTNOISE_SQL] = DPLibraries.SMARTNOISE_SQL
    """Result type description."""
    df: Annotated[
        pd.DataFrame,
        PlainSerializer(dataframe_to_dict),
        PlainValidator(dataframe_from_dict),
    ]
    """Dataframe containing the query result."""


# SmartnoiseSynth
class SmartnoiseSynthModel(BaseModel):
    """Type for smartnoise_synth result when it is a pickled model."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    """Result type description."""
    type: Literal[DPLibraries.SMARTNOISE_SYNTH] = DPLibraries.SMARTNOISE_SYNTH
    """Synthetic data generator model."""
    # model: Annotated[Synthesizer, PlainSerializer(serialize_model), PlainValidator(deserialize_model)]


class SmartnoiseSynthSamples(BaseModel):
    """Type for smartnoise_synth result when it is a dataframe of samples."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    type: Literal["sn_synth_samples"] = "sn_synth_samples"
    """Result type description."""
    df_samples: Annotated[
        pd.DataFrame,
        PlainSerializer(dataframe_to_dict),
        PlainValidator(dataframe_from_dict),
    ]
    """Dataframe containing the generated synthetic samples."""


# OpenDP
class OpenDPQueryResult(BaseModel):
    """Type for opendp result."""

    type: Literal[DPLibraries.OPENDP] = DPLibraries.OPENDP
    """Result type description."""
    value: int | float | list[int | float]
    """The result value of the query."""


class OpenDPPolarsQueryResult(BaseModel):
    """Type for opendp Polars result."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    type: Literal[DPLibraries.OPENDP_POLARS] = DPLibraries.OPENDP_POLARS
    """Result type description."""
    # order of PlainValidator and PlainSerializer matters in that case:
    # https://github.com/pydantic/pydantic/issues/8512
    value: Annotated[
        pl.DataFrame,
        PlainValidator(polars_df_from_str),
        PlainSerializer(polars_df_to_str),
    ]
    """The result value of the query."""


# Response object
QueryResultT = Annotated[
    DiffPrivLibQueryResult
    | SmartnoiseSQLQueryResult
    | SmartnoiseSynthModel
    | SmartnoiseSynthSamples
    | OpenDPQueryResult
    | OpenDPPolarsQueryResult,
    Discriminator("type"),
]


class QueryResponse(CostResponse):
    """Response to Lomas queries."""

    requested_by: str
    """The user that triggered the query."""
    result: QueryResultT
    """The query result object."""


class Job(ResponseModel):
    """Scheduled Job Response."""

    uid: UUID = Field(default_factory=uuid4)
    """Job unique identifier."""
    requested_by: str | None = None
    """Name of the user that requested this job."""
    status: Literal["in_progress", "failed", "complete"] = "in_progress"
    """Job status."""
    result: QueryResponse | CostResponse | None = None
    """Job result, if available."""
    error: LomasServerExceptionType | None = None
    """Job error, if any."""
    status_code: int = 200
    """Status code for job response."""
