from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field


class SamplingOptions(BaseModel):
    steps: int = Field(default=25, ge=1)
    sampler_name: str = Field(default="Euler")
    cfg_scale: float = Field(default=6, ge=0)
    seed: int | None = None


class SizeOptions(BaseModel):
    width: int = Field(default=1024, ge=1)
    height: int = Field(default=1024, ge=1)


class ConditioningOptions(BaseModel):
    mode: Literal["ImagePrompt"] | str = Field(default="ImagePrompt")
    weight: float = Field(default=0.35, ge=0, le=1)
    stop_at: int = Field(default=1, ge=0)
    loading: bool = False


class GenerateOptions(BaseModel):
    anime_enhance: int = 2
    mode: int = 0
    gen_mode: int = 0
    prompt_magic_mode: int = 2


class EstimateOptions(BaseModel):
    """Options for the ``estimate`` task method. All fields are optional and have sensible defaults."""

    channel_id: str = Field(default="")
    speed_type: int = Field(default=1)
    negative_prompt: str = Field(default="")
    n_iter: int = Field(default=1, ge=1)

    size: SizeOptions = Field(default_factory=SizeOptions)
    sampling: SamplingOptions = Field(default_factory=SamplingOptions)
    conditioning: ConditioningOptions = Field(default_factory=ConditioningOptions)
    generate: GenerateOptions = Field(default_factory=GenerateOptions)

    lora_models: list[Any] = Field(default_factory=list)
    embeddings: list[Any] = Field(default_factory=list)
