import os
import torch
from shared.utils.hf import build_hf_url


class family_handler:
    @staticmethod
    def query_supported_types():
        return ["longcat_video", "longcat_avatar"]

    @staticmethod
    def query_family_maps():
        return {}, {}

    @staticmethod
    def query_model_family():
        return "longcat"

    @staticmethod
    def query_family_infos():
        return {"longcat": (60, "LongCat")}

    @staticmethod
    def register_lora_cli_args(parser):
        parser.add_argument(
            "--lora-dir-longcat",
            type=str,
            default=os.path.join("loras", "longcat"),
            help="Path to a directory that contains LongCat Video LoRAs",
        )
        parser.add_argument(
            "--lora-dir-longcat-avatar",
            type=str,
            default=os.path.join("loras", "longcat_avatar"),
            help="Path to a directory that contains LongCat Avatar LoRAs",
        )

    @staticmethod
    def get_lora_dir(base_model_type, args):
        if base_model_type == "longcat_avatar":
            return args.lora_dir_longcat_avatar
        return args.lora_dir_longcat

    @staticmethod
    def query_model_def(base_model_type, model_def):
        extra_model_def = {
            "frames_minimum": 5,
            "frames_steps": 4,
            "sliding_window": True,
            "guidance_max_phases": 1,
            "image_prompt_types_allowed": "TSVL",
            "video_continuation": True,
            "sample_solvers": [
                ("Auto (Continuation = Enhanced HF)", "auto"),
                ("Default", ""),
                ("Enhanced HF", "enhance_hf"),
                ("Distill", "distill"),
            ],
        }
        text_encoder_folder = "umt5-xxl"
        extra_model_def["text_encoder_URLs"] = [
            build_hf_url("DeepBeepMeep/Wan2.1", text_encoder_folder, "models_t5_umt5-xxl-enc-bf16.safetensors"),
            build_hf_url("DeepBeepMeep/Wan2.1", text_encoder_folder, "models_t5_umt5-xxl-enc-quanto_int8.safetensors"),
        ]
        extra_model_def["text_encoder_folder"] = text_encoder_folder

        if base_model_type == "longcat_video":
            extra_model_def.update(
                {
                    "fps": 15,
                    "profiles_dir": ["longcat_video"],
                }
            )
        elif base_model_type == "longcat_avatar":
            extra_model_def.update(
                {
                    "fps": 16,
                    "profiles_dir": [base_model_type],
                    "audio_guide_label": "Voice to follow",
                    "audio_guide2_label": "Voice to follow #2",
                    "audio_guidance": True,
                    "any_audio_prompt": True,
                    "audio_prompt_choices": True,                
                    "image_ref_choices": {
                        "choices": [("None", ""), ("Anchor Reference Image", "KI")],
                        "letters_filter": "KI",
                        "visible": True,
                        "label": "Anchor Reference Image",
                    },
                    "reference_image_enabled": True,
                    "no_background_removal": True,
                    "image_prompt_types_allowed": "TSVL",
                }
            )


        return extra_model_def

    @staticmethod
    def get_rgb_factors(base_model_type):
        from shared.RGB_factors import get_rgb_factors

        return get_rgb_factors("wan")

    @staticmethod
    def query_model_files(computeList, base_model_type, model_def=None):
        download_def = [
            {
                "repoId": "DeepBeepMeep/Wan2.1",
                "sourceFolderList": ["umt5-xxl", "chinese-wav2vec2-base"],
                "fileList": [
                    ["special_tokens_map.json", "spiece.model", "tokenizer.json", "tokenizer_config.json"],
                    [
                        "config.json",
                        "preprocessor_config.json",
                        "pytorch_model.bin",
                        "readme.txt",
                    ],
                ],
            }
        ]
        download_def += [
            {
                "repoId": "DeepBeepMeep/Wan2.1",
                "sourceFolderList": [""],
                "fileList": [["Wan2.1_VAE_bf16.safetensors"]],
            }
        ]
        return download_def

    @staticmethod
    def load_model(
        model_filename,
        model_type,
        base_model_type,
        model_def,
        quantizeTransformer=False,
        text_encoder_quantization=None,
        dtype=torch.bfloat16,
        VAE_dtype=torch.float32,
        mixed_precision_transformer=False,
        save_quantized=False,
        submodel_no_list=None,
        text_encoder_filename=None,
        **kwargs,
    ):
        from .longcat_main import LongCatModel

        longcat_model = LongCatModel(
            checkpoint_dir="ckpts",
            model_filename=model_filename,
            model_type=model_type,
            model_def=model_def,
            base_model_type=base_model_type,
            text_encoder_filename=text_encoder_filename,
            quantizeTransformer=quantizeTransformer,
            dtype=dtype,
            VAE_dtype=VAE_dtype,
            mixed_precision_transformer=mixed_precision_transformer,
            save_quantized=save_quantized,
        )

        pipe = {
            "transformer": longcat_model.transformer,
            "vae": longcat_model.vae,
            "text_encoder": longcat_model.text_encoder.model,
        }
        if longcat_model.audio_encoder is not None:
            pipe["wav2vec"] = longcat_model.audio_encoder

        return longcat_model, pipe

    @staticmethod
    def update_default_settings(base_model_type, model_def, ui_defaults):
        ui_defaults.update(
            {
                "guidance_scale": 4.0,
                "num_inference_steps": 50,
                "audio_guidance_scale": 4.0,
                "sliding_window_overlap": 13,
                "sliding_window_size": 93,                 
            }
        )
        if base_model_type == "longcat_video":
            ui_defaults.update({"video_length": 93})

        if base_model_type in ["longcat_avatar"]:
            ui_defaults.update({"video_length": 93, "video_prompt_type": ""})

        if ui_defaults.get("sample_solver", "") == "":
            ui_defaults["sample_solver"] = "auto"