from __future__ import annotations

from typing import TYPE_CHECKING

from ...options import (
    ControlnetToImgOptions,
    EstimateOptions,
    ImgToImgOptions,
    RemoveBgOptions,
    TextToImgOptions,
    UpscaleHdFixOptions,
    UpscaleOptions,
)
from ...options.upscale_images import (
    GenerateOptions,
    LabUpscaleOptions,
    SamplingOptions,
    SizeOptions,
)
from .._schemas.images import (
    Cond,
    ControlnetUnit,
    CreateImageToTextBody,
    CreateSmartEditImageBody,
    Generate,
    HiResArg,
    ImageUrlBody,
    Img2TextMeta,
    Input,
    LabBase,
    Meta,
    TaskV2Body,
)
from .._schemas.tasks import (
    CreateTaskBody,
    CreateTaskMetadata,
    RemoveBackground,
)

if TYPE_CHECKING:
    from ..._enums import ImageSize


def build_text_to_img_body(
    *,
    model_no: str,
    model_ver_no: str,
    prompt: str,
    image_url: str | None,
    seed_value: int,
    gen_mode: str,
    image_size: ImageSize,
    num_images: int,
    free_creation: bool | None,
    prompt_magic: bool | None,
    options: TextToImgOptions,
) -> TaskV2Body:
    gen_mode_value: int
    gen_mode_value = 1 if gen_mode == "quality" else 0

    speed_type_value: int
    if free_creation is None:
        speed_type_value = 1
    elif free_creation is False:
        speed_type_value = 2
    elif free_creation is True:
        speed_type_value = 3

    prompt_magic_mode_value: int
    if prompt_magic is None:
        prompt_magic_mode_value = 2
    elif prompt_magic is False:
        prompt_magic_mode_value = 1
    elif prompt_magic is True:
        prompt_magic_mode_value = 0

    return TaskV2Body(
        model_no=model_no,
        model_ver_no=model_ver_no,
        channel_id=options.channel_id,
        speed_type=speed_type_value,
        meta=Meta(
            prompt=prompt,
            negative_prompt=options.negative_prompt,
            width=image_size.width,
            height=image_size.height,
            steps=options.sampling.steps,
            cfg_scale=options.sampling.cfg_scale,
            sampler_name=options.sampling.sampler_name,
            n_iter=num_images,
            lora_models=options.lora_models,
            vae=options.vae,
            clip_skip=options.clip_skip,
            controlnet_units=[
                ControlnetUnit(
                    control_mode=cn.control_mode,
                    guidance=cn.guidance,
                    guidance_end=cn.guidance_end,
                    input_image=cn.input_image,
                    model=cn.model,
                    module=cn.module,
                    processor_res=cn.processor_res,
                    threshold_a=cn.threshold_a,
                    type=cn.type,
                    weight=cn.weight,
                )
                for cn in options.controlnet.units
            ]
            if options.controlnet.units
            else None,
            seed=seed_value,
            restore_faces=options.restore_faces,
            embeddings=options.embeddings,
            lab_base=LabBase(
                conds=(
                    [
                        Cond(
                            image=image_url,
                            weight=options.conditioning.weight,
                            mode=options.conditioning.mode,
                            stop_at=options.conditioning.stop_at,
                        )
                    ]
                    if image_url
                    else []
                )
            ),
            generate=Generate(
                anime_enhance=options.generate.anime_enhance,
                mode=options.generate.mode,
                gen_mode=gen_mode_value,
                prompt_magic_mode=prompt_magic_mode_value,
            ),
        ),
        ss=options.ss,
        source=options.source,
    )


