from __future__ import annotations

from typing import Any

from .base import APIDataModel

_DEFAULT_SIZES: dict[str, tuple[int, int]] = {
    "1:2": (512, 1024),
    "9:16": (592, 1024),
    "2:3": (688, 1024),
    "3:4": (768, 1024),
    "4:5": (768, 960),
    "2:1": (1024, 512),
    "16:9": (1024, 592),
    "3:2": (1024, 688),
    "4:3": (1024, 768),
    "5:4": (960, 768),
    "1:1": (1024, 1024),
}


# Response models
class ModelInfoV2(APIDataModel):
    id: str
    desc: str


class Meta(APIDataModel):
    seed: int
    guidance_scale: int
    local_prompt: str
    width: int
    height: int
    smart_edit: Any
    images: Any
    extra_prompt: str
    prompt: str
    init_images: Any
    refiner_mode: int
    embeddings: Any
    steps: int
    left_margin: int
    up_margin: int
    lcm_mode: int
    generate: Any
    lora_models: Any
    hi_res_arg: Any
    image: str


class TaskOperation(APIDataModel):
    is_make_variations: bool
    is_img_2_img: bool
    is_upscale: bool
    is_animation: bool
    is_extra_upscale: bool
    is_remove_background: bool
    is_video_extend: bool
    is_digital_human_chat: bool
    is_img_2_video: bool


class Sample(APIDataModel):
    green: int
    support_remix: bool
    id: str
    url: str
    width: int
    height: int
    generation_process: str
    support_detail: bool
    cover: str
    nsfw: int
    meta: Meta
    collect: Any
    task_operation: TaskOperation


class RecGenerateParam(APIDataModel):
    width: int | None = None
    prompt: str
    sampler: str
    vae: str | None = None
    cfg_scale: float
    height: int | None = None
    negative_prompt: str | None = None
    steps: int
    clip_skip: int


class TrainInfo(APIDataModel):
    epoch: int
    step: int
    clip_skip: int
    train_ui_mode: int


class ModelVerInfoV2(APIDataModel):
    base_model_name: str
    model_type: str | None = None
    rec_generate_param: RecGenerateParam | None = None
    train_info: TrainInfo | None = None
    trigger_words: list[str]
    model_id: str
    desc: str
    create_at: int
    publish_at: int
    update_at: int
    rating: float
    model_ver_id: str
    name: str


class Same(APIDataModel):
    primary: int
    train_model_no: str
    train_ui_mode: int


class ModelVerNo(APIDataModel):
    id: str
    version_name: str
    same: Same


class Author(APIDataModel):
    head: str
    avatar_frame: str
    is_follow: bool
    follower_cnt: int
    num_of_view: int
    id: str
    name: str
    status: int
    num_of_collect: int
    is_active: bool
    join_stimulate: int


class Tag(APIDataModel):
    show_in: int
    id: str
    title: str
    cover: str
    is_collect: bool


class Interaction(APIDataModel):
    rating: float
    rating_time: int


class TrainParam(APIDataModel):
    epoch: int
    step: int
    clip_skip: int
    train_ui_mode: int


class Stat(APIDataModel):
    num_of_artwork: int
    rating: float
    num_of_comment: int | None = None
    num_of_rating: int
    num_of_like: int
    num_of_collection: int
    num_of_task: int
    num_of_view: int


class ModelDetails(APIDataModel):
    model_info_v2: ModelInfoV2
    create_at: int
    publish_status: int
    update_at: int
    commercial: int
    hashs: Any
    is_auditing: int
    cover: str
    desc: str
    samples: list[Sample]
    allow_use_in_haiyi: int
    model_ver_no: str
    model_ver_info_v2: ModelVerInfoV2
    origin_model: Any
    type: str
    file_id: str
    is_agency_fee: bool
    nsfw_level: int
    channel: int
    allow_creation: int
    can_subscribe: int
    rank_level_list: list[Any]
    audit_status: Any
    download: int
    cover_width: int
    ver: str
    meta: Any
    payment_way: int
    version_name: str
    categories: Any
    status: int
    nsfw: int
    allow_sell: int
    appeal_count: int
    cover_v2: str
    model_ver_nos: list[ModelVerNo]
    mix: int
    agency_fee: int
    id: str
    name: str
    author: Author
    pre_set_meta: Any
    sys_tag: list[str] | None = None
    tags: list[Tag]
    refer: str
    allow_resell: int
    interaction: Interaction
    collect: Any
    cover_height: int
    sub_type: int
    train_param: TrainParam
    trained_words: list[str] | None = None
    base_model: str
    stat: Stat
    is_preset: int


