from __future__ import annotations

from enum import StrEnum

from pydantic import BaseModel, ConfigDict, Field
from pydantic.alias_generators import to_camel


class Pd4castrBaseModel(BaseModel):
    model_config = ConfigDict(
        extra="ignore",
        alias_generator=to_camel,
        populate_by_name=True,
    )


class FileFormat(StrEnum):
    json = "json"
    csv = "csv"
    parquet = "parquet"


class ModelRunMode(StrEnum):
    automatic = "AUTOMATIC"
    on_demand = "ON_DEMAND"


class TimeHorizon(Pd4castrBaseModel):
    name: str
    short_name: str
    key: str


class ModelMetadata(Pd4castrBaseModel):
    description: str | None = None
    release_date: str | None = None
    resolution: str | None = None


class UISpecification(Pd4castrBaseModel):
    series_key: bool | None = None
    colours: list[str] | None = None


class ColumnSpecification(Pd4castrBaseModel):
    name: str
    key: str | None = None
    type: str
    ui: UISpecification | None = None


class OutputVariable(Pd4castrBaseModel):
    name: str
    key: str
    type: str
    series_key: bool | None = None
    colours: list[str] | None = None


class SensitivitySpecification(Pd4castrBaseModel):
    id: str
    name: str
    key: str


class ForecastVariableSpecification(Pd4castrBaseModel):
    name: str
    short_name: str
    key: str


class ModelInputSpecification(Pd4castrBaseModel):
    id: str
    key: str


class Model(Pd4castrBaseModel):
    id: str
    model_group_id: str
    name: str
    display_name: str
    notes: str
    revision: int
    docker_image: str
    created_at: str
    run_mode: ModelRunMode
    horizon: TimeHorizon
    model_metadata: ModelMetadata
    output_variables: list[OutputVariable] = []
    # DEPRECATED: use output_variables
    output_specification: list[ColumnSpecification] = []
    inputs: list[ModelInputSpecification] | None = None
    tags: list[str]
    sensitivities: list[SensitivitySpecification]
    output_file_format: FileFormat
    display_timezone: str
    forecast_variable: ForecastVariableSpecification


class ModelGroup(Pd4castrBaseModel):
    id: str
    name: str
    latest_revision: int
    public: bool
    created_at: str


class ModelRunStatus(StrEnum):
    pending = "pending"
    input_files_prepared = "input_files_prepared"
    running = "running"
    completed = "completed"
    failed = "failed"


class Sensitivity(Pd4castrBaseModel):
    id: str
    name: str
    key: str


class ModelRun(Pd4castrBaseModel):
    id: str
    status: ModelRunStatus
    created_at: str
    run_datetime: str
    started_at: str | None = None
    completed_at: str | None = None
    errored_at: str | None = None
    sensitivity: Sensitivity | None = None
    error: str | None = None
    container_id: str | None = None


class ModelRunPage(Pd4castrBaseModel):
    data: list[ModelRun]
    next_cursor: str | None = None
    has_more: bool


class ModelSummary(Pd4castrBaseModel):
    id: str
    display_name: str
    display_timezone: str


class ModelRunOutputData(Pd4castrBaseModel):
    model_config = ConfigDict(extra="allow")

    forecast_datetime: str = Field(alias="forecast_datetime")


class ModelRunWithOutput(Pd4castrBaseModel):
    id: str
    run_datetime: str
    sensitivity: Sensitivity | None = None
    model: ModelSummary
    colours: dict[str, str | None]
    data: list[ModelRunOutputData]


class ModelRunOutputResult(Pd4castrBaseModel):
    run: ModelRunWithOutput
    comparison_runs: dict[str, ModelRunWithOutput]


class Organisation(Pd4castrBaseModel):
    id: str
    display_name: str
    slug: str
    domains: list[str]
    permissions: list[str]


class User(Pd4castrBaseModel):
    id: str
    email: str
    organisation: Organisation | None = None
