from __future__ import annotations

from typing import Any

from pydantic import BaseModel

from .base import APIDataModel
from .history import DeductionDetail
from .images import (
    Banner,
    Generate,
    HiResArg,
    ImgUri,
    Meta,
)


# Request models
class ComfyuiInput(BaseModel):
    node_id: str | None = None
    node_type: str | None = None
    field: str | None = None
    val: str | None = None
    type: int | None = None
    extra_data: None


class RemoveBackground(BaseModel):
    uri: str | None = None


class CreateTaskMetadata(BaseModel):
    extra_prompt: str | None = None
    prompt: str | None = None
    local_prompt: str | None = None
    width: int | None = None
    height: int | None = None
    steps: int | None = None
    batch_size: int | None = None
    init_images: Any | None = None
    seed: int | None = None
    n_iter: int | None = None
    cfg_scale: float | None = None
    lora_models: Any | 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: list[str] | None = None
    vae: str | None = None
    refiner_mode: int | None = None
    lcm_mode: int | None = None
    embeddings: Any | None = None
    generate: Generate | None = None
    comfyui_inputs: list[ComfyuiInput] | None = None
    remove_background: RemoveBackground | None = None
    hi_res_arg: HiResArg | None = None
    sampler_name: str | None = None
    clip_skip: int | None = None


class CreateTaskBody(BaseModel):
    action: int
    art_work_no: str | None = None
    parent_task_id: str | None = None
    category: int
    meta: CreateTaskMetadata
    task_flow_version: str | None = None
    pre_task_id: str | None = None
    pre_art_work_id: str | None = None
    task_uri_index: int | None = None
    channel_id: str | None = None
    ss: int | None = None


# Response models
class TaskReturn(APIDataModel):
    sub_status: int
    used: int
    hashrate: float
    temp_hashrate: float
    free_num: int
    remain_num: int
    is_can_guide: bool
    task_run_time: int
    task_push_timestamp: int
    task_queue_position: int
    total_queue_position: int


class NewTask(APIDataModel):
    id: str
    artwork_no: str
    type: int
    art_model_no: str
    art_model_name: str
    art_model_ver_no: str
    base_model_name: str
    meta: Meta
    local_prompt: str
    process: int
    status: int
    status_desc: str
    create_at: int
    img_uris: Any
    reverse_prompts: Any
    reverse_local_prompts: Any
    channel_id: str
    category: int
    deduction_detail: Any
    pub_community: int
    expiration_time: int
    speed_type: int
    banner: Banner | None = None
    nsfw: int
    pub_artwork_no: str
    pub_artwork_nos: Any
    collect: Any
    preprocess_images: Any


class TaskOperationResult(APIDataModel):
    new_task: NewTask
    task_return: TaskReturn


class GuideInfo(APIDataModel):
    id: str
    is_show: bool
    title: str
    content: str
    button: str
    media_url: str
    url: str


class EstimateTime(APIDataModel):
    min_time: int
    max_time: int
    avg_time: float


class TaskItem(APIDataModel):
    task_id: str
    type: int
    sub_type: int
    task_domain_type: int
    process: int

    img_url: str
    banner: Any | None = None

    status: int
    status_desc: str
    sub_status: int

    pub_artwork_no: str
    pub_artwork_nos: list[str] | None = None

    nsfw: int
    img_uris: list[ImgUri] | None = None
    reverse_prompts: list[str] | None = None

    category: int
    deduction_detail: DeductionDetail

    pub_community: int
    expiration_time: int
    create_at: int
    speed_type: int
    reason: int
    total_num: int
    current_num: int

    task_return: TaskReturn
    estimate_time: EstimateTime
    is_show: bool
    guide_info: GuideInfo

    reverse_local_prompts: list[str] | None = None

    @property
    def percent(self) -> int:
        """Generation progress as a percentage (0-100), aliased from ``process``."""
        return self.process

    @property
    def is_done(self) -> bool:
        """``True`` when generation progress has reached 100 %."""
        return self.process >= 100

    @property
    def image_url(self) -> str | None:
        """URL of the first generated image, or ``None`` if not yet available."""
        if self.img_uris:
            return self.img_uris[0].url
        return None

    @property
    def image_urls(self) -> list[str]:
        """URLs of all generated images; empty list if not yet available."""
        if self.img_uris:
            return [uri.url for uri in self.img_uris]
        return []

    @property
    def reverse_prompt(self) -> str | None:
        """Reverse-engineered prompt describing the generated image, or ``None`` if absent.

        Returns the second element of ``reverse_prompts`` when more than one
        is present, otherwise the first.
        """
        if self.reverse_prompts:
            return (
                self.reverse_prompts[1]
                if len(self.reverse_prompts) > 1
                else self.reverse_prompts[0]
            )
        return None


class TasksCollection(APIDataModel):
    items: list[TaskItem]


class TaskSpeedUpResult(APIDataModel):
    num: int