def build_img_to_img_body(
    *,
    model_no: str,
    model_ver_no: str,
    prompt: str,
    image_url: str,
    strength: float,
    seed_value: int,
    prompt_magic: bool | None,
    image_size: ImageSize,
    num_images: int,
    free_creation: bool | None,
    options: ImgToImgOptions,
) -> TaskV2Body:
    speed_type_value: int
    if free_creation is None:
        speed_type_value = 1
    elif free_creation is False:
        speed_type_value = 2
    elif free_creation is True:
        speed_type_value = 3

    prompt_magic_mode_value: int
    if prompt_magic is None:
        prompt_magic_mode_value = 2
    elif prompt_magic is False:
        prompt_magic_mode_value = 1
    elif prompt_magic is True:
        prompt_magic_mode_value = 0

    return TaskV2Body(
        model_no=model_no,
        model_ver_no=model_ver_no,
        channel_id=options.channel_id,
        speed_type=speed_type_value,
        meta=Meta(
            prompt=prompt,
            negative_prompt=options.negative_prompt,
            width=image_size.width,
            height=image_size.height,
            steps=options.sampling.steps,
            cfg_scale=options.sampling.cfg_scale,
            sampler_name=options.sampling.sampler_name,
            n_iter=num_images,
            lora_models=options.lora_models,
            vae=options.vae,
            clip_skip=options.clip_skip,
            seed=seed_value,
            init_images=[image_url],
            denoising_strength=strength,
            embeddings=options.embeddings,
            generate=Generate(
                anime_enhance=options.generate.anime_enhance,
                mode=options.generate.mode,
                gen_mode=options.generate.gen_mode,
                prompt_magic_mode=prompt_magic_mode_value,
            ),
        ),
        ss=options.ss,
    )


def build_img_to_text_body(
    *,
    url: str,
    ss: int | None = None,
) -> CreateImageToTextBody:
    return CreateImageToTextBody(
        meta=Img2TextMeta(img_2_text=ImageUrlBody(uri=url)), ss=ss
    )


def build_controlnet_to_img_body(
    *,
    model_no: str,
    model_ver_no: str,
    prompt: str,
    control_input_image: str,
    weight: float,
    seed_value: int,
    gen_mode: str,
    image_size: ImageSize,
    num_images: int,
    free_creation: bool | None,
    prompt_magic: bool | None,
    options: ControlnetToImgOptions,
) -> TaskV2Body:
    cn = options.controlnet

    gen_mode_value: int
    gen_mode_value = 1 if gen_mode == "quality" else 0

    speed_type_value: int
    if free_creation is None:
        speed_type_value = 1
    elif free_creation is False:
        speed_type_value = 2
    elif free_creation is True:
        speed_type_value = 3

    prompt_magic_mode_value: int
    if prompt_magic is None:
        prompt_magic_mode_value = 2
    elif prompt_magic is False:
        prompt_magic_mode_value = 1
    elif prompt_magic is True:
        prompt_magic_mode_value = 0

    return TaskV2Body(
        model_no=model_no,
        model_ver_no=model_ver_no,
        channel_id=options.channel_id,
        speed_type=speed_type_value,
        pre_art_work_id=options.pre_art_work_id,
        meta=Meta(
            prompt=prompt,
            negative_prompt=options.negative_prompt,
            width=image_size.width,
            height=image_size.height,
            steps=options.sampling.steps,
            cfg_scale=options.sampling.cfg_scale,
            sampler_name=options.sampling.sampler_name,
            n_iter=num_images,
            lora_models=options.lora_models,
            vae=options.vae,
            clip_skip=options.clip_skip,
            controlnet_units=[
                ControlnetUnit(
                    control_mode=cn.control_mode,
                    guidance=cn.guidance,
                    guidance_end=cn.guidance_end,
                    input_image=control_input_image,
                    model=cn.model,
                    module=cn.module,
                    processor_res=cn.processor_res,
                    threshold_a=cn.threshold_a,
                    type=cn.type,
                    weight=weight,
                )
            ],
            seed=seed_value,
            embeddings=options.embeddings,
            generate=Generate(
                anime_enhance=options.generate.anime_enhance,
                mode=options.generate.mode,
                gen_mode=gen_mode_value,
                prompt_magic_mode=prompt_magic_mode_value,
            ),
        ),
        ss=options.ss,
    )


def build_smart_edit_body(
    *,
    prompt: str,
    image_url: str,
    apply_id: str,
    is_use_unlimited_free: bool,
    task_flow_version: str,
    node_type_2: str,
    ss: int | None = None,
) -> CreateSmartEditImageBody:
    return CreateSmartEditImageBody(
        apply_id=apply_id,
        inputs=[
            Input(
                field="image",
                node_id="38",
                node_type="LoadImage",
                val=image_url,
            ),
            Input(
                field="prompt",
                node_id="89",
                node_type=node_type_2,
                val=prompt,
            ),
        ],
        is_use_unlimited_free=is_use_unlimited_free,
        task_flow_version=task_flow_version,
        ss=ss,
    )


