from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_validator


NonEmptyStr = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
PositiveInt = Annotated[int, Field(gt=0)]
NonNegativeInt = Annotated[int, Field(ge=0)]
PositiveFloat = Annotated[float, Field(gt=0)]
Probability = Annotated[float, Field(ge=0, le=1)]
NonNegativeFloat = Annotated[float, Field(ge=0)]


class RLXModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class EnvConfig(RLXModel):
    id: NonEmptyStr
    num_envs: PositiveInt


class AlgoConfig(RLXModel):
    name: Literal["ppo"]
    total_timesteps: PositiveInt
    rollout_steps: PositiveInt
    batch_size: PositiveInt
    learning_rate: PositiveFloat
    gamma: Probability
    gae_lambda: Probability
    clip_range: NonNegativeFloat
    entropy_coef: NonNegativeFloat
    value_coef: NonNegativeFloat
    update_epochs: PositiveInt


class PolicyConfig(RLXModel):
    type: Literal["mlp"]
    hidden_sizes: list[PositiveInt] = Field(min_length=1)

    @field_validator("hidden_sizes")
    @classmethod
    def validate_hidden_sizes(cls, value: list[int]) -> list[int]:
        if not value:
            raise ValueError("must contain at least one layer size")
        return value


class CheckpointConfig(RLXModel):
    save_every: PositiveInt


class EvalConfig(RLXModel):
    every: PositiveInt
    episodes: PositiveInt
    deterministic: bool


class ExperimentConfig(RLXModel):
    run_name: NonEmptyStr
    seed: NonNegativeInt
    device: NonEmptyStr
    env: EnvConfig
    algo: AlgoConfig
    policy: PolicyConfig
    checkpoint: CheckpointConfig
    eval: EvalConfig

