import gc
import os
import torch
from pathlib import Path

import folder_paths

NODE_DIR    = Path(__file__).parent
CONFIGS_DIR = NODE_DIR / "walkyrie_configs"

_pipe_cache = {}


def _split_state_dict(merged_sd):
    components = {"transformer": {}, "text_encoder": {}, "vae": {}}
    for k, v in merged_sd.items():
        for comp in components:
            if k.startswith(comp + "."):
                components[comp][k[len(comp) + 1:]] = v
                break
    return components


def _unload_cached_pipe():
    global _pipe_cache
    if not _pipe_cache:
        return
    print("[Walkyrie] Clearing previous pipe from memory...")
    for pipe in _pipe_cache.values():
        pipe.to("cpu")
        del pipe
    _pipe_cache.clear()
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("[Walkyrie] Memory cleared.")


def _load_pipe(ckpt_path, torch_dtype):
    """
    Build a pipeline_walkyrie instance from a merged safetensors checkpoint.
    Mirrors what merge_to_single.py produces:
        transformer.*  /  text_encoder.*  /  vae.*
    """
    from safetensors.torch import load_file
    from transformers import AutoTokenizer, UMT5EncoderModel, AutoConfig
    from diffusers import (
        AutoencoderKLWan,
        WanTransformer3DModel,
        FlowMatchEulerDiscreteScheduler,
    )
    import sys, pathlib
    sys.path.insert(0, str(pathlib.Path(__file__).parent))
    from pipeline_walkyrie import pipeline_walkyrie

    cache_key = (ckpt_path, torch_dtype)
    if cache_key in _pipe_cache:
        print("[Walkyrie] Using cached pipe.")
        return _pipe_cache[cache_key]

    if _pipe_cache:
        _unload_cached_pipe()

    configs = str(CONFIGS_DIR)
    print(f"[Walkyrie] Loading checkpoint: {ckpt_path}")
    merged_sd = load_file(ckpt_path, device="cpu")
    sd        = _split_state_dict(merged_sd)

    tokenizer = AutoTokenizer.from_pretrained(configs, subfolder="tokenizer")

    te_config    = AutoConfig.from_pretrained(os.path.join(configs, "text_encoder"))
    text_encoder = UMT5EncoderModel(te_config).to(torch_dtype)
    text_encoder.load_state_dict(sd["text_encoder"], strict=False)
    text_encoder.eval()

    vae_config = AutoencoderKLWan.load_config(configs, subfolder="vae")
    vae        = AutoencoderKLWan.from_config(vae_config).to(torch_dtype)
    vae.load_state_dict(sd["vae"], strict=False)
    vae.eval()

    tr_config   = WanTransformer3DModel.load_config(configs, subfolder="transformer")
    transformer = WanTransformer3DModel.from_config(tr_config).to(torch_dtype)
    transformer.load_state_dict(sd["transformer"], strict=False)
    transformer.eval()

    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(configs, subfolder="scheduler")

    pipe = pipeline_walkyrie(
        tokenizer    = tokenizer,
        text_encoder = text_encoder,
        vae          = vae,
        transformer  = transformer,
        scheduler    = scheduler,
    )

    _pipe_cache[cache_key] = pipe
    print("[Walkyrie] Pipe ready.")
    return pipe


class WalkyrieLoader:
    """Loads a merged Walkyrie safetensors checkpoint into a pipeline."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
                "dtype"    : (["bfloat16", "float16", "float32"], {"default": "bfloat16"}),
            },
        }

    RETURN_TYPES = ("WALKYRIE_PIPE",)
    RETURN_NAMES = ("pipe",)
    FUNCTION     = "load"
    CATEGORY     = "Walkyrie"

    def load(self, ckpt_name, dtype):
        dtype_map   = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
        torch_dtype = dtype_map[dtype]
        ckpt_path   = folder_paths.get_full_path("checkpoints", ckpt_name)
        pipe        = _load_pipe(ckpt_path, torch_dtype)
        return (pipe,)


class WalkyrieLoraLoader:
    """
    Calls pipe.load_lora_weights() then pipe.fuse_lora() —
    identical to the reference usage:
        pipe.load_lora_weights("lora.safetensors")
        pipe.fuse_lora(lora_scale=1.0)
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "pipe"      : ("WALKYRIE_PIPE",),
                "lora_name" : (folder_paths.get_filename_list("loras"),),
                "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
            },
        }

    RETURN_TYPES = ("WALKYRIE_PIPE",)
    RETURN_NAMES = ("pipe",)
    FUNCTION     = "load_lora"
    CATEGORY     = "Walkyrie"

    def load_lora(self, pipe, lora_name, lora_scale):
        lora_path = folder_paths.get_full_path("loras", lora_name)
        print(f"[Walkyrie] Loading LoRA: {lora_path}  (scale={lora_scale})")
        pipe.load_lora_weights(lora_path)
        pipe.fuse_lora(lora_scale=lora_scale)
        print("[Walkyrie] LoRA fused.")
        return (pipe,)


class WalkyrieSampler:
    """Runs pipe() directly. Output is converted to ComfyUI IMAGE format."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "pipe"          : ("WALKYRIE_PIPE",),
                "prompt"        : ("STRING", {"multiline": True, "default": ""}),
                "height"        : ("INT",   {"default": 1024, "min": 64,  "max": 4096, "step": 16}),
                "width"         : ("INT",   {"default": 1024, "min": 64,  "max": 4096, "step": 16}),
                "steps"         : ("INT",   {"default": 20,   "min": 1,   "max": 200}),
                "guidance_scale": ("FLOAT", {"default": 1.0,  "min": 0.0, "max": 20.0, "step": 0.1}),
                "seed"          : ("INT",   {"default": 0,    "min": 0,   "max": 0xffffffffffffffff}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION     = "sample"
    CATEGORY     = "Walkyrie"

    def sample(self, pipe, prompt, height, width, steps, guidance_scale, seed):
        from comfy.utils import ProgressBar

        device    = "cuda" if torch.cuda.is_available() else "cpu"
        pipe.to(device)
        generator = torch.Generator(device=device).manual_seed(seed)

        comfy_pbar = ProgressBar(steps)

        def _step_callback(pipeline, i, t, callback_kwargs):
            comfy_pbar.update(1)
            return callback_kwargs

        output = pipe(
            prompt                 = prompt,
            height                 = height,
            width                  = width,
            num_inference_steps    = steps,
            guidance_scale         = guidance_scale,
            generator              = generator,
            output_type            = "np",
            callback_on_step_end   = _step_callback,
        )

        import numpy as np
        frames = output.frames                              
        image  = torch.from_numpy(np.stack(frames))        
        return (image,)



NODE_CLASS_MAPPINGS = {
    "WalkyrieLoader"    : WalkyrieLoader,
    "WalkyrieLoraLoader": WalkyrieLoraLoader,
    "WalkyrieSampler"   : WalkyrieSampler,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "WalkyrieLoader"    : "Walkyrie Loader",
    "WalkyrieLoraLoader": "Walkyrie LoRA Loader",
    "WalkyrieSampler"   : "Walkyrie Sampler",
}