from __future__ import annotations

import gc
import sys
from copy import copy
from dataclasses import dataclass
from pathlib import Path
from types import SimpleNamespace
from typing import Any

import torch


NODE_ROOT = Path(__file__).resolve().parent
SRC_ROOT = NODE_ROOT / "src"
if not (SRC_ROOT / "rum_xpred").is_dir():
    raise ImportError(f"Bundled RUM Anima X-Pred package not found: {SRC_ROOT / 'rum_xpred'}")
if str(SRC_ROOT) not in sys.path:
    sys.path.insert(0, str(SRC_ROOT))

from rum_xpred.adapters.anima_sd_scripts import create_adapter
from rum_xpred.anima import anima_euler_step, make_shifted_sigma_schedule, xpred_to_anima_v
from rum_xpred.cache_batches import make_seeded_eps_batch


def _model_names(folder_name: str, placeholder: str) -> list[str]:
    try:
        import folder_paths

        names = folder_paths.get_filename_list(folder_name)
        return names if names else [placeholder]
    except Exception:
        return [placeholder]


def _model_path(folder_name: str, filename: str) -> Path:
    try:
        import folder_paths

        return Path(folder_paths.get_full_path_or_raise(folder_name, filename))
    except Exception as exc:
        raise FileNotFoundError(f"Could not resolve {filename!r} from ComfyUI models/{folder_name}") from exc


