from __future__ import annotations

from typing import Any

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 GenerateOptions(BaseModel):
    anime_enhance: int = 0
    mode: int = 0
    gen_mode: int = 0
    prompt_magic_mode: int = 2


class LabUpscaleOptions(BaseModel):
    refiner_switch: int = 0
    hdr_enabled: bool = False
    hdr_intensity: int = 0
    auto_brightness_open: bool = False
    auto_brightness_strength: int = 0
    auto_adjust_open: bool = False
    auto_adjust_strength: int = 0
    auto_adjust_brightness: int = 0
    film_turbo: int = 0
    expand_token_length: int = 0
    custom_height_width: bool = False
    semantics_open: bool = False
    sharpness: int = 0
    quality_open: bool = False
    quality_detail_offset: int = 0
    quality_type: str = ""
    quality_mode: int = 3
    intelligence_open: bool = False
    prompt_magic: bool = False


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

    channel_id: str = Field(default="")
    parent_task_id: str = Field(default="")
    speed_type: int = Field(default=1)
    pre_task_id: str = Field(default="")
    task_uri_index: int = Field(default=0)
    ss: int | None = None

    extra_prompt: str = Field(default="")
    local_prompt: str = Field(default="")

    size: SizeOptions = Field(default_factory=SizeOptions)
    sampling: SamplingOptions = Field(default_factory=SamplingOptions)
    batch_size: int = Field(default=1, ge=1)
    n_iter: int = Field(default=1, ge=1)

    smart_edit: bool | None = None
    guidance_scale: int = Field(default=0)
    left_margin: int = Field(default=0)
    up_margin: int = Field(default=0)

    image: str = Field(default="")
    images: list[Any] | None = None

    vae: str | None = None
    refiner_mode: int = Field(default=0)
    lcm_mode: int = Field(default=0)

    lab: LabUpscaleOptions = Field(default_factory=LabUpscaleOptions)

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

    generate: GenerateOptions = Field(default_factory=GenerateOptions)

    enable_hr: bool = Field(default=False)
    controlnet_units: list[Any] | None = None
