import torch
import os
import kornia
import pathlib
import typing
import cv2
import numpy as np
import tempfile
import shutil
import warnings
import logging
from functools import cmp_to_key
from packaging import version
from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames, generic_frame_loop, InterpolationStateList
from comfy.model_management import get_torch_device
from insightface.app import FaceAnalysis
from insightface.model_zoo import get_model
from urllib.request import urlretrieve
import sys

# -------------------------------
# Logging & warning setup
# -------------------------------
DEBUG = False
os.environ["ORT_LOG_LEVEL"] = "3"  # ONNX only shows errors
os.environ["ORT_CUDA_MEMORY_LIMIT"] = "0"
warnings.simplefilter("ignore", category=FutureWarning)
warnings.simplefilter("ignore", category=DeprecationWarning)
warnings.simplefilter("ignore", category=UserWarning)
logging.getLogger("insightface").setLevel(logging.ERROR)
logging.getLogger("urllib3").setLevel(logging.ERROR)
logging.getLogger("onnxruntime").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")  # suppress all remaining warnings
logging.basicConfig(level=logging.ERROR)

# -------------------------------
# Context manager to suppress stdout/stderr temporarily
# -------------------------------
class suppress_stdout_stderr:
    """Suppress all console output temporarily"""
    def __enter__(self):
        self._stdout = sys.stdout
        self._stderr = sys.stderr
        sys.stdout = open(os.devnull, "w")
        sys.stderr = open(os.devnull, "w")
    def __exit__(self, exc_type, exc_val, exc_tb):
        sys.stdout.close()
        sys.stderr.close()
        sys.stdout = self._stdout
        sys.stderr = self._stderr

def log(msg, progress=None):
    """Unified log function"""
    if DEBUG:
        print(msg)
    if progress is not None:
        # Optionally integrate with ComfyUI progress here
        pass

# -------------------------------
# RIFE VFI Node
# -------------------------------
MODEL_TYPE = pathlib.Path(__file__).parent.name
CKPT_NAME_VER_DICT = {
    "rife40.pth": "4.0",
    "rife41.pth": "4.0",
    "rife42.pth": "4.2",
    "rife43.pth": "4.3",
    "rife44.pth": "4.3",
    "rife45.pth": "4.5",
    "rife46.pth": "4.6",
    "rife47.pth": "4.7",
    "rife48.pth": "4.7",
    "rife49.pth": "4.7"
}

class RIFE_VFI:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "ckpt_name": (
                    sorted(list(CKPT_NAME_VER_DICT.keys()), key=lambda ckpt_name: version.parse(CKPT_NAME_VER_DICT[ckpt_name])),
                    {"default": "rife47.pth"}
                ),
                "frames": ("IMAGE", ),
                "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}),
                "multiplier": ("INT", {"default": 2, "min": 1}),
                "fast_mode": ("BOOLEAN", {"default":True}),
                "ensemble": ("BOOLEAN", {"default":True}),
                "scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0})
            },
            "optional": {
                "optional_interpolation_states": ("INTERPOLATION_STATES", )
            }
        }

    RETURN_TYPES = ("IMAGE", )
    FUNCTION = "vfi"
    CATEGORY = "bmgjet"

    def vfi(
        self,
        ckpt_name: typing.AnyStr,
        frames: torch.Tensor,
        clear_cache_after_n_frames = 10,
        multiplier: typing.SupportsInt = 2,
        fast_mode = False,
        ensemble = False,
        scale_factor = 1.0,
        optional_interpolation_states: InterpolationStateList = None,
        **kwargs
    ):
        from .rife_arch import IFNet
        model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name)
        arch_ver = CKPT_NAME_VER_DICT[ckpt_name]
        interpolation_model = IFNet(arch_ver=arch_ver)

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        state = torch.load(model_path, map_location=device)
        interpolation_model.load_state_dict(state)
        interpolation_model = interpolation_model.to(device).eval()
        frames = preprocess_frames(frames).to(device)
        
        def return_middle_frame(frame_0, frame_1, timestep, model, scale_list, in_fast_mode, in_ensemble):
            return model(frame_0, frame_1, timestep, scale_list, in_fast_mode, in_ensemble)
        
        scale_list = [8 / scale_factor, 4 / scale_factor, 2 / scale_factor, 1 / scale_factor] 
        
        args = [interpolation_model, scale_list, fast_mode, ensemble]
        out = postprocess_frames(
            generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, 
                               interpolation_states=optional_interpolation_states, dtype=torch.float32)
        )
        return (out,)