def _sigma_like(value: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
    return value.reshape(1, 1, 1, 1).to(device=z.device, dtype=z.dtype).expand(z.shape[0], 1, 1, 1)


def _velocity_from_xpred(student_forward, z: torch.Tensor, sigma: torch.Tensor, eps_floor: float) -> torch.Tensor:
    x_pred = student_forward(z, sigma)
    return xpred_to_anima_v(z, x_pred, sigma, eps_floor)


def _sample_euler(
    student_forward,
    eps_latent: torch.Tensor,
    sigmas: torch.Tensor,
    *,
    prediction_type: str,
    eps_floor: float,
) -> torch.Tensor:
    z = eps_latent
    for index, sigma_value in enumerate(sigmas[:-1]):
        sigma = _sigma_like(sigma_value, z)
        sigma_next = sigmas[index + 1].reshape(1, 1, 1, 1).to(device=z.device, dtype=z.dtype)
        if prediction_type == "x":
            v = _velocity_from_xpred(student_forward, z, sigma, eps_floor)
        else:
            v = student_forward(z, sigma)
        z = anima_euler_step(z, v, sigma, sigma_next)
        if not torch.isfinite(z).all():
            raise FloatingPointError(f"non-finite latent during {prediction_type}-pred Euler sampling")
    return z


def _sample_heun(
    student_forward,
    eps_latent: torch.Tensor,
    sigmas: torch.Tensor,
    *,
    prediction_type: str,
    eps_floor: float,
) -> torch.Tensor:
    z = eps_latent
    if sigmas.numel() <= 2:
        return _sample_euler(student_forward, z, sigmas, prediction_type=prediction_type, eps_floor=eps_floor)

    for index in range(sigmas.numel() - 2):
        sigma = _sigma_like(sigmas[index], z)
        sigma_next = _sigma_like(sigmas[index + 1], z)
        if prediction_type == "x":
            v = _velocity_from_xpred(student_forward, z, sigma, eps_floor)
        else:
            v = student_forward(z, sigma)
        z_euler = anima_euler_step(z, v, sigma, sigma_next)
        if prediction_type == "x":
            v_next = _velocity_from_xpred(student_forward, z_euler, sigma_next, eps_floor)
        else:
            v_next = student_forward(z_euler, sigma_next)
        z = anima_euler_step(z, 0.5 * (v + v_next), sigma, sigma_next)
        if not torch.isfinite(z).all():
            raise FloatingPointError(f"non-finite latent during {prediction_type}-pred Heun sampling")

    sigma = _sigma_like(sigmas[-2], z)
    sigma_next = sigmas[-1].reshape(1, 1, 1, 1).to(device=z.device, dtype=z.dtype)
    if prediction_type == "x":
        v = _velocity_from_xpred(student_forward, z, sigma, eps_floor)
    else:
        v = student_forward(z, sigma)
    return anima_euler_step(z, v, sigma, sigma_next)


def _comfy_device() -> torch.device:
    try:
        import comfy.model_management as model_management

        return torch.device(model_management.get_torch_device())
    except Exception:
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _comfy_intermediate_device() -> torch.device:
    try:
        import comfy.model_management as model_management

        return torch.device(model_management.intermediate_device())
    except Exception:
        return torch.device("cpu")


def _soft_empty_cache() -> None:
    gc.collect()
    try:
        import comfy.model_management as model_management

        model_management.soft_empty_cache()
    except Exception:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def _dtype_from_name(name: str) -> torch.dtype:
    if name == "bf16":
        return torch.bfloat16
    if name == "fp16":
        return torch.float16
    if name == "fp32":
        return torch.float32
    raise ValueError(f"unsupported precision: {name!r}")


@dataclass
class LoadedAnimaXPred:
    args: Any
    adapter: Any
    student: torch.nn.Module
    device: torch.device
    dtype: torch.dtype
    prediction_type: str


class AnimaXPredModelLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "checkpoint": (_model_names("diffusion_models", "put_xpred_checkpoint_in_models_diffusion_models.safetensors"),),
                "text_encoder": (_model_names("text_encoders", "put_qwen_text_encoder_in_models_text_encoders.safetensors"),),
                "vae": (_model_names("vae", "put_qwen_image_vae_in_models_vae.safetensors"),),
                "prediction_type": (["x", "v"], {"default": "x"}),
                "precision": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
                "attn_mode": (["torch", "flash", "sageattn", "xformers"], {"default": "flash"}),
                "text_encoder_cpu": ("BOOLEAN", {"default": False}),
                "fp8": ("BOOLEAN", {"default": False}),
                "fp8_scaled": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("RUM_ANIMA_XPRED",)
    RETURN_NAMES = ("model",)
    FUNCTION = "load_model"
    CATEGORY = "RUM/Anima XPred"

    def load_model(
        self,
        checkpoint: str,
        text_encoder: str,
        vae: str,
        prediction_type: str,
        precision: str,
        attn_mode: str,
        text_encoder_cpu: bool,
        fp8: bool,
        fp8_scaled: bool,
    ):
        checkpoint_path = _model_path("diffusion_models", checkpoint)
        text_encoder_path = _model_path("text_encoders", text_encoder)
        vae_path = _model_path("vae", vae)
        for label, path in {
            "checkpoint": checkpoint_path,
            "text_encoder": text_encoder_path,
            "vae": vae_path,
        }.items():
            if not path.is_file():
                raise FileNotFoundError(f"{label} not found: {path}")

        device = _comfy_device()
        dtype = _dtype_from_name(precision)
        args = SimpleNamespace(
            dit=str(checkpoint_path),
            student_init=str(checkpoint_path),
            text_encoder=str(text_encoder_path),
            vae=str(vae_path),
            output_dir=str(NODE_ROOT / "outputs/comfyui-rum-anima-xpred"),
            negative_prompt="",
            flow_shift=3.0,
            teacher_steps=40,
            attn_mode=attn_mode,
            fp8=fp8,
            fp8_scaled=fp8_scaled,
            text_encoder_cpu=text_encoder_cpu,
            teacher_lora=None,
            teacher_lora_weight=1.0,
            vae_spatial_chunk_size=None,
            vae_disable_cache=False,
        )
        adapter = create_adapter(args, device=device, dtype=dtype)
        student = adapter.load_student_xpred(init_checkpoint=str(checkpoint_path))
        student.to(device=device, dtype=dtype).eval().requires_grad_(False)
        return (LoadedAnimaXPred(args=args, adapter=adapter, student=student, device=device, dtype=dtype, prediction_type=prediction_type),)


class AnimaXPredLoraLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("RUM_ANIMA_XPRED",),
                "lora": (_model_names("loras", "put_anima_lora_in_models_loras.safetensors"),),
                "strength": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.05}),
            }
        }

    RETURN_TYPES = ("RUM_ANIMA_XPRED",)
    RETURN_NAMES = ("model",)
    FUNCTION = "load_lora"
    CATEGORY = "RUM/Anima XPred"

    def load_lora(self, model: LoadedAnimaXPred, lora: str, strength: float):
        lora_path = _model_path("loras", lora)
        if not lora_path.is_file():
            raise FileNotFoundError(f"lora not found: {lora_path}")

        args = copy(model.args)
        args.teacher_lora = str(lora_path)
        args.teacher_lora_weight = float(strength)

        adapter = create_adapter(args, device=model.device, dtype=model.dtype)
        checkpoint_path = getattr(args, "student_init", None) or getattr(args, "dit", None)
        student = adapter.load_student_xpred(init_checkpoint=checkpoint_path)
        student.to(device=model.device, dtype=model.dtype).eval().requires_grad_(False)
        return (
            LoadedAnimaXPred(
                args=args,
                adapter=adapter,
                student=student,
                device=model.device,
                dtype=model.dtype,
                prediction_type=model.prediction_type,
            ),
        )


