import os
import time
import json
import torch
import numpy as np
from PIL import Image
import folder_paths

try:
    from server import PromptServer
    from aiohttp import web
    SERVER_AVAILABLE = True
except ImportError:
    SERVER_AVAILABLE = False

PENDING_DECISIONS = {}

class RSCropImage:
    @classmethod
    def INPUT_TYPES(cls):
        inputs = {
            "required": {
                "image": ("IMAGE",),
            },
            "optional": {
                "crop_data": ("STRING", {"multiline": False, "default": "{}"}),
            }
        }
        if SERVER_AVAILABLE:
            inputs["hidden"] = {
                "unique_id": "UNIQUE_ID",
                "prompt_id": "PROMPT_ID",
            }
        return inputs

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("CROPPED_IMAGE", "CROP_DATA")
    FUNCTION = "crop_image"
    CATEGORY = "🦊 RaykoStudio"
    OUTPUT_NODE = True

    def crop_image(self, image, crop_data="{}", unique_id=None, prompt_id=None):
        unique_id = str(unique_id) if unique_id else "unknown"
        
        if SERVER_AVAILABLE and unique_id:
            if unique_id not in PENDING_DECISIONS:
                PENDING_DECISIONS[unique_id] = {
                    "status": "pending",
                    "crop_data": crop_data
                }
            
            if prompt_id is None:
                import uuid
                prompt_id = str(uuid.uuid4())

            try:
                img_tensor = image[0]
                i = 255. * img_tensor.cpu().numpy()
                img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
                
                filename = f"rscrop_{unique_id}_{prompt_id}.png"
                subfolder = "rscrop"
                full_output_folder = os.path.join(folder_paths.get_temp_directory(), subfolder)
                os.makedirs(full_output_folder, exist_ok=True)
                img.save(os.path.join(full_output_folder, filename))

                PromptServer.instance.send_sync("rayko.rscrop.show", {
                    "node_id": unique_id,
                    "image_url": f"/view?filename={filename}&type=temp&subfolder={subfolder}",
                    "image_width": img.width,
                    "image_height": img.height
                })
            except Exception as e:
                print(f"[RS Crop 🦊] Error saving image: {e}")

            while True:
                state = PENDING_DECISIONS.get(unique_id, {})
                current_status = state.get("status", "pending")
                current_crop_data = state.get("crop_data", crop_data)
                
                if current_status == "approved":
                    crop_data = current_crop_data
                    if unique_id in PENDING_DECISIONS:
                        del PENDING_DECISIONS[unique_id]
                    break
                
                if current_status == "cancelled":
                    if unique_id in PENDING_DECISIONS:
                        del PENDING_DECISIONS[unique_id]
                    return (image, "{}")
                    
                if current_status == "removed":
                    if unique_id in PENDING_DECISIONS:
                        del PENDING_DECISIONS[unique_id]
                    return (image, "{}")
                
                if current_status == "rejected":
                    PENDING_DECISIONS[unique_id]["status"] = "pending"
                    crop_data = current_crop_data
                
                time.sleep(0.3)

        try:
            crop_dict = json.loads(crop_data.replace("'", '"'))
        except Exception:
            crop_dict = {}
        
        x = int(crop_dict.get("x", 0))
        y = int(crop_dict.get("y", 0))
        width = int(crop_dict.get("width", 512))
        height = int(crop_dict.get("height", 512))
        
        if image is None or len(image) == 0:
            h, w = 512, 512
            crop_output = {
                "source_width": w,
                "source_height": h,
                "x": 0,
                "y": 0,
                "width": 512,
                "height": 512
            }
            return (torch.zeros((1, h, w, 3)), json.dumps(crop_output))
        
        img_tensor = image[0]
        img_height = img_tensor.shape[0]
        img_width = img_tensor.shape[1]
        
        x = max(0, min(x, img_width - 1))
        y = max(0, min(y, img_height - 1))
        x2 = min(x + width, img_width)
        y2 = min(y + height, img_height)
        
        actual_width = max(1, x2 - x)
        actual_height = max(1, y2 - y)
        
        if len(img_tensor.shape) == 3:
            cropped = img_tensor[y:y2, x:x2, :].unsqueeze(0)
        else:
            cropped = img_tensor[:, y:y2, x:x2, :]
        
        crop_output = {
            "source_width": img_width,
            "source_height": img_height,
            "x": x,
            "y": y,
            "width": actual_width,
            "height": actual_height
        }
        
        return (cropped, json.dumps(crop_output))


if SERVER_AVAILABLE:
    @PromptServer.instance.routes.post("/rayko/rscrop/decision")
    async def rscrop_decision(request):
        try:
            data = await request.json()
            node_id = str(data.get("node_id"))
            decision = data.get("decision")
            crop_data = data.get("crop_data")
            
            if node_id in PENDING_DECISIONS:
                if decision == "approve":
                    PENDING_DECISIONS[node_id]["status"] = "approved"
                    if crop_data:
                        PENDING_DECISIONS[node_id]["crop_data"] = crop_data
                
                elif decision == "reject":
                    PENDING_DECISIONS[node_id]["status"] = "rejected"
                    if crop_data is not None:
                        PENDING_DECISIONS[node_id]["crop_data"] = crop_data
                
                elif decision == "cancel":
                    PENDING_DECISIONS[node_id]["status"] = "cancelled"
                    
                return web.Response(status=200, text="Decision recorded")
            else:
                return web.Response(status=404, text=f"Node {node_id} not waiting")
                
        except Exception as e:
            return web.Response(status=500, text=str(e))

    @PromptServer.instance.routes.post("/rayko/rscrop/cleanup")
    async def rscrop_cleanup(request):
        try:
            data = await request.json()
            node_id = str(data.get("node_id"))
            
            if node_id in PENDING_DECISIONS:
                PENDING_DECISIONS[node_id]["status"] = "removed"
                return web.Response(status=200, text="Cleanup recorded")
            else:
                return web.Response(status=200, text="Node not found")
                
        except Exception as e:
            return web.Response(status=500, text=str(e))


NODE_CLASS_MAPPINGS = {"RSCropImage": RSCropImage}
NODE_DISPLAY_NAME_MAPPINGS = {"RSCropImage": "🦊 RS Crop Image"}