# -------------------------------
# InputSwitch Node
# -------------------------------
class InputSwitch:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mode": (["Image", "Video"],),
                "max_width": ("INT", {"default": 1024, "min": 64, "max": 2048, "step": 2}),
                "max_height": ("INT", {"default": 1024, "min": 64, "max": 2048, "step": 2}),
            },
            "optional": {
                "Image": ("IMAGE",),
                "Video": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("IMAGE", "INT", "INT", "INT", "IMAGE")
    RETURN_NAMES = ("image", "width", "height", "frame_count", "images")
    FUNCTION = "load_media"
    CATEGORY = "bmgjet"

    def load_media(self, mode, max_width, max_height, Image=None, Video=None, on_progress=None):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        def log(msg, progress=None):
            if on_progress:
                on_progress(msg, progress)
            else:
                print("[InputSwitch]", msg)
        def round2(x):
            return x - (x % 2)

        max_width = round2(max_width)
        max_height = round2(max_height)

        # Helper resize function
        def resize_if_needed(tensor):
            if tensor is None:
                raise ValueError("[InputSwitch] Received None tensor for resizing")
            is_video = (tensor.ndim == 4)
            if is_video:
                N, H, W, C = tensor.shape
            else:
                H, W, C = tensor.shape
                N = 1
            if W <= max_width and H <= max_height:
                return tensor, H, W
            scale = min(max_width / W, max_height / H)
            new_w = round2(int(W * scale))
            new_h = round2(int(H * scale))
            if is_video:
                t = tensor.permute(0, 3, 1, 2)
            else:
                t = tensor.unsqueeze(0).permute(0, 3, 1, 2)
            t_resized = kornia.geometry.transform.resize(
                t, (new_h, new_w), interpolation="bicubic", antialias=True
            )
            if is_video:
                out = t_resized.permute(0, 2, 3, 1)
            else:
                out = t_resized.permute(0, 2, 3, 1).squeeze(0)
            return out, new_h, new_w

        if mode == "Image":
            if Image is None:
                raise ValueError("No image provided for 'Image' mode.")
            Image = Image.to(device)
            log("Resizing image...", 0.0)
            Image, h_new, w_new = resize_if_needed(Image)
            Image = Image.unsqueeze(0) if Image.ndim == 3 else Image
            Image_dup = torch.cat([Image, Image.clone()], dim=0)
            log(f"Width={w_new}, Height={h_new}, Frames=1", 1.0)
            return (Image, w_new, h_new, 1, Image_dup)

        if mode == "Video":
            if Video is None:
                raise ValueError("No video input provided for 'Video' mode.")
            if isinstance(Video, torch.Tensor):
                if Video.ndim == 3:
                    Video = Video.unsqueeze(0)
                Video = Video.to(device)
            elif isinstance(Video, (list, tuple)):
                Video = torch.stack([v.to(device) if isinstance(v, torch.Tensor) else torch.tensor(v, device=device) for v in Video])
            else:
                raise TypeError("Unsupported video input type.")
            resized_frames = []
            total_frames = Video.shape[0]
            for idx, f in enumerate(Video):
                f_resized, h_new, w_new = resize_if_needed(f)
                resized_frames.append(f_resized)
                if on_progress:
                    on_progress(f"Resizing frame {idx+1}/{total_frames}", (idx+1)/total_frames)
            Video = torch.stack(resized_frames)
            frame_count = Video.shape[0]
            last_frame = Video[-1:].clone()
            video_without_last = Video[:-1].clone()
            log(f"Width={w_new}, Height={h_new}, Frames={frame_count}", 1.0)
            if last_frame.ndim == 3:
                last_frame = last_frame.unsqueeze(0)
            if video_without_last.ndim == 3:
                video_without_last = video_without_last.unsqueeze(0)
            return (last_frame, w_new, h_new, frame_count, video_without_last)

class MoveToDevice:
    @classmethod
    def INPUT_TYPES(cls):
        # Accepts ANY ComfyUI input type
        return {
            "required": {
                "data": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "move"
    CATEGORY = "bmgjet"

    def move(self, data):
        device = get_torch_device()

        def move_any(x):
            if isinstance(x, torch.Tensor):
                return x.to(device)
            if isinstance(x, list):
                return [move_any(v) for v in x]
            if isinstance(x, tuple):
                return tuple(move_any(v) for v in x)
            if isinstance(x, dict):
                return {k: move_any(v) for k, v in x.items()}
            return x

        return (move_any(data),)

class RoundToEven:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": ("INT", {"default": 1, "min": 2, "max": 999999}),
            }
        }

    RETURN_TYPES = ("INT",)
    FUNCTION = "round_even"
    CATEGORY = "bmgjet"

    def round_even(self, value):
        # Round to nearest number divisible by 2
        # (i.e., nearest even number)
        rounded = int(2 * round(value / 2))
        return (rounded,)

class ColorMatch:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "reference_image": ("IMAGE",),
                "match_strength": (
                    "FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}
                )
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "match_color"
    CATEGORY = "bmgjet"

    def match_color(self, images, reference_image, match_strength, on_progress=None):
        import torch
        import kornia
        import os
        import sys
        import time

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        images = images.to(device)
        reference_image = reference_image.to(device)
        images = images.permute(0, 3, 1, 2)
        reference_image = reference_image.permute(0, 3, 1, 2)

        ref_lab = kornia.color.rgb_to_lab(reference_image)
        ref_mean = ref_lab.mean(dim=[2, 3], keepdim=True)
        ref_std = ref_lab.std(dim=[2, 3], keepdim=True) + 1e-5

        num_threads = os.cpu_count() or 4
        batch_size = min(128, max(4, num_threads))
        outputs = []

        total_batches = (images.shape[0] + batch_size - 1) // batch_size
        bar_length = 30
        start_time = time.time()

        for batch_idx, start in enumerate(range(0, images.shape[0], batch_size)):
            batch = images[start:start + batch_size]
            img_lab = kornia.color.rgb_to_lab(batch)
            img_mean = img_lab.mean(dim=[2, 3], keepdim=True)
            img_std = img_lab.std(dim=[2, 3], keepdim=True) + 1e-5
            matched_lab = (img_lab - img_mean) / img_std * ref_std + ref_mean
            blended_lab = match_strength * matched_lab + (1.0 - match_strength) * img_lab
            output = kornia.color.lab_to_rgb(blended_lab)
            outputs.append(output)

            # -------------------------------
            # Console progress bar (single-line)
            # -------------------------------
            progress = (batch_idx + 1) / total_batches
            filled = int(bar_length * progress)
            bar = "█" * filled + "-" * (bar_length - filled)
            elapsed = time.time() - start_time
            avg_time = elapsed / (batch_idx + 1)
            remaining = avg_time * (total_batches - batch_idx - 1)
            msg = f"[ColorMatch] Progress: |{bar}| {batch_idx + 1}/{total_batches} ETA: {remaining:.1f}s"

            if on_progress:
                on_progress(msg, progress)
            else:
                print(f"\r{msg}", end="")
                sys.stdout.flush()

        output = torch.cat(outputs, dim=0).clamp(0.0, 1.0)
        output = output.permute(0, 2, 3, 1)
        output = output.to(device)

        if on_progress is None:
            print()  # move to next line after completion
            print("[ColorMatch] Color matching complete.")
        else:
            on_progress("[ColorMatch] Color matching complete.", 1.0)

        return (output,)


class FaceSwapNode:
    face_app = None
    swapper_model = None
    cached_src_face = None
    cached_src_image = None  # store last source image for comparison

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_image": ("IMAGE",),
                "target_images": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "swap_faces"
    CATEGORY = "bmgjet"

    def swap_faces(self, source_image, target_images):
        import os
        import numpy as np
        import torch
        import time
        import sys
        from urllib.request import urlretrieve
        from insightface.app import FaceAnalysis
        from insightface.model_zoo import get_model

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # -------------------------------
        # Helpers
        # -------------------------------
        def ensure_rgb(img):
            if img is None:
                raise ValueError("[FaceSwapNode] Received None image input")
            if isinstance(img, torch.Tensor):
                img = img.detach().cpu().numpy()
                if img.ndim == 4:
                    img = img[0]
                if img.ndim == 3 and img.shape[0] in (1,3):
                    img = np.transpose(img, (1,2,0))
                img = (img*255).astype(np.uint8)
            if img.ndim == 2:
                img = np.stack([img]*3, axis=-1)
            elif img.ndim == 3 and img.shape[2] == 1:
                img = np.repeat(img, 3, axis=2)
            return np.ascontiguousarray(img)

        def to_tensor(img: np.ndarray):
            img = ensure_rgb(img)
            img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
            img = cv2.flip(img, 1)  # 1 = horizontal flip
            img = np.ascontiguousarray(img)
            t = torch.from_numpy(img).permute(1,0,2).float() / 255.0
            return t.to(device)

        def images_equal(img1, img2):
            if img1 is None or img2 is None:
                return False
            return np.array_equal(img1, img2)

        # -------------------------------
        # Load inswapper and FaceAnalysis models silently
        # -------------------------------
        INSWAPPER_URL = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
        base_folder = os.path.join(os.path.dirname(__file__), "Models")
        os.makedirs(base_folder, exist_ok=True)
        model_path = os.path.join(base_folder, "inswapper_128.onnx")
        if not os.path.exists(model_path):
            print("[FaceSwapNode] Downloading inswapper_128.onnx...")
            urlretrieve(INSWAPPER_URL, model_path)

        with suppress_stdout_stderr():
            if FaceSwapNode.swapper_model is None:
                FaceSwapNode.swapper_model = get_model(model_path, download=False)

            if FaceSwapNode.face_app is None:
                ctx = 0 if torch.cuda.is_available() else -1
                FaceSwapNode.face_app = FaceAnalysis(name="buffalo_l")
                FaceSwapNode.face_app.prepare(ctx_id=ctx, det_size=(640,640))

        app = FaceSwapNode.face_app
        swapper = FaceSwapNode.swapper_model

        # -------------------------------
        # Prepare source face
        # -------------------------------
        src_img = ensure_rgb(source_image)
        if FaceSwapNode.cached_src_face is None or not images_equal(FaceSwapNode.cached_src_image, src_img):
            with suppress_stdout_stderr():
                faces = app.get(src_img) or []
            FaceSwapNode.cached_src_face = faces[0] if faces else None
            FaceSwapNode.cached_src_image = src_img

        src_face = FaceSwapNode.cached_src_face

        # -------------------------------
        # Prepare target images
        # -------------------------------
        tgt_list = []
        if isinstance(target_images, torch.Tensor):
            if target_images.ndim == 3:
                tgt_list = [target_images]
            elif target_images.ndim == 4:
                tgt_list = [target_images[i] for i in range(target_images.shape[0])]
        elif isinstance(target_images, (list, tuple)):
            tgt_list = [t for t in target_images if t is not None]
        else:
            tgt_list = [target_images]

        tgt_list = [ensure_rgb(t) for t in tgt_list if t is not None]
        if len(tgt_list) == 0:
            raise ValueError("[FaceSwapNode] No valid target images to process.")

        # -------------------------------
        # Face swapping loop
        # -------------------------------
        swapped_list = []
        total = len(tgt_list)
        bar_length = 30
        start_time = time.time()

        for idx, tgt in enumerate(tgt_list, 1):
            with suppress_stdout_stderr():
                if src_face is None:
                    swapped_list.append(tgt)
                else:
                    tgt_faces = app.get(tgt) or []
                    if not tgt_faces:
                        swapped_list.append(tgt)
                    else:
                        swapped_list.append(swapper.get(tgt, tgt_faces[0], src_face, paste_back=True))

            # Progress bar (single line)
            elapsed = time.time() - start_time
            avg_time = elapsed / idx
            remaining = avg_time * (total - idx)
            filled = int(bar_length * idx // total)
            bar = "█" * filled + "-" * (bar_length - filled)
            print(f"\r[FaceSwapNode] Progress: |{bar}| {idx}/{total} ETA: {remaining:.1f}s", end="")
            sys.stdout.flush()

        print("\n[FaceSwapNode] Face swapping complete.")

        # Convert swapped images to tensor
        swapped_tensors = [to_tensor(img) for img in swapped_list]
        video_tensor = torch.stack(swapped_tensors, dim=0)
        return (video_tensor,)

# -------------------------------
# Basic Math Nodes
# -------------------------------

class MathAdd:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "a": ("FLOAT", {"default": 0.0}),
                "b": ("FLOAT", {"default": 0.0}),
            }
        }

    RETURN_TYPES = ("FLOAT",)
    FUNCTION = "add"
    CATEGORY = "bmgjet/math"

    def add(self, a, b):
        return (float(a + b),)