class AdditionalModelPanel(APIDataModel):
    enable_panel: bool
    support_base_model_names: list[str]
    component_block_items: list[Any]


class AllOption(APIDataModel):
    aspect_ratio: str
    width: int
    height: int
    quality_width: int
    quality_height: int
    max_width: int
    max_height: int
    min_width: int
    min_height: int


class ResolutionParamItem(APIDataModel):
    max_width: int
    max_height: int
    min_width: int
    min_height: int
    dft_resolution_option: Any
    dft_resolution_options: Any
    options: Any
    all_options: list[AllOption]
    enable_custom_resolution: bool
    enable_more_custom_resolution: bool

    def get_size(self, aspect_ratio: str) -> tuple[int, int] | None:
        """Return ``(width, height)`` for the given aspect ratio string.

        Checks the server-provided ``all_options`` list first; falls back to
        the built-in ``_DEFAULT_SIZES`` table if the ratio is not listed.

        Args:
            aspect_ratio: Ratio string in ``"W:H"`` format, e.g. ``"16:9"``.

        Returns:
            A ``(width, height)`` pixel tuple, or ``None`` if the ratio is
            not found in either the server options or the default table.
        """
        for option in self.all_options:
            if option.aspect_ratio == aspect_ratio:
                return option.width, option.height
        return _DEFAULT_SIZES.get(aspect_ratio)


class ComponentBlockItem(APIDataModel):
    component_type: str
    enable_block: bool
    value: bool | None = None
    value_type: int
    max: int
    min: int
    step: int
    accuracy: int
    option_values: Any
    resolution_param: ResolutionParamItem | None = None


class BasicPanel(APIDataModel):
    enable_panel: bool
    component_block_items: list[ComponentBlockItem]


class ComponentBlockItem1(APIDataModel):
    component_type: str
    enable_block: bool
    value: bool | int | str | None = None
    value_type: int
    max: int
    min: int
    step: float
    accuracy: int
    option_values: list[str] | None = None
    resolution_param: Any


class AdvancedPanel(APIDataModel):
    enable_panel: bool
    component_block_items: list[ComponentBlockItem1]


class ComponentBlockItem2(APIDataModel):
    component_type: str
    enable_block: bool
    value: Any
    value_type: int
    max: int
    min: int
    step: int
    accuracy: int
    option_values: Any
    resolution_param: Any


class ImageGenParamPanel(APIDataModel):
    enable_panel: bool
    component_block_items: list[ComponentBlockItem2]


class ImgParam(APIDataModel):
    enable_config_option: bool
    additional_model_panel: AdditionalModelPanel
    basic_panel: BasicPanel
    advanced_panel: AdvancedPanel
    image_gen_param_panel: ImageGenParamPanel


class GenerationParameters(APIDataModel):
    model_type: str | None = None
    max_niter: int | None = None
    t2v_param: Any | None = None
    i2v_param: Any | None = None
    img_param: ImgParam | None = None
    v2v_param: Any | None = None
    i2v_model_param: str | None = None
    t2v_model_param: str | None = None
    video_edit_param: Any | None = None
    gen_num_param: Any | None = None
    resolution_param: Any | None = None
    hide_img_cond_panel: bool | None = None

    def get_size(self, aspect_ratio: str) -> tuple[int, int] | None:
        """Return ``(width, height)`` for the given aspect ratio string.

        Uses resolution options from ``img_param`` when available; falls back
        to the built-in ``_DEFAULT_SIZES`` table otherwise.

        Args:
            aspect_ratio: Ratio string in ``"W:H"`` format, e.g. ``"16:9"``.

        Returns:
            A ``(width, height)`` pixel tuple, or ``None`` if the ratio is
            not found in either the model's resolution options or the default
            table.
        """
        if self.img_param is not None:
            resolution_block = next(
                (
                    block
                    for block in self.img_param.basic_panel.component_block_items
                    if block.resolution_param is not None
                ),
                None,
            )
            if (
                resolution_block is not None
                and resolution_block.resolution_param is not None
            ):
                return resolution_block.resolution_param.get_size(aspect_ratio)

        return _DEFAULT_SIZES.get(aspect_ratio)
