from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field


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


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)


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


class ControlnetUnitOptions(BaseModel):
    control_mode: int | None = None
    guidance: int | None = None
    guidance_end: int | None = None
    input_image: str | None = None
    model: str | None = None
    module: str | None = None
    processor_res: int | None = None
    threshold_a: Any | None = None
    type: str | None = None
    weight: float | None = None


class ControlnetOptions(BaseModel):
    units: list[ControlnetUnitOptions] | None = None


class TextToImgOptions(BaseModel):
    """Options for the ``text_to_img`` generation 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="")

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

    clip_skip: int = Field(default=2)
    vae: str | None = None
    restore_faces: bool = False

    controlnet: ControlnetOptions = Field(default_factory=ControlnetOptions)

    source: int | None = None
    ss: int | None = None

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

    # @classmethod
    # def fast(cls) -> TextToImgOptions: #TODO: is it necessary
    #     """Return a preset with a reduced step count for faster generation."""
    #     return cls(sampling=SamplingOptions(steps=20))