def build_upscale_body(
    *,
    model_no: str,
    model_ver_no: str,
    prompt: str,
    image_url: str,
    options: UpscaleOptions,
    seed_value: int,
) -> TaskV2Body:
    return TaskV2Body(
        model_no=model_no,
        model_ver_no=model_ver_no,
        channel_id=options.channel_id,
        parent_task_id=options.parent_task_id,
        speed_type=options.speed_type,
        pre_task_id=options.pre_task_id,
        task_uri_index=options.task_uri_index,
        meta=Meta(
            extra_prompt=options.extra_prompt,
            prompt=prompt,
            local_prompt=options.local_prompt,
            width=options.size.width,
            height=options.size.height,
            steps=options.sampling.steps,
            cfg_scale=options.sampling.cfg_scale,
            sampler_name=options.sampling.sampler_name,
            batch_size=options.batch_size,
            init_images=[image_url],
            seed=seed_value,
            n_iter=options.n_iter,
            lora_models=options.lora_models,
            smart_edit=options.smart_edit,
            guidance_scale=options.guidance_scale,
            left_margin=options.left_margin,
            up_margin=options.up_margin,
            image=options.image,
            images=options.images,
            vae=options.vae,
            refiner_mode=options.refiner_mode,
            lcm_mode=options.lcm_mode,
            lab_base=LabBase(
                refiner_switch=options.lab.refiner_switch,
                hdr_enabled=options.lab.hdr_enabled,
                hdr_intensity=options.lab.hdr_intensity,
                auto_brightness_open=options.lab.auto_brightness_open,
                auto_brightness_strength=options.lab.auto_brightness_strength,
                auto_adjust_open=options.lab.auto_adjust_open,
                auto_adjust_strength=options.lab.auto_adjust_strength,
                auto_adjust_brightness=options.lab.auto_adjust_brightness,
                film_turbo=options.lab.film_turbo,
                expand_token_length=options.lab.expand_token_length,
                custom_height_width=options.lab.custom_height_width,
                semantics_open=options.lab.semantics_open,
                sharpness=options.lab.sharpness,
                quality_open=options.lab.quality_open,
                quality_detail_offset=options.lab.quality_detail_offset,
                quality_type=options.lab.quality_type,
                quality_mode=options.lab.quality_mode,
                intelligence_open=options.lab.intelligence_open,
                prompt_magic=options.lab.prompt_magic,
            ),
            embeddings=options.embeddings,
            generate=Generate(
                anime_enhance=options.generate.anime_enhance,
                mode=options.generate.mode,
                gen_mode=options.generate.gen_mode,
                prompt_magic_mode=options.generate.prompt_magic_mode,
            ),
            enable_hr=options.enable_hr,
            controlnet_units=options.controlnet_units,
        ),
        ss=options.ss,
    )


def build_upscale_hd_fix_body(
    *,
    image_url: str,
    upscaler_1: str,
    upscaling_resize: float,
    options: UpscaleHdFixOptions,
) -> CreateTaskBody:
    return CreateTaskBody(
        action=6,
        category=options.category,
        meta=CreateTaskMetadata(
            width=options.width,
            height=options.height,
            hi_res_arg=HiResArg(
                image=image_url,
                upscaler_1=upscaler_1,
                upscaling_resize=upscaling_resize,
            ),
        ),
        task_flow_version=options.task_flow_version,
        pre_task_id=options.pre_task_id or None,
        pre_art_work_id=options.pre_art_work_id or None,
        task_uri_index=options.task_uri_index,
        ss=options.ss,
    )


def build_remove_bg_body(
    *,
    image_url: str,
    options: RemoveBgOptions,
) -> CreateTaskBody:
    return CreateTaskBody(
        action=7,
        art_work_no=options.art_work_no or None,
        parent_task_id=options.parent_task_id,
        category=options.category,
        meta=CreateTaskMetadata(
            extra_prompt=options.extra_prompt,
            prompt=options.prompt,
            local_prompt=options.local_prompt,
            width=options.width,
            height=options.height,
            steps=options.steps,
            cfg_scale=options.cfg_scale,
            batch_size=options.batch_size,
            init_images=options.init_images,
            seed=options.seed,
            n_iter=options.n_iter,
            lora_models=options.lora_models,
            smart_edit=options.smart_edit,
            guidance_scale=options.guidance_scale,
            left_margin=options.left_margin,
            up_margin=options.up_margin,
            image=options.image,
            images=None,
            vae=options.vae,
            refiner_mode=options.refiner_mode,
            lcm_mode=options.lcm_mode,
            embeddings=options.embeddings,
            generate=Generate(
                anime_enhance=options.generate.anime_enhance,
                mode=options.generate.mode,
                gen_mode=options.generate.gen_mode,
                prompt_magic_mode=options.generate.prompt_magic_mode,
            ),
            remove_background=RemoveBackground(uri=image_url),
            sampler_name=options.sampler_name,
            clip_skip=options.clip_skip,
            hi_res_arg=options.hi_res_arg,
        ),
        channel_id=options.channel_id,
        ss=options.ss,
    )


