from __future__ import annotations

from datetime import datetime
from typing import Literal

from pydantic import BaseModel, Field, model_validator
from one_run.constants import JSON_SCHEMA_VERSION


class ExperimentConfig(BaseModel):
    name: str = Field(min_length=1)
    seed: int
    tags: list[str] | None = None
    color: Literal["black", "red", "green", "yellow", "blue", "magenta", "cyan", "white"] | None = None


class EntrypointConfig(BaseModel):
    command: str = Field(min_length=1)


class EnvironmentConfig(BaseModel):
    kind: Literal["local", "docker"]
    image: str | None = None
    workdir: str | None = None
    mounts: list[str] | None = None
    env: dict[str, str] | None = None
    env_file: str | None = None
    pull_policy: Literal["always", "if_missing", "never"] | None = None
    labels: dict[str, str] | None = None

    @model_validator(mode="after")
    def validate_docker(self) -> "EnvironmentConfig":
        if self.kind == "docker" and not self.image:
            raise ValueError("environment.image is required when environment.kind=docker")
        return self


class ResourcesConfig(BaseModel):
    gpus: int | None = None
    cpus: int | None = None
    mem_gb: int | None = None
    time_min: int | None = None


class SlurmConfig(BaseModel):
    partition: str | None = None
    retry_max_attempts: int | None = None
    retry_backoff_sec: int | None = None
    retry_on_states: list[Literal["failed", "timeout", "cancelled"]] | None = None


class BackendConfig(BaseModel):
    kind: Literal["local", "slurm"]
    slurm: SlurmConfig | None = None


class ScheduleConfig(BaseModel):
    cron: str = Field(min_length=1)
    enabled: bool = True


class DependsOnConfig(BaseModel):
    run_ids: list[str] | None = None
    tags: list[str] | None = None
    timeout: str | None = None
    interval_sec: float = 2.0

    @model_validator(mode="after")
    def validate_depends_on(self) -> "DependsOnConfig":
        has_run_ids = bool(self.run_ids)
        has_tags = bool(self.tags)
        if not has_run_ids and not has_tags:
            raise ValueError("depends_on requires run_ids and/or tags")
        if self.interval_sec <= 0:
            raise ValueError("depends_on.interval_sec must be > 0")
        return self


class ListenerConfig(BaseModel):
    kind: Literal["local", "webhook"]
    name: str | None = None
    action: str | None = None
    events: list[str] | None = None
    states: list[str] | None = None
    tags: list[str] | None = None
    handler: str | None = None
    webhook_url: str | None = None
    timeout_sec: float | None = None

    @model_validator(mode="after")
    def validate_listener(self) -> "ListenerConfig":
        if self.kind == "webhook" and not self.webhook_url:
            raise ValueError("listeners[].webhook_url is required when kind=webhook")
        if self.timeout_sec is not None and self.timeout_sec <= 0:
            raise ValueError("listeners[].timeout_sec must be > 0")
        return self


class Manifest(BaseModel):
    experiment: ExperimentConfig
    entrypoint: EntrypointConfig
    environment: EnvironmentConfig
    resources: ResourcesConfig | None = None
    backend: BackendConfig
    schedule: ScheduleConfig | None = None
    depends_on: DependsOnConfig | None = None
    listeners: list[ListenerConfig] | None = None


class RunMetadata(BaseModel):
    schema_version: str = JSON_SCHEMA_VERSION
    run_id: str
    timestamp: datetime
    hostname: str
    cwd: str
    project_root: str
    git_commit: str | None
    git_dirty: bool | None
    experiment_name: str | None = None
    tags: list[str] | None = None
    experiment_color: Literal["black", "red", "green", "yellow", "blue", "magenta", "cyan", "white"] | None = None
    backend_type: Literal["local", "slurm"]
    docker_image: str | None
    docker_image_digest: str | None = None
    slurm_job_id: str | None = None
    slurm_attempt: int | None = None


class RunStatus(BaseModel):
    schema_version: str = JSON_SCHEMA_VERSION
    state: Literal["pending", "running", "success", "failed"]
    started_at: datetime | None = None
    ended_at: datetime | None = None
    exit_code: int | None = None
    message: str | None = None


class RunRuntime(BaseModel):
    schema_version: str = JSON_SCHEMA_VERSION
    collected_at: datetime
    duration_sec: float | None = None
    peak_rss_mb: float | None = None
    cpu_count: int | None = None
    total_ram_gb: float | None = None
    slurm_elapsed: str | None = None
    slurm_max_rss: str | None = None
    slurm_alloc_cpus: int | None = None
