# roop_custom_nodes.py
import os
import subprocess
import glob
from PIL import Image
import numpy as np
import torch
import hmac
import hashlib
import requests

# Helper: save tensor image
def save_tensor_as_image(tensor, path):
    np_img = (tensor[0].cpu().numpy() * 255).astype(np.uint8)
    if np_img.shape[0] == 3:  # If format is CHW
        np_img = np.transpose(np_img, (1, 2, 0))  # Convert to HWC for PIL
    img = Image.fromarray(np_img)
    img.save(path)

# Helper: load image as tensor
def load_image_as_tensor(path):
    img = Image.open(path).convert("RGB")
    np_img = np.array(img).astype(np.float32) / 255.0
    return torch.from_numpy(np_img).unsqueeze(0)

def get_unique_filename(path):
    base, ext = os.path.splitext(path)
    counter = 1
    new_path = path
    while os.path.exists(new_path):
        new_path = f"{base}_{counter}{ext}"
        counter += 1
    return new_path

def send_webhook_image(webhook_url, webhook_secret, output_path, extra_data=None):
    if not webhook_url:
        return

    with open(output_path, 'rb') as f:
        image_data = f.read()

    # GitHub-style HMAC signature
    signature = 'sha256=' + hmac.new(
        webhook_secret.encode('utf-8'),
        image_data,
        hashlib.sha256
    ).hexdigest()

    headers = {
        'X-Hub-Signature-256': signature
    }

    files = {
        'file': ('swapped.png', image_data, 'image/png')
    }

    data = extra_data or {}

    try:
        resp = requests.post(webhook_url, headers=headers, files=files, data=data)
        resp.raise_for_status()
        print(f"[Webhook] Sent successfully to {webhook_url}")
    except Exception as e:
        print(f"[Webhook Error] {e}")


class RoopFaceSwap:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_image": ("IMAGE",),
                "target_image": ("IMAGE",),
                "roop_dir": ("STRING", {"default": "/content/roop"}),
                "output_name": ("STRING", {"default": "roop_output.png"}),
                "many_faces": ("BOOLEAN", {"default": False})
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("swapped_image",)
    FUNCTION = "run"
    CATEGORY = "Roop/Basic"

    def run(self, source_image, target_image, roop_dir, output_name, many_faces):
        temp_dir = os.path.join(roop_dir, "temp_io")
        os.makedirs(temp_dir, exist_ok=True)

        source_path = os.path.join(temp_dir, "source.png")
        target_path = os.path.join(temp_dir, "target.png")
        output_path = get_unique_filename(os.path.join(temp_dir, output_name))

        save_tensor_as_image(source_image, source_path)
        save_tensor_as_image(target_image, target_path)

        cmd = [
            "python", "run.py",
            "-s", source_path,
            "-t", target_path,
            "-o", output_path,
            "--execution-provider", "cuda",
            "--frame-processor", "face_swapper"
        ]
        if many_faces:
            cmd.append("--many-faces")

        subprocess.run(cmd, check=True, cwd=roop_dir)

        if not os.path.exists(output_path):
            print(f"[Warning] Roop did not produce output: {output_path}")
            blank = torch.zeros_like(target_image)
            return (blank,)

        return (load_image_as_tensor(output_path),)


class RoopFaceSwapWithEnhancer:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_image": ("IMAGE",),
                "target_image": ("IMAGE",),
                "roop_dir": ("STRING", {"default": "/content/roop"}),
                "output_name": ("STRING", {"default": "roop_output.png"}),
                "many_faces": ("BOOLEAN", {"default": False})
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("swapped_image",)
    FUNCTION = "run"
    CATEGORY = "Roop/Enhanced"

    def run(self, source_image, target_image, roop_dir, output_name, many_faces):
        temp_dir = os.path.join(roop_dir, "temp_io")
        os.makedirs(temp_dir, exist_ok=True)

        source_path = os.path.join(temp_dir, "source.png")
        target_path = os.path.join(temp_dir, "target.png")
        output_path = get_unique_filename(os.path.join(temp_dir, output_name))

        save_tensor_as_image(source_image, source_path)
        save_tensor_as_image(target_image, target_path)

        cmd = [
            "python", "run.py",
            "-s", source_path,
            "-t", target_path,
            "-o", output_path,
            "--execution-provider", "cuda",
            "--frame-processor", "face_swapper", "face_enhancer"
        ]
        if many_faces:
            cmd.append("--many-faces")

        subprocess.run(cmd, check=True, cwd=roop_dir)

        if not os.path.exists(output_path):
            print(f"[Warning] Roop did not produce output: {output_path}")
            blank = torch.zeros_like(target_image)
            return (blank,)

        return (load_image_as_tensor(output_path),)