def build_estimate_body(
    *,
    model_no: str | None,
    model_ver_no: str | None,
    prompt: str,
    image_url: str | None,
    options: EstimateOptions,
    seed_value: int,
) -> TaskV2Body:
    return TaskV2Body(
        model_no=model_no,
        model_ver_no=model_ver_no,
        channel_id=options.channel_id,
        speed_type=options.speed_type,
        meta=Meta(
            cfg_scale=options.sampling.cfg_scale,
            width=options.size.width,
            height=options.size.height,
            n_iter=options.n_iter,
            prompt=prompt,
            seed=seed_value,
            steps=options.sampling.steps,
            sampler_name=options.sampling.sampler_name,
            negative_prompt=options.negative_prompt,
            lab_base=LabBase(
                conds=(
                    [
                        Cond(
                            image=image_url,
                            stop_at=options.conditioning.stop_at,
                            weight=options.conditioning.weight,
                            mode=options.conditioning.mode,
                            loading=options.conditioning.loading,
                        )
                    ]
                    if image_url
                    else []
                )
            ),
            lora_models=options.lora_models,
            embeddings=options.embeddings,
            generate=Generate(
                anime_enhance=options.generate.anime_enhance,
                mode=options.generate.mode,
                gen_mode=options.generate.gen_mode,
                prompt_magic_mode=options.generate.prompt_magic_mode,
            ),
        ),
    )


def build_upscale_options_from_meta(
    *,
    meta: Meta,
    defaults: UpscaleOptions,
) -> UpscaleOptions:
    lab = meta.lab_base
    gen = meta.generate
    _d = defaults
    return UpscaleOptions(
        extra_prompt=meta.extra_prompt or _d.extra_prompt,
        local_prompt=meta.local_prompt or _d.local_prompt,
        size=SizeOptions(width=meta.width, height=meta.height),
        sampling=SamplingOptions(
            steps=int(meta.steps) if meta.steps > 0 else _d.sampling.steps,
            sampler_name=meta.sampler_name or _d.sampling.sampler_name,
            cfg_scale=meta.cfg_scale
            if meta.cfg_scale is not None
            else _d.sampling.cfg_scale,
            seed=meta.seed,
        ),
        batch_size=meta.batch_size if meta.batch_size is not None else _d.batch_size,
        n_iter=meta.n_iter if meta.n_iter is not None else _d.n_iter,
        smart_edit=meta.smart_edit,
        guidance_scale=meta.guidance_scale
        if meta.guidance_scale is not None
        else _d.guidance_scale,
        left_margin=meta.left_margin
        if meta.left_margin is not None
        else _d.left_margin,
        up_margin=meta.up_margin if meta.up_margin is not None else _d.up_margin,
        image=meta.image or _d.image,
        images=meta.images,
        vae=meta.vae,
        refiner_mode=meta.refiner_mode
        if meta.refiner_mode is not None
        else _d.refiner_mode,
        lcm_mode=meta.lcm_mode if meta.lcm_mode is not None else _d.lcm_mode,
        lora_models=meta.lora_models
        if meta.lora_models is not None
        else _d.lora_models,
        embeddings=meta.embeddings if meta.embeddings is not None else _d.embeddings,
        enable_hr=bool(meta.enable_hr) if meta.enable_hr is not None else _d.enable_hr,
        controlnet_units=meta.controlnet_units,
        lab=LabUpscaleOptions(
            refiner_switch=lab.refiner_switch
            if lab.refiner_switch is not None
            else _d.lab.refiner_switch,
            hdr_enabled=lab.hdr_enabled
            if lab.hdr_enabled is not None
            else _d.lab.hdr_enabled,
            hdr_intensity=lab.hdr_intensity
            if lab.hdr_intensity is not None
            else _d.lab.hdr_intensity,
            auto_brightness_open=lab.auto_brightness_open
            if lab.auto_brightness_open is not None
            else _d.lab.auto_brightness_open,
            auto_brightness_strength=lab.auto_brightness_strength
            if lab.auto_brightness_strength is not None
            else _d.lab.auto_brightness_strength,
            auto_adjust_open=lab.auto_adjust_open
            if lab.auto_adjust_open is not None
            else _d.lab.auto_adjust_open,
            auto_adjust_strength=lab.auto_adjust_strength
            if lab.auto_adjust_strength is not None
            else _d.lab.auto_adjust_strength,
            auto_adjust_brightness=lab.auto_adjust_brightness
            if lab.auto_adjust_brightness is not None
            else _d.lab.auto_adjust_brightness,
            film_turbo=lab.film_turbo
            if lab.film_turbo is not None
            else _d.lab.film_turbo,
            expand_token_length=lab.expand_token_length
            if lab.expand_token_length is not None
            else _d.lab.expand_token_length,
            custom_height_width=lab.custom_height_width
            if lab.custom_height_width is not None
            else _d.lab.custom_height_width,
            semantics_open=lab.semantics_open
            if lab.semantics_open is not None
            else _d.lab.semantics_open,
            sharpness=lab.sharpness if lab.sharpness is not None else _d.lab.sharpness,
            quality_open=lab.quality_open
            if lab.quality_open is not None
            else _d.lab.quality_open,
            quality_detail_offset=lab.quality_detail_offset
            if lab.quality_detail_offset is not None
            else _d.lab.quality_detail_offset,
            quality_type=lab.quality_type or _d.lab.quality_type,
            quality_mode=lab.quality_mode
            if lab.quality_mode is not None
            else _d.lab.quality_mode,
            intelligence_open=lab.intelligence_open
            if lab.intelligence_open is not None
            else _d.lab.intelligence_open,
            prompt_magic=lab.prompt_magic
            if lab.prompt_magic is not None
            else _d.lab.prompt_magic,
        )
        if lab
        else _d.lab,
        generate=GenerateOptions(
            anime_enhance=gen.anime_enhance,
            mode=gen.mode,
            gen_mode=gen.gen_mode,
            prompt_magic_mode=gen.prompt_magic_mode,
        )
        if gen
        else _d.generate,
    )