class MathSubtract:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "a": ("FLOAT", {"default": 0.0}),
                "b": ("FLOAT", {"default": 0.0}),
            }
        }

    RETURN_TYPES = ("FLOAT",)
    FUNCTION = "subtract"
    CATEGORY = "bmgjet/math"

    def subtract(self, a, b):
        return (float(a - b),)


class MathMultiply:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "a": ("FLOAT", {"default": 0.0}),
                "b": ("FLOAT", {"default": 0.0}),
            }
        }

    RETURN_TYPES = ("FLOAT",)
    FUNCTION = "multiply"
    CATEGORY = "bmgjet/math"

    def multiply(self, a, b):
        return (float(a * b),)


class MathDivide:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "a": ("FLOAT", {"default": 0.0}),
                "b": ("FLOAT", {"default": 1.0}),
            }
        }

    RETURN_TYPES = ("FLOAT",)
    FUNCTION = "divide"
    CATEGORY = "bmgjet/math"

    def divide(self, a, b):
        if b == 0:
            return (0.0,)
        return (float(a / b),)

# -------------------------------
# Float → Int Node
# -------------------------------

class FloatToInt:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": ("FLOAT", {"default": 0.0}),
            }
        }

    RETURN_TYPES = ("INT",)
    FUNCTION = "convert"
    CATEGORY = "bmgjet/math"

    def convert(self, value):
        return (int(value),)

# Node registration
NODE_CLASS_MAPPINGS = {"InputSwitch": InputSwitch, "ColorMatch": ColorMatch, "RIFE_VFI": RIFE_VFI, "RoundToEven": RoundToEven,"MoveToDevice": MoveToDevice, "FaceSwapNode": FaceSwapNode, "MathAdd": MathAdd, "MathSubtract": MathSubtract, "MathMultiply": MathMultiply, "MathDivide": MathDivide, "FloatToInt": FloatToInt,}
NODE_DISPLAY_NAME_MAPPINGS = {"InputSwitch": "Input Switch", "ColorMatch": "Color Match", "RIFE VFI": "RIFE VFI", "RoundToEven": "Round To Even","MoveToDevice": "Move To Device","FaceSwapNode": "Face Swap",    "MathAdd": "Math Add", "MathSubtract": "Math Subtract", "MathMultiply": "Math Multiply", "MathDivide": "Math Divide", "FloatToInt": "Float To Int",}