class RoopBatchFaceSwap:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_image": ("IMAGE",),
                "input_dir": ("STRING", {"default": "/input/images"}),
                "output_dir": ("STRING", {"default": "/output/images"}),
                "roop_dir": ("STRING", {"default": "/content/roop"}),
                "use_enhancer": ("BOOLEAN", {"default": False}),
                "many_faces": ("BOOLEAN", {"default": False})
            }
        }

    RETURN_TYPES = ()
    RETURN_NAMES = ()
    FUNCTION = "run"
    CATEGORY = "Roop/Batch"

    def run(self, source_image, input_dir, output_dir, roop_dir, use_enhancer, many_faces):
        os.makedirs(output_dir, exist_ok=True)
        temp_dir = os.path.join(roop_dir, "temp_io")
        os.makedirs(temp_dir, exist_ok=True)

        source_path = os.path.join(temp_dir, "source.png")
        save_tensor_as_image(source_image, source_path)

        image_paths = glob.glob(os.path.join(input_dir, "*.jpg")) + \
                      glob.glob(os.path.join(input_dir, "*.png"))

        for img_path in image_paths:
            target_name = os.path.basename(img_path)
            target_path = os.path.join(temp_dir, "target.png")
            output_path = os.path.join(output_dir, f"out_{target_name}")

            Image.open(img_path).save(target_path)

            cmd = [
                "python", "run.py",
                "-s", source_path,
                "-t", target_path,
                "-o", output_path,
                "--execution-provider", "cuda",
                "--frame-processor", "face_swapper"
            ]

            if use_enhancer:
                cmd[-1] += " face_enhancer"
            if many_faces:
                cmd.append("--many-faces")

            subprocess.run(cmd, check=True, cwd=roop_dir)

            if not os.path.exists(output_path):
                print(f"[Skipped] NSFW or error: {img_path} -> No output generated.")
                continue

        return ()


class RoopFaceSwapVideo:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_image": ("IMAGE",),
                "target_video_path": ("STRING", {"default": "/path/to/video.mp4"}),
                "roop_dir": ("STRING", {"default": "/content/roop"}),
                "output_name": ("STRING", {"default": "swapped_video.mp4"}),
                "many_faces": ("BOOLEAN", {"default": False})
            }
        }

    RETURN_TYPES = ()
    RETURN_NAMES = ()
    FUNCTION = "run"
    CATEGORY = "Roop/Video"

    def run(self, source_image, target_video_path, roop_dir, output_name, many_faces):
        temp_dir = os.path.join(roop_dir, "temp_io")
        os.makedirs(temp_dir, exist_ok=True)

        source_path = os.path.join(temp_dir, "source.png")
        output_path = get_unique_filename(os.path.join(temp_dir, output_name))
        save_tensor_as_image(source_image, source_path)

        cmd = [
            "python", "run.py",
            "-s", source_path,
            "-t", target_video_path,
            "-o", output_path,
            "--keep-fps",
            "--keep-frames",
            "--execution-provider", "cuda",
            "--frame-processor", "face_swapper"
        ]
        if many_faces:
            cmd.append("--many-faces")

        subprocess.run(cmd, check=True, cwd=roop_dir)

        if not os.path.exists(output_path):
            print(f"[Warning] Roop did not produce video output: {output_path}")

        return ()