def build_remove_bg_options_from_meta(
    *,
    meta: Meta,
    defaults: RemoveBgOptions,
) -> RemoveBgOptions:
    gen = meta.generate
    _d = defaults
    return RemoveBgOptions(
        extra_prompt=meta.extra_prompt or _d.extra_prompt,
        prompt=meta.prompt or _d.prompt,
        local_prompt=meta.local_prompt or _d.local_prompt,
        width=meta.width,
        height=meta.height,
        hi_res_arg=meta.hi_res_arg if meta.hi_res_arg is not None else _d.hi_res_arg,
        clip_skip=meta.clip_skip if meta.clip_skip is not None else _d.clip_skip,
        steps=meta.steps,
        batch_size=meta.batch_size if meta.batch_size is not None else _d.batch_size,
        init_images=meta.init_images
        if meta.init_images is not None
        else _d.init_images,
        sampler_name=meta.sampler_name or _d.sampler_name,
        seed=meta.seed,
        n_iter=meta.n_iter if meta.n_iter is not None else _d.n_iter,
        cfg_scale=meta.cfg_scale if meta.cfg_scale is not None else _d.cfg_scale,
        lora_models=meta.lora_models
        if meta.lora_models is not None
        else _d.lora_models,
        smart_edit=meta.smart_edit,
        guidance_scale=meta.guidance_scale
        if meta.guidance_scale is not None
        else _d.guidance_scale,
        left_margin=meta.left_margin
        if meta.left_margin is not None
        else _d.left_margin,
        up_margin=meta.up_margin if meta.up_margin is not None else _d.up_margin,
        image=meta.image or _d.image,
        vae=meta.vae,
        refiner_mode=meta.refiner_mode
        if meta.refiner_mode is not None
        else _d.refiner_mode,
        lcm_mode=meta.lcm_mode if meta.lcm_mode is not None else _d.lcm_mode,
        embeddings=meta.embeddings if meta.embeddings is not None else _d.embeddings,
        generate=type(_d.generate)(
            anime_enhance=gen.anime_enhance,
            mode=gen.mode,
            gen_mode=gen.gen_mode,
            prompt_magic_mode=gen.prompt_magic_mode,
        )
        if gen
        else _d.generate,
    )
