"""
ComfyUI Pixel Sort Node
Integrates ASDFPixelSort Processing script with ComfyUI workflows
"""

import subprocess
import os
import tempfile
import torch
import numpy as np
from PIL import Image

class PixelSortNode:
    """
    A ComfyUI node that applies pixel sorting effects using the ASDFPixelSort Processing script.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "mode": (["white", "black", "bright", "dark"], {"default": "white"}),
                "white_threshold": ("INT", {"default": -12345678, "min": -16777216, "max": 0, "step": 1000}),
                "black_threshold": ("INT", {"default": -3456789, "min": -16777216, "max": 0, "step": 1000}),
                "bright_threshold": ("INT", {"default": 127, "min": 0, "max": 255, "step": 1}),
                "dark_threshold": ("INT", {"default": 223, "min": 0, "max": 255, "step": 1}),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "apply_pixel_sort"
    CATEGORY = "image/effects"
    
    def __init__(self):
        # Path to Processing installation
        self.processing_path = "/root/Processing/processing-4.3/processing-java"
        self.sketch_path = "/root/Processing/ASDFPixelSort"
        
    def apply_pixel_sort(self, image, mode, white_threshold, black_threshold, bright_threshold, dark_threshold):
        """
        Apply pixel sorting to the input image using the Processing script.
        """
        
        # Convert mode string to number
        mode_map = {"white": 0, "black": 1, "bright": 2, "dark": 3}
        mode_num = mode_map[mode]
        
        # Create temporary files
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as input_file:
            input_path = input_file.name
        
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as output_file:
            output_path = output_file.name
        
        try:
            # Convert tensor to PIL Image and save
            img_tensor = image[0]  # Remove batch dimension
            img_array = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
            img_pil = Image.fromarray(img_array)
            img_pil.save(input_path)
            
            # Prepare command arguments
            cmd = [
                self.processing_path,
                f"--sketch={self.sketch_path}",
                "--run",
                input_path,
                output_path,
                str(mode_num),
                str(white_threshold),
                str(black_threshold),
                str(bright_threshold),
                str(dark_threshold)
            ]
            
            # Set up environment for headless operation
            env = os.environ.copy()
            env["DISPLAY"] = ":99"
            
            # Run the Processing script
            print(f"Running pixel sort with mode: {mode}")
            result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120)
            
            if result.returncode != 0:
                print(f"Processing script error: {result.stderr}")
                raise RuntimeError(f"Pixel sort failed: {result.stderr}")
            
            # Load the processed image
            if not os.path.exists(output_path):
                raise RuntimeError("Output image was not created")
                
            img_sorted = Image.open(output_path).convert("RGB")
            img_array = np.array(img_sorted).astype(np.float32) / 255.0
            img_tensor = torch.from_numpy(img_array).unsqueeze(0)  # Add batch dimension
            
            print(f"Pixel sort completed successfully")
            return (img_tensor,)
            
        except subprocess.TimeoutExpired:
            raise RuntimeError("Pixel sort operation timed out")
        except Exception as e:
            raise RuntimeError(f"Pixel sort failed: {str(e)}")
        finally:
            # Clean up temporary files
            try:
                if os.path.exists(input_path):
                    os.unlink(input_path)
                if os.path.exists(output_path):
                    os.unlink(output_path)
            except:
                pass

class PixelSortAdvancedNode:
    """
    Advanced pixel sort node with more control options.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "mode": (["white", "black", "bright", "dark"], {"default": "white"}),
            },
            "optional": {
                "white_threshold": ("INT", {"default": -12345678, "min": -16777216, "max": 0, "step": 1000}),
                "black_threshold": ("INT", {"default": -3456789, "min": -16777216, "max": 0, "step": 1000}),
                "bright_threshold": ("INT", {"default": 127, "min": 0, "max": 255, "step": 1}),
                "dark_threshold": ("INT", {"default": 223, "min": 0, "max": 255, "step": 1}),
                "custom_sketch_path": ("STRING", {"default": ""}),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "apply_pixel_sort_advanced"
    CATEGORY = "image/effects"
    
    def apply_pixel_sort_advanced(self, image, mode, white_threshold=-12345678, black_threshold=-3456789, 
                                bright_threshold=127, dark_threshold=223, custom_sketch_path=""):
        """
        Advanced pixel sorting with optional custom sketch path.
        """
        node = PixelSortNode()
        if custom_sketch_path:
            node.sketch_path = custom_sketch_path
        
        return node.apply_pixel_sort(image, mode, white_threshold, black_threshold, 
                                   bright_threshold, dark_threshold)

# Node registration
NODE_CLASS_MAPPINGS = {
    "PixelSortNode": PixelSortNode,
    "PixelSortAdvancedNode": PixelSortAdvancedNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PixelSortNode": "Pixel Sort (ASDF)",
    "PixelSortAdvancedNode": "Pixel Sort Advanced (ASDF)"
}