from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel

from .base import APIDataModel
from .history import (
    ActionCapability,
    DeductionDetail,
)

Upscaler = Literal[
    "4x-UltraSharp",
    "Lanczos",
    "Nearest",
    "R-ESRGAN 4x+",
    "R-ESRGAN 4x+ Anime6B",
    "8x_NMKD-Superscale_150000_G",
    "4x_Valar_v1",
]


# Request models
class TaskV2Body(BaseModel):
    model_no: str | None = None
    model_ver_no: str | None = None
    channel_id: str
    parent_task_id: str | None = None
    speed_type: int
    meta: Meta
    pre_task_id: str | None = None
    task_uri_index: int | None = None
    ss: int | None = None
    source: int | None = None
    pre_art_work_id: str | None = None


class ImageUrlBody(BaseModel):
    url: str | None = None
    uri: str | None = None


class Img2TextMeta(BaseModel):
    img_2_text: ImageUrlBody


class CreateImageToTextBody(BaseModel):
    meta: Img2TextMeta
    ss: int | None = None


class Input(BaseModel):
    field: str
    node_id: str
    node_type: str
    val: str


class CreateSmartEditImageBody(BaseModel):
    apply_id: str
    inputs: list[Input]
    is_use_unlimited_free: bool
    task_flow_version: str
    ss: int | None = None


# Response models
class Banner(APIDataModel):
    width: int
    height: int
    url: str
    nsfw: int
    nsfw_level: int | None = None
    is_nsfw_plus: bool
    green: bool | None = None
    has_meta: bool | None = None
    stream_url: str | None = None
    cover_url: str


class Ckpt(APIDataModel):
    id: str
    banner: Banner
    name: str
    model_no: str
    model_ver_no: str
    version_name: str
    model_type: str
    base_model_name: str
    weight: float | None = None
    tags: list[str] | None = None
    trained_words: list[str] | None = None


class TaskMetadata(APIDataModel):
    id: str
    prompt: str
    local_prompt: str
    type: int
    sub_type: int
    task_domain_type: int | None = None
    sub_mode_category: int
    width: int
    height: int
    ckpts: list[Ckpt]
    loras: list[Any]
    embeddings: list[Any]
    create_at: int


class HiResArg(APIDataModel):
    resize_mode: int | None = None
    show_extras_results: bool | None = None
    gfpgan_visibility: int | None = None
    codeformer_visibility: int | None = None
    codeformer_weight: int | None = None
    upscaling_resize: int | float
    upscaling_resize_w: int | None = None
    upscaling_resize_h: int | None = None
    upscaling_crop: bool | None = None
    upscaler_1: str
    upscaler_2: str | None = None
    extras_upscaler2visibility: int | None = None
    upscale_first: bool | None = None
    image: str


class Cond(APIDataModel):
    image: str
    stop_at: int | None = None
    weight: float
    mode: str
    loading: bool | None = None


class LabBase(APIDataModel):
    refiner_switch: int | None = None
    hdr_enabled: bool | None = None
    hdr_intensity: int | None = None
    auto_brightness_open: bool | None = None
    auto_brightness_strength: int | None = None
    auto_adjust_open: bool | None = None
    auto_adjust_strength: int | None = None
    auto_adjust_brightness: int | None = None
    film_turbo: int | None = None
    expand_token_length: int | None = None
    custom_height_width: bool | None = None
    semantics_open: bool | None = None
    sharpness: int | None = None
    quality_open: bool | None = None
    quality_detail_offset: int | None = None
    quality_type: str | None = None
    quality_mode: int | None = None
    intelligence_open: bool | None = None
    prompt_magic: bool | None = None
    conds: list[Cond] | None = None


class Generate(APIDataModel):
    anime_enhance: int
    mode: int
    gen_mode: int
    prompt_magic_mode: int


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


class Meta(APIDataModel):
    extra_prompt: str | None = None
    prompt: str | None = None
    local_prompt: str | None = None
    negative_prompt: str | None = None
    width: int
    height: int
    steps: int
    cfg_scale: float | None = None
    sampler_name: str | None = None
    batch_size: int | None = None
    init_images: list[Any] | None = None
    denoising_strength: float | None = None
    seed: int
    n_iter: int | None = None
    lora_models: list[Any] | None = None
    hi_res_arg: HiResArg | None = None
    smart_edit: Any | None = None
    guidance_scale: int | None = None
    left_margin: int | None = None
    up_margin: int | None = None
    image: str | None = None
    images: Any | None = None
    clip_skip: int | None = None
    vae: str | None = None
    refiner_mode: int | None = None
    lcm_mode: int | None = None
    lab_base: LabBase | None = None
    restore_faces: bool | None = None
    embeddings: list[Any] | None = None
    generate: Generate | None = None
    enable_hr: int | None = None
    controlnet_units: list[ControlnetUnit] | None = None


class ImageToTextResult(APIDataModel):
    id: str
    prompt: str
    local_prompt: str
    type: int
    sub_type: int
    sub_mode_category: int
    width: int
    height: int
    ckpts: Any
    loras: Any
    embeddings: Any
    create_at: int


class SmartGenerationId(APIDataModel):
    id: str


class ImageDownload(APIDataModel):
    task_id: str
    url: str


class UploadedImage(APIDataModel):
    url: str
    signed: str


class UploadedImageUrl(APIDataModel):
    access_url: str


class Stat(APIDataModel):
    num_of_like: int
    num_of_collection: int
    num_of_task: int
    num_of_view: int
    num_of_artwork: int
    rating: int


class ArtModel(APIDataModel):
    id: str
    banner: Banner
    author: Any | None = None
    account_no: str
    name: str
    nsfw: int
    model_ver_no: str
    version_name: str
    tags: list[str] | None = None
    trained_words: list[str] | None = None
    type: str
    status: int
    stat: Stat
    collect: Any | None = None
    base_model: str
    publish_status: int
    primary: int
    is_agency_fee: bool
    agency_fee: int
    payment_way: int
    channel: int
    channel_type: str
    rank_score: int
    is_preset: int
    sys_tag: Any | None = None
    pre_set_template_id: str


class ImgUri(APIDataModel):
    id: str
    width: int
    height: int
    cover_url: str
    title: str
    is_blur: bool
    url: str
    nsfw: int
    is_nsfw_plus: bool
    index: int
    expiration_time: int
    is_submit: bool
    is_saved_personal: bool
    rembg_x: int
    rembg_y: int
    node_key: str | None = None
    type: int
    sub_type: int
    operation_labels: list[str] | None = None
    action_capability: ActionCapability | None = None
    is_nsfw_appeal: bool | None = None
    is_video_has_audio: bool | None = None
    artwork_no: str | None = None
    nsfw_level: int | None = None


class ArtworkDetails(APIDataModel):
    id: str
    template_id: str
    artwork_no: str
    type: int
    sub_type: int
    retry_target_type: int
    art_model_no: str
    art_model_ver_no: str
    base_model_name: str
    art_model_name: str
    art_model: ArtModel | None = None
    meta: Meta
    local_prompt: str
    category: int
    create_at: int
    img_uris: list[ImgUri] | None = None
    status: int
    status_desc: str
    deduction_detail: DeductionDetail | None = None
    task_domain_type: int
    retry_target_domain_type: int