class AnimaXPredSampler:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("RUM_ANIMA_XPRED",),
                "prompt": ("STRING", {"multiline": True, "default": "hatsune miku, 1girl, masterpiece, best quality"}),
                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
                "seed": ("INT", {"default": 20260701, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
                "steps": ("INT", {"default": 10, "min": 1, "max": 200}),
                "sampler": (["heun", "euler"], {"default": "heun"}),
                "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.1}),
                "width": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 8}),
                "height": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 8}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 16}),
                "flow_shift": ("FLOAT", {"default": 3.0, "min": 0.01, "max": 100.0, "step": 0.01}),
                "eps_floor": ("FLOAT", {"default": 1e-4, "min": 1e-8, "max": 1.0, "step": 1e-5}),
                "offload_after_sample": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("IMAGE", "RUM_ANIMA_LATENT")
    RETURN_NAMES = ("image", "latent")
    FUNCTION = "sample"
    CATEGORY = "RUM/Anima XPred"

    def sample(
        self,
        model: LoadedAnimaXPred,
        prompt: str,
        negative_prompt: str,
        seed: int,
        steps: int,
        sampler: str,
        cfg: float,
        width: int,
        height: int,
        batch_size: int,
        flow_shift: float,
        eps_floor: float,
        offload_after_sample: bool,
    ):
        if width % 8 != 0 or height % 8 != 0:
            raise ValueError("width and height must be divisible by 8")

        model.args.prompt = prompt
        model.args.negative_prompt = negative_prompt
        model.args.width = width
        model.args.height = height
        model.args.flow_shift = flow_shift
        model.student.to(device=model.device, dtype=model.dtype).eval()

        with torch.no_grad():
            sigmas = make_shifted_sigma_schedule(steps, flow_shift, device=model.device, dtype=model.dtype)
            eps_latent = make_seeded_eps_batch(
                list(range(batch_size)),
                seed=int(seed),
                height=height,
                width=width,
                device=model.device,
                dtype=model.dtype,
            )
            cond_embed, uncond_embed = model.adapter._encode_prompts([prompt] * batch_size, cfg=cfg, anima_model=model.student)
            text_conditioning = {
                "prompt_embeds": cond_embed.detach(),
                "negative_prompt_embeds": uncond_embed.detach(),
            }

            def student_forward(z: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
                return model.adapter.student_forward_xpred(
                    model.student,
                    z,
                    sigma,
                    text_conditioning,
                    guidance_scale=cfg,
                )

            if sampler == "euler":
                latent = _sample_euler(
                    student_forward,
                    eps_latent,
                    sigmas,
                    prediction_type=model.prediction_type,
                    eps_floor=eps_floor,
                )
            else:
                latent = _sample_heun(
                    student_forward,
                    eps_latent,
                    sigmas,
                    prediction_type=model.prediction_type,
                    eps_floor=eps_floor,
                )

            vae = model.adapter._load_vae()
            vae.to(model.device)
            pixels = vae.decode_to_pixels(latent.to(model.device, dtype=vae.dtype))
            if pixels.ndim == 5:
                pixels = pixels.squeeze(2)
            images = ((pixels.clamp(-1.0, 1.0) + 1.0) * 0.5).to(torch.float32)
            images = images.movedim(1, -1).to(device=_comfy_intermediate_device())
            latent_out = {
                "samples": latent.detach().to(device=_comfy_intermediate_device(), dtype=torch.float32),
                "sigmas": sigmas.detach().to(device="cpu", dtype=torch.float32),
            }

        if offload_after_sample:
            model.student.to("cpu")
            if model.adapter.vae is not None:
                model.adapter.vae.to("cpu")
            _soft_empty_cache()

        return (images, latent_out)


NODE_CLASS_MAPPINGS = {
    "AnimaXPredModelLoader": AnimaXPredModelLoader,
    "AnimaXPredLoraLoader": AnimaXPredLoraLoader,
    "AnimaXPredSampler": AnimaXPredSampler,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "AnimaXPredModelLoader": "Load Anima XPred Model",
    "AnimaXPredLoraLoader": "Load Anima XPred LoRA",
    "AnimaXPredSampler": "Sample Anima XPred",
}