from .keysync_wrapper.download_utils import download_or_check_models
from .keysync_wrapper.model_utils import load_keysync_model_config
from .keysync_wrapper.runner import run_keysync


DEFAULT_MODELS_DIR = "D:/ComfyUI/models/keysync"
DEFAULT_KEYSYNC_REPO_PATH = "D:/ComfyUI/custom_nodes/ComfyUI-Keysync-wrapper/keysync"


class KeySyncDownloadModels:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "download": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("models_dir",)
    FUNCTION = "run"
    CATEGORY = "video/lipsync"

    def run(self, download):
        resolved_models_dir = download_or_check_models(
            models_dir=DEFAULT_MODELS_DIR,
            repo_id="toninio19/keysync",
            download=download,
        )
        return (resolved_models_dir,)


class KeySyncModelLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {}}

    RETURN_TYPES = ("KEYSYNC_MODEL",)
    RETURN_NAMES = ("keysync_model",)
    FUNCTION = "load"
    CATEGORY = "video/lipsync"

    def load(self):
        model = load_keysync_model_config(
            keysync_repo_path=DEFAULT_KEYSYNC_REPO_PATH,
            models_dir=DEFAULT_MODELS_DIR,
            keyframes_ckpt="checkpoints/keyframe_dub.pt",
            interpolation_ckpt="checkpoints/interpolation_dub.pt",
        )
        return (model,)


class KeySyncLipSync:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "keysync_model": ("KEYSYNC_MODEL",),
                "images": ("IMAGE",),
                "audio": ("AUDIO",),
                "input_fps": ("INT", {"default": 25, "min": 1, "max": 120, "step": 1}),
                "compute_until": ("INT", {"default": 30, "min": 1, "max": 300, "step": 1}),
            }
        }

    RETURN_TYPES = ("IMAGE", "AUDIO", "INT", "STRING")
    RETURN_NAMES = ("images", "audio", "frame_rate", "output_video_path")
    FUNCTION = "run"
    CATEGORY = "video/lipsync"

    def run(
        self,
        keysync_model,
        images,
        audio,
        input_fps,
        compute_until,
    ):
        return run_keysync(
            keysync_model=keysync_model,
            images=images,
            audio=audio,
            input_fps=input_fps,
            compute_until=compute_until,
            fix_occlusion=False,
            position="None",
            start_frame=0,
            force_25fps_16k_audio=True,
            trim_audio_to_frames=True,
            keep_intermediate_outputs=True,
            subprocess_timeout_sec=0,
            use_subprocess=True,
        )


class KeySyncLipSyncAdvanced(KeySyncLipSync):
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "keysync_model": ("KEYSYNC_MODEL",),
                "images": ("IMAGE",),
                "audio": ("AUDIO",),
                "input_fps": ("INT", {"default": 25, "min": 1, "max": 120, "step": 1}),
                "compute_until": ("INT", {"default": 30, "min": 1, "max": 300, "step": 1}),
                "fix_occlusion": ("BOOLEAN", {"default": False}),
                "position": ("STRING", {"default": "None", "multiline": False}),
                "start_frame": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
            }
        }

    FUNCTION = "run_advanced"
    CATEGORY = "video/lipsync"

    def run_advanced(
        self,
        keysync_model,
        images,
        audio,
        input_fps,
        compute_until,
        fix_occlusion,
        position,
        start_frame,
    ):
        return run_keysync(
            keysync_model=keysync_model,
            images=images,
            audio=audio,
            input_fps=input_fps,
            compute_until=compute_until,
            fix_occlusion=fix_occlusion,
            position=position,
            start_frame=start_frame,
            force_25fps_16k_audio=True,
            trim_audio_to_frames=True,
            keep_intermediate_outputs=True,
            subprocess_timeout_sec=0,
            use_subprocess=True,
        )


NODE_CLASS_MAPPINGS = {
    "KeySyncDownloadModels": KeySyncDownloadModels,
    "KeySyncModelLoader": KeySyncModelLoader,
    "KeySyncLipSync": KeySyncLipSync,
    "KeySyncLipSyncAdvanced": KeySyncLipSyncAdvanced,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "KeySyncDownloadModels": "KeySync Download Models",
    "KeySyncModelLoader": "KeySync Model Loader",
    "KeySyncLipSync": "KeySync LipSync",
    "KeySyncLipSyncAdvanced": "KeySync LipSync Advanced",
}