class RoopSendWebhookImage:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_tensor": ("IMAGE",),
                "filename": ("STRING", {"default": "output.png"}),
                "webhook_url": ("STRING", {"default": ""}),
                "webhook_secret": ("STRING", {"default": ""}),
                "enable_webhook": ("BOOLEAN", {"default": True}),
                "roop_dir": ("STRING", {"default": "/content/roop"})
            }
        }

    RETURN_TYPES = ()
    RETURN_NAMES = ()
    FUNCTION = "run"
    CATEGORY = "Roop/Webhook"

    def run(self, image_tensor, filename, webhook_url, webhook_secret, enable_webhook, roop_dir):
        if not enable_webhook or not webhook_url:
            print("[WebhookImage] Disabled or URL not set — skipping.")
            return ()

        temp_dir = os.path.join(roop_dir, "temp_io")
        os.makedirs(temp_dir, exist_ok=True)
        output_path = os.path.join(temp_dir, filename)

        save_tensor_as_image(image_tensor, output_path)

        try:
            with open(output_path, 'rb') as f:
                file_data = f.read()

            headers = {}
            if webhook_secret:
                signature = 'sha256=' + hmac.new(
                    webhook_secret.encode('utf-8'),
                    file_data,
                    hashlib.sha256
                ).hexdigest()
                headers['X-Hub-Signature-256'] = signature

            files = {
                'file': (filename, file_data, 'image/png')
            }

            resp = requests.post(webhook_url, headers=headers, files=files)
            resp.raise_for_status()
            print(f"[WebhookImage] Sent image: {filename} → {webhook_url}")
        except Exception as e:
            print(f"[WebhookImage Error] {e}")

        return ()


class RoopSendWebhookFile:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "file_path": ("STRING",),
                "filename": ("STRING", {"default": "output.mp4"}),
                "webhook_url": ("STRING", {"default": ""}),
                "webhook_secret": ("STRING", {"default": ""}),
                "enable_webhook": ("BOOLEAN", {"default": True})
            }
        }

    RETURN_TYPES = ()
    RETURN_NAMES = ()
    FUNCTION = "run"
    CATEGORY = "Roop/Webhook"

    def run(self, file_path, filename, webhook_url, webhook_secret, enable_webhook):
        if not enable_webhook or not webhook_url:
            print("[WebhookFile] Disabled or URL not set — skipping.")
            return ()

        if not os.path.exists(file_path):
            print(f"[WebhookFile] File does not exist: {file_path}")
            return ()

        try:
            with open(file_path, 'rb') as f:
                file_data = f.read()

            headers = {}
            if webhook_secret:
                signature = 'sha256=' + hmac.new(
                    webhook_secret.encode('utf-8'),
                    file_data,
                    hashlib.sha256
                ).hexdigest()
                headers['X-Hub-Signature-256'] = signature

            content_type = "video/mp4" if filename.endswith(".mp4") else "application/octet-stream"
            files = {
                'file': (filename, file_data, content_type)
            }

            resp = requests.post(webhook_url, headers=headers, files=files)
            resp.raise_for_status()
            print(f"[WebhookFile] Sent file: {filename} → {webhook_url}")
        except Exception as e:
            print(f"[WebhookFile Error] {e}")

        return ()


# Register with ComfyUI
NODE_CLASS_MAPPINGS = {
    "RoopFaceSwap": RoopFaceSwap,
    "RoopFaceSwapWithEnhancer": RoopFaceSwapWithEnhancer,
    "RoopBatchFaceSwap": RoopBatchFaceSwap,
    "RoopFaceSwapVideo": RoopFaceSwapVideo,
    "RoopSendWebhookImage": RoopSendWebhookImage,
    "RoopSendWebhookFile": RoopSendWebhookFile,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "RoopFaceSwap": "Roop Face Swap (Image)",
    "RoopFaceSwapWithEnhancer": "Roop Face Swap + Enhancer",
    "RoopBatchFaceSwap": "Roop Batch Image Folder",
    "RoopFaceSwapVideo": "Roop Face Swap (Video)",
    "RoopSendWebhookImage": "Roop Webhook: Image Tensor",
    "RoopSendWebhookFile": "Roop Webhook: File Path",
}