# SPDX-License-Identifier: Apache-2.0
# Copyright 2025-2026 Raykosan (RaykoStudio)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import torch
import numpy as np
import cv2
import json
import time
import folder_paths
from PIL import Image
from collections import OrderedDict
from server import PromptServer
from aiohttp import web
import aiohttp

class LRUCache:
    def __init__(self, max_size: int = 50):
        self.cache = OrderedDict()
        self.max_size = max_size

    def get(self, key):
        if key not in self.cache:
            return None
        self.cache.move_to_end(key)
        return self.cache[key]

    def put(self, key, value):
        if key in self.cache:
            self.cache.move_to_end(key)
        self.cache[key] = value
        if len(self.cache) > self.max_size:
            self.cache.popitem(last=False)

    def clear(self):
        self.cache.clear()

PENDING_ADJUSTMENTS = LRUCache(max_size=50)
PREVIEW_CACHE = LRUCache(max_size=20)
LUT_CACHE = LRUCache(max_size=10)
BATCH_PARAMS = {}  # Глобальное хранилище для параметров батча {node_id: adjustments}

def tensor2pil(tensor):
    arr = (tensor.cpu().numpy() * 255.0).astype(np.uint8)
    if arr.ndim == 4:
        arr = arr[0]
    if arr.shape[2] >= 4:
        return Image.fromarray(arr[:, :, :4], 'RGBA').convert('RGB')
    return Image.fromarray(arr, 'RGB')

def pil2tensor(pil_img):
    arr = np.array(pil_img).astype(np.float32) / 255.0
    return torch.from_numpy(arr).unsqueeze(0)

def apply_basic_adjustments(img_bgr, brightness, contrast, hue, saturation):
    if brightness != 0 or contrast != 0:
        alpha = 1.0 + (contrast / 100.0)
        beta = brightness
        img_bgr = img_bgr.astype(np.float32)
        img_bgr = img_bgr * alpha + beta
        img_bgr = np.clip(img_bgr, 0, 255).astype(np.uint8)
    
    if hue != 0 or saturation != 0:
        img_hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
        img_hsv[:, :, 0] = (img_hsv[:, :, 0] + hue) % 180
        img_hsv[:, :, 1] = np.clip(img_hsv[:, :, 1] * (1.0 + saturation / 100.0), 0, 255)
        img_bgr = cv2.cvtColor(img_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
    
    return img_bgr

def apply_sharpen(img_bgr, amount):
    if amount <= 0:
        return img_bgr
    
    amount = min(amount, 200) / 100.0
    blurred = cv2.GaussianBlur(img_bgr, (0, 0), 3)
    sharpened = cv2.addWeighted(img_bgr, 1.0 + amount, blurred, -amount, 0)
    return sharpened

def apply_vibrance(img_bgr, amount):
    if amount == 0:
        return img_bgr
    
    amount = amount / 100.0
    img_float = img_bgr.astype(np.float32) / 255.0
    
    max_rgb = np.max(img_float, axis=2)
    avg_rgb = np.mean(img_float, axis=2)
    saturation = max_rgb - avg_rgb
    
    mask = 1.0 - saturation
    
    result = img_float.copy()
    for c in range(3):
        result[:, :, c] = img_float[:, :, c] + (img_float[:, :, c] - avg_rgb[:, :, np.newaxis][:, :, 0]) * amount * mask
    
    result = np.clip(result * 255, 0, 255).astype(np.uint8)
    return result

def apply_clarity(img_bgr, amount):
    if amount == 0:
        return img_bgr
    
    amount = amount / 100.0
    blurred = cv2.GaussianBlur(img_bgr, (0, 0), 30)
    clarified = cv2.addWeighted(img_bgr, 1.0 + amount, blurred, -amount, 0)
    return clarified

def parse_cube_lut(lut_path):
    try:
        mtime = os.path.getmtime(lut_path)
        cache_key = f"{lut_path}_{int(mtime)}"
    except:
        cache_key = lut_path
    
    cached = LUT_CACHE.get(cache_key)
    if cached is not None:
        return cached
    
    with open(lut_path, 'r') as f:
        lines = f.readlines()
    
    size = 33
    lut_data = []
    for line in lines:
        line = line.strip()
        if line.startswith('LUT_3D_SIZE'):
            size = int(line.split()[1])
        elif line and not line.startswith('#') and not line.startswith('TITLE') and not line.startswith('DOMAIN'):
            try:
                lut_data.append([float(x) for x in line.split()])
            except ValueError:
                continue
    
    lut_array = np.array(lut_data, dtype=np.float32).reshape((size, size, size, 3))
    lut_array = np.transpose(lut_array, (2, 1, 0, 3))
    LUT_CACHE.put(cache_key, lut_array)
    return lut_array

def apply_lut(img_bgr, lut_path, intensity):
    if intensity <= 0.01:
        return img_bgr
    
    if not os.path.exists(lut_path):
        temp_dir = folder_paths.get_temp_directory()
        full_path = os.path.join(temp_dir, "rs_adjustments", lut_path)
        if not os.path.exists(full_path):
            return img_bgr
        lut_path = full_path
    
    lut_array = parse_cube_lut(lut_path)
    size = lut_array.shape[0]
    
    img_float = img_bgr.astype(np.float32) / 255.0
    img_rgb = img_float[:, :, ::-1]
    
    coords = img_rgb * (size - 1)
    idx = np.floor(coords).astype(np.int32)
    idx = np.clip(idx, 0, size - 2)
    frac = coords - idx
    
    i0, i1 = idx[:, :, 0], idx[:, :, 0] + 1
    j0, j1 = idx[:, :, 1], idx[:, :, 1] + 1
    k0, k1 = idx[:, :, 2], idx[:, :, 2] + 1
    
    f0, f1, f2 = frac[:, :, 0], frac[:, :, 1], frac[:, :, 2]
    
    w000 = (1 - f0) * (1 - f1) * (1 - f2)
    w001 = (1 - f0) * (1 - f1) * f2
    w010 = (1 - f0) * f1 * (1 - f2)
    w011 = (1 - f0) * f1 * f2
    w100 = f0 * (1 - f1) * (1 - f2)
    w101 = f0 * (1 - f1) * f2
    w110 = f0 * f1 * (1 - f2)
    w111 = f0 * f1 * f2
    
    result = np.zeros_like(img_rgb)
    for c in range(3):
        lut_c = lut_array[:, :, :, c]
        result[:, :, c] = (
            w000 * lut_c[i0, j0, k0] +
            w001 * lut_c[i0, j0, k1] +
            w010 * lut_c[i0, j1, k0] +
            w011 * lut_c[i0, j1, k1] +
            w100 * lut_c[i1, j0, k0] +
            w101 * lut_c[i1, j0, k1] +
            w110 * lut_c[i1, j1, k0] +
            w111 * lut_c[i1, j1, k1]
        )
    
    result = np.clip(result, 0, 1)
    result_bgr = result[:, :, ::-1]
    
    if intensity < 1.0:
        result_bgr = img_float * (1.0 - intensity) + result_bgr * intensity
        return np.clip(result_bgr * 255, 0, 255).astype(np.uint8)
    
    return (result_bgr * 255).astype(np.uint8)

def apply_levels(img_bgr, input_black, gamma, input_white):
    if input_black == 0 and gamma == 1.0 and input_white == 255:
        return img_bgr
    
    img_float = img_bgr.astype(np.float32) / 255.0
    in_min = input_black / 255.0
    in_max = input_white / 255.0
    in_range = max(0.001, in_max - in_min)
    
    normalized = (img_float - in_min) / in_range
    normalized = np.clip(normalized, 0, 1)
    
    if gamma != 1.0:
        normalized = np.power(normalized, 1.0 / gamma)
    
    result = (normalized * 255).astype(np.uint8)
    return result

def apply_exposure(img_bgr, exposure, offset):
    if exposure == 0 and offset == 0:
        return img_bgr
    
    img_float = img_bgr.astype(np.float32) / 255.0
    multiplier = 2.0 ** (exposure / 100.0)
    img_float = img_float * multiplier + (offset / 100.0)
    
    result = np.clip(img_float * 255, 0, 255).astype(np.uint8)
    return result

def apply_color_balance(img_bgr, shadows_cyan_red, shadows_magenta_green, shadows_yellow_blue,
                        midtones_cyan_red, midtones_magenta_green, midtones_yellow_blue,
                        highlights_cyan_red, highlights_magenta_green, highlights_yellow_blue):
    if all(v == 0 for v in [shadows_cyan_red, shadows_magenta_green, shadows_yellow_blue,
                            midtones_cyan_red, midtones_magenta_green, midtones_yellow_blue,
                            highlights_cyan_red, highlights_magenta_green, highlights_yellow_blue]):
        return img_bgr
    
    img_float = img_bgr.astype(np.float32) / 255.0
    luminance = 0.299 * img_float[:, :, 2] + 0.587 * img_float[:, :, 1] + 0.114 * img_float[:, :, 0]
    
    shadows_mask = np.clip(1.0 - luminance * 2.0, 0, 1)[:, :, np.newaxis]
    midtones_mask = np.clip(1.0 - np.abs(luminance - 0.5) * 4.0, 0, 1)[:, :, np.newaxis]
    highlights_mask = np.clip(luminance * 2.0 - 1.0, 0, 1)[:, :, np.newaxis]
    
    result = img_float.copy()
    
    result[:, :, 2] += shadows_mask[:, :, 0] * (shadows_cyan_red / 100.0)
    result[:, :, 1] += shadows_mask[:, :, 0] * (shadows_magenta_green / 100.0)
    result[:, :, 0] += shadows_mask[:, :, 0] * (shadows_yellow_blue / 100.0)
    
    result[:, :, 2] += midtones_mask[:, :, 0] * (midtones_cyan_red / 100.0)
    result[:, :, 1] += midtones_mask[:, :, 0] * (midtones_magenta_green / 100.0)
    result[:, :, 0] += midtones_mask[:, :, 0] * (midtones_yellow_blue / 100.0)
    
    result[:, :, 2] += highlights_mask[:, :, 0] * (highlights_cyan_red / 100.0)
    result[:, :, 1] += highlights_mask[:, :, 0] * (highlights_magenta_green / 100.0)
    result[:, :, 0] += highlights_mask[:, :, 0] * (highlights_yellow_blue / 100.0)
    
    result = np.clip(result * 255, 0, 255).astype(np.uint8)
    return result

def apply_black_and_white(img_bgr, red, yellow, green, cyan, blue, magenta):
    img_float = img_bgr.astype(np.float32) / 255.0
    img_rgb = img_float[:, :, ::-1]
    img_hsv = cv2.cvtColor((img_rgb * 255).astype(np.uint8), cv2.COLOR_RGB2HSV)
    hue = img_hsv[:, :, 0].astype(np.float32) * 2.0

    w_red = red / 100.0
    w_yellow = yellow / 100.0
    w_green = green / 100.0
    w_cyan = cyan / 100.0
    w_blue = blue / 100.0
    w_magenta = magenta / 100.0

    ref_hues = np.array([0, 30, 60, 90, 120, 150], dtype=np.float32)
    ref_weights = np.array([
        [w_red, 0, 0],
        [w_yellow, w_yellow, 0],
        [0, w_green, 0],
        [0, w_cyan, w_cyan],
        [0, 0, w_blue],
        [w_magenta, 0, w_magenta]
    ], dtype=np.float32)

    ref_hues_ext = np.array([0, 30, 60, 90, 120, 150, 360])
    ref_weights_ext = np.concatenate([ref_weights, ref_weights[:1]], axis=0)

    hue_norm = hue % 360
    idx = np.digitize(hue_norm, ref_hues_ext, right=False) - 1
    idx = np.clip(idx, 0, 5)

    left_weights = ref_weights[idx]
    right_weights = ref_weights_ext[idx + 1]
    left_angle = ref_hues[idx]
    right_angle = ref_hues_ext[idx + 1]
    t = (hue_norm - left_angle) / (right_angle - left_angle + 1e-6)
    weights = left_weights + t[:, :, np.newaxis] * (right_weights - left_weights)

    r = img_rgb[:, :, 0]
    g = img_rgb[:, :, 1]
    b = img_rgb[:, :, 2]
    gray = r * weights[:, :, 0] + g * weights[:, :, 1] + b * weights[:, :, 2]
    gray = np.clip(gray, 0, 1)

    result = (gray * 255).astype(np.uint8)
    result_bgr = cv2.cvtColor(result, cv2.COLOR_GRAY2BGR)
    return result_bgr

def apply_channel_mixer(img_bgr, red_red, red_green, red_blue,
                        green_red, green_green, green_blue,
                        blue_red, blue_green, blue_blue):
    if (red_red == 100 and red_green == 0 and red_blue == 0 and
        green_red == 0 and green_green == 100 and green_blue == 0 and
        blue_red == 0 and blue_green == 0 and blue_blue == 100):
        return img_bgr
    
    img_float = img_bgr.astype(np.float32) / 255.0
    img_rgb = img_float[:, :, ::-1]
    
    r, g, b = img_rgb[:, :, 0], img_rgb[:, :, 1], img_rgb[:, :, 2]
    
    new_r = (r * (red_red / 100.0) + g * (red_green / 100.0) + b * (red_blue / 100.0))
    new_g = (r * (green_red / 100.0) + g * (green_green / 100.0) + b * (green_blue / 100.0))
    new_b = (r * (blue_red / 100.0) + g * (blue_green / 100.0) + b * (blue_blue / 100.0))
    
    result_rgb = np.stack([new_r, new_g, new_b], axis=-1)
    result_rgb = np.clip(result_rgb, 0, 1)
    result_bgr = (result_rgb[:, :, ::-1] * 255).astype(np.uint8)
    
    return result_bgr

def apply_selective_color(img_bgr, color_name, cyan, magenta, yellow, black):
    if cyan == 0 and magenta == 0 and yellow == 0 and black == 0:
        return img_bgr
    
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    img_hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
    
    h_deg = img_hsv[:, :, 0] * 2.0
    s_norm = img_hsv[:, :, 1] / 255.0
    v_norm = img_hsv[:, :, 2] / 255.0
    
    color_centers = {
        "Reds": 0, "Yellows": 60, "Greens": 120, 
        "Cyans": 180, "Blues": 240, "Magentas": 300
    }
    
    if color_name in color_centers:
        center = color_centers[color_name]
        diff = np.abs(h_deg - center)
        diff = np.minimum(diff, 360.0 - diff)
        mask_h = np.exp(-(diff ** 2) / (2 * 40.0 ** 2))
        mask_s = 1.0 - np.exp(-(s_norm ** 2) / (2 * 0.15 ** 2))
        mask = mask_h * mask_s
    else:
        if color_name == "Whites":
            mask_v = np.exp(-((v_norm - 1.0) ** 2) / (2 * 0.15 ** 2))
            mask_s = np.exp(-(s_norm ** 2) / (2 * 0.25 ** 2))
            mask = mask_v * mask_s
        elif color_name == "Neutrals":
            mask_v = np.exp(-((v_norm - 0.5) ** 2) / (2 * 0.15 ** 2))
            mask_s = np.exp(-(s_norm ** 2) / (2 * 0.25 ** 2))
            mask = mask_v * mask_s
        elif color_name == "Blacks":
            mask_v = np.exp(-(v_norm ** 2) / (2 * 0.15 ** 2))
            mask_s = np.exp(-(s_norm ** 2) / (2 * 0.25 ** 2))
            mask = mask_v * mask_s
        else:
            mask = np.zeros_like(h_deg)

    c, m, y, k = cyan / 100.0, magenta / 100.0, yellow / 100.0, black / 100.0
    
    img_rgb[:, :, 0] += mask * (-c - k)
    img_rgb[:, :, 1] += mask * (-m - k)
    img_rgb[:, :, 2] += mask * (-y - k)
    
    img_rgb = np.clip(img_rgb, 0, 1)
    result_bgr = (cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) * 255).astype(np.uint8)
    return result_bgr

def process_image(img_bgr, adjustments):
    brightness = adjustments.get("brightness", 0)
    contrast = adjustments.get("contrast", 0)
    hue = adjustments.get("hue", 0)
    saturation = adjustments.get("saturation", 0)
    
    result_bgr = apply_basic_adjustments(img_bgr, brightness, contrast, hue, saturation)
    
    sharpen = adjustments.get("sharpen", 0)
    if sharpen != 0:
        result_bgr = apply_sharpen(result_bgr, sharpen)
    
    lut_path = adjustments.get("lut_path", "")
    lut_intensity = adjustments.get("lut_intensity", 100)
    if lut_path:
        result_bgr = apply_lut(result_bgr, lut_path, lut_intensity / 100.0)

    vibrance = adjustments.get("vibrance", 0)
    if vibrance != 0:
        result_bgr = apply_vibrance(result_bgr, vibrance)
    
    clarity = adjustments.get("clarity", 0)
    if clarity != 0:
        result_bgr = apply_clarity(result_bgr, clarity)
    
    levels = adjustments.get("levels", {})
    result_bgr = apply_levels(
        result_bgr,
        levels.get("input_black", 0),
        levels.get("gamma", 1.0),
        levels.get("input_white", 255)
    )
    
    exposure = adjustments.get("exposure", {})
    result_bgr = apply_exposure(
        result_bgr,
        exposure.get("exposure", 0),
        exposure.get("offset", 0)
    )
    
    color_balance = adjustments.get("color_balance", {})
    result_bgr = apply_color_balance(
        result_bgr,
        color_balance.get("shadows_cyan_red", 0),
        color_balance.get("shadows_magenta_green", 0),
        color_balance.get("shadows_yellow_blue", 0),
        color_balance.get("midtones_cyan_red", 0),
        color_balance.get("midtones_magenta_green", 0),
        color_balance.get("midtones_yellow_blue", 0),
        color_balance.get("highlights_cyan_red", 0),
        color_balance.get("highlights_magenta_green", 0),
        color_balance.get("highlights_yellow_blue", 0)
    )
    
    bw = adjustments.get("black_white", {})
    bw_enabled = adjustments.get("bw_enabled", False)
    if bw_enabled:
        result_bgr = apply_black_and_white(
            result_bgr,
            bw.get("bw_red", 40),
            bw.get("bw_yellow", 60),
            bw.get("bw_green", 40),
            bw.get("bw_cyan", 60),
            bw.get("bw_blue", 20),
            bw.get("bw_magenta", 80)
        )

    channel_mixer = adjustments.get("channel_mixer", {})
    result_bgr = apply_channel_mixer(
        result_bgr,
        channel_mixer.get("red_red", 100),
        channel_mixer.get("red_green", 0),
        channel_mixer.get("red_blue", 0),
        channel_mixer.get("green_red", 0),
        channel_mixer.get("green_green", 100),
        channel_mixer.get("green_blue", 0),
        channel_mixer.get("blue_red", 0),
        channel_mixer.get("blue_green", 0),
        channel_mixer.get("blue_blue", 100)
    )
    
    selective_color = adjustments.get("selective_color", {})
    result_bgr = apply_selective_color(
        result_bgr,
        selective_color.get("color_name", "Reds"),
        selective_color.get("sc_cyan", 0),
        selective_color.get("sc_magenta", 0),
        selective_color.get("sc_yellow", 0),
        selective_color.get("sc_black", 0)
    )
    
    return result_bgr

def cleanup_temp_files(node_id, temp_dir, exclude_file=None):
    prefix = f"rs_adjustments_{node_id}_"
    for f in os.listdir(temp_dir):
        if f.startswith(prefix) and f != exclude_file:
            try:
                os.remove(os.path.join(temp_dir, f))
            except:
                pass

@PromptServer.instance.routes.post("/rayko/rs_adjustments")
async def adjustments_handler(request):
    try:
        data = await request.json()
        node_id = str(data.get("id"))
        adjustments = data.get("adjustments", {})
        is_batch = data.get("is_batch", False)
        
        if node_id is None or adjustments is None:
            return web.json_response({"error": "Missing data"}, status=400)
        
        # Если включен Batch - сохраняем параметры глобально
        if is_batch:
            BATCH_PARAMS[node_id] = adjustments
            
        PENDING_ADJUSTMENTS.put(node_id, {
            "status": "completed",
            "adjustments": adjustments
        })
        
        temp_dir = folder_paths.get_temp_directory()
        cleanup_temp_files(node_id, temp_dir)
        
        return web.json_response({"status": "ok"})
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)

@PromptServer.instance.routes.post("/rayko/rs_adjustments/cancel")
async def cancel_handler(request):
    try:
        data = await request.json()
        node_id = str(data.get("node_id"))
        if node_id:
            PENDING_ADJUSTMENTS.put(node_id, {"status": "cancelled"})
            # Очищаем batch-параметры при отмене
            BATCH_PARAMS.pop(node_id, None)
            temp_dir = folder_paths.get_temp_directory()
            cleanup_temp_files(node_id, temp_dir)
        return web.json_response({"status": "cancelled"})
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)

@PromptServer.instance.routes.post("/rayko/rs_adjustments/preview")
async def preview_handler(request):
    try:
        data = await request.json()
        node_id = str(data.get("node_id", "unknown"))
        adjustments = data.get("adjustments", {})
        image_file = data.get("image_file", "")
        
        cache_key = f"{node_id}_{image_file}_{json.dumps(adjustments, sort_keys=True)}"
        cached = PREVIEW_CACHE.get(cache_key)
        if cached:
            return web.json_response(cached)
        
        temp_dir = folder_paths.get_temp_directory()
        img_path = os.path.join(temp_dir, image_file)
        
        if not os.path.exists(img_path):
            return web.json_response({"error": "Image not found"}, status=404)
        
        img_pil = Image.open(img_path).convert('RGB')
        img_bgr = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
        
        result_bgr = process_image(img_bgr, adjustments)
        
        ts = int(time.time() * 1000)
        filename = f"rs_adjustments_{node_id}_{ts}_preview.png"
        filepath = os.path.join(temp_dir, filename)
        
        result_pil = Image.fromarray(cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB))
        result_pil.save(filepath)
        
        response_data = {
            "preview_file": filename,
            "timestamp": ts
        }
        
        PREVIEW_CACHE.put(cache_key, response_data)
        
        return web.json_response(response_data)
    except Exception as e:
        import traceback
        traceback.print_exc()
        return web.json_response({"error": str(e)}, status=500)

@PromptServer.instance.routes.get("/rayko/rs_adjustments/ws")
async def ws_preview_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    
    print("[RS Adjustments] WebSocket connection opened")
    
    try:
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                    node_id = str(data.get("node_id", "unknown"))
                    adjustments = data.get("adjustments", {})
                    image_file = data.get("image_file", "")
                    
                    temp_dir = folder_paths.get_temp_directory()
                    img_path = os.path.join(temp_dir, image_file)
                    
                    if not os.path.exists(img_path):
                        await ws.send_json({"error": "Image not found"})
                        continue
                    
                    img_pil = Image.open(img_path).convert('RGB')
                    img_bgr = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
                    
                    max_size = 1024
                    h, w = img_bgr.shape[:2]
                    if h > max_size or w > max_size:
                        scale = max_size / max(h, w)
                        new_w = int(w * scale)
                        new_h = int(h * scale)
                        img_bgr = cv2.resize(img_bgr, (new_w, new_h), interpolation=cv2.INTER_AREA)
                    
                    result_bgr = process_image(img_bgr, adjustments)
                    
                    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 90]
                    _, jpeg_buffer = cv2.imencode('.jpg', result_bgr, encode_param)
                    jpeg_bytes = jpeg_buffer.tobytes()
                    
                    await ws.send_bytes(jpeg_bytes)
                    
                except Exception as e:
                    print(f"[RS Adjustments] WebSocket error: {e}")
                    import traceback
                    traceback.print_exc()
                    await ws.send_json({"error": str(e)})
            
            elif msg.type == aiohttp.WSMsgType.ERROR:
                print(f"[RS Adjustments] WebSocket error: {ws.exception()}")
    
    except Exception as e:
        print(f"[RS Adjustments] WebSocket connection error: {e}")
    
    print("[RS Adjustments] WebSocket connection closed")
    return ws

class RS_ImageAdjustments:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "brightness": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
                "contrast": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
                "sharpen": ("INT", {"default": 0, "min": 0, "max": 200, "step": 1}),
                "hue": ("INT", {"default": 0, "min": -180, "max": 180, "step": 1}),
                "saturation": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
                "advanced_params": "STRING"
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "apply_adjustments"
    CATEGORY = "🦊 RaykoStudio"
    OUTPUT_NODE = True
    DESCRIPTION = "Interactive image adjustments with real-time preview"

    def apply_adjustments(self, image, brightness, contrast, sharpen, hue, saturation, unique_id=None, advanced_params=None):
        unique_id = str(unique_id) if unique_id is not None else "unknown"
        
        # Создаем безопасный fallback-тензор на случай ошибок/отмены
        safe_fallback = torch.zeros((1, 512, 512, 3)) 
        
        if image.numel() == 0:
            return (safe_fallback,)
        
        h, w = image.shape[1], image.shape[2]
        if h < 8 or w < 8:
            return (safe_fallback,)
        
        PENDING_ADJUSTMENTS.cache.pop(unique_id, None)
        
        bg_pil = tensor2pil(image)
        bg_w, bg_h = bg_pil.size
        
        ts = int(time.time() * 1000)
        td = folder_paths.get_temp_directory()
        
        bfn = f"rs_adjustments_{unique_id}_{ts}_bg.png"
        bg_pil.save(os.path.join(td, bfn))
        
        cleanup_temp_files(unique_id, td, exclude_file=bfn)
        
        # === BATCH MODE CHECK ===
        # Проверяем ДО входа в режим ожидания
        if unique_id in BATCH_PARAMS:
            print(f"[RS Adjustments] Node {unique_id}: Applying saved BATCH parameters...")
            adjustments = BATCH_PARAMS[unique_id].copy()
            
            # Гарантируем наличие базовых параметров
            adjustments["brightness"] = adjustments.get("brightness", brightness)
            adjustments["contrast"] = adjustments.get("contrast", contrast)
            adjustments["sharpen"] = adjustments.get("sharpen", sharpen)
            adjustments["hue"] = adjustments.get("hue", hue)
            adjustments["saturation"] = adjustments.get("saturation", saturation)
            
            img_bgr = cv2.cvtColor(np.array(bg_pil), cv2.COLOR_RGB2BGR)
            result_bgr = process_image(img_bgr, adjustments)
            
            result_pil = Image.fromarray(cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB))
            cleanup_temp_files(unique_id, td)
            
            # Удаляем batch-параметры ТОЛЬКО после успешной обработки
            # Это гарантирует, что следующий кадр уже не будет batch
            BATCH_PARAMS.pop(unique_id, None)
            
            return (pil2tensor(result_pil),)
        
        # === ОБЫЧНЫЙ РЕЖИМ ===
        PENDING_ADJUSTMENTS.put(unique_id, {"status": "pending"})
        
        time.sleep(0.5)
        
        PromptServer.instance.send_sync("rs-adjustments-start", {
            "id": unique_id,
            "bg_file": bfn,
            "bg_width": bg_w,
            "bg_height": bg_h,
            "timestamp": ts,
            "brightness": brightness,
            "contrast": contrast,
            "sharpen": sharpen,
            "hue": hue,
            "saturation": saturation,
            "advanced_params": advanced_params or "{}"
        })
        
        print(f"[RS Image Adjustments] Node {unique_id} waiting for user input (Apply/Cancel)...")
        
        while unique_id not in PENDING_ADJUSTMENTS.cache or PENDING_ADJUSTMENTS.cache[unique_id].get("status") == "pending":
            if unique_id in PENDING_ADJUSTMENTS.cache and PENDING_ADJUSTMENTS.cache[unique_id].get("status") == "cancelled":
                break
            time.sleep(0.05)
        
        if unique_id in PENDING_ADJUSTMENTS.cache:
            if PENDING_ADJUSTMENTS.cache[unique_id].get("status") == "cancelled":
                print(f"[RS Image Adjustments] Cancelled by user for node {unique_id}")
                PENDING_ADJUSTMENTS.cache.pop(unique_id, None)
                BATCH_PARAMS.pop(unique_id, None)  # Очищаем batch при отмене
                cleanup_temp_files(unique_id, td)
                return (safe_fallback,)
        
        print(f"[RS Image Adjustments] Input received for node {unique_id}")
        
        data = PENDING_ADJUSTMENTS.cache.pop(unique_id, None)
        if data is None:
            cleanup_temp_files(unique_id, td)
            return (safe_fallback,)
        
        adjustments = data.get("adjustments", {})
        
        brightness = adjustments.get("brightness", brightness)
        contrast = adjustments.get("contrast", contrast)
        sharpen = adjustments.get("sharpen", sharpen)
        hue = adjustments.get("hue", hue)
        saturation = adjustments.get("saturation", saturation)
        
        img_bgr = cv2.cvtColor(np.array(bg_pil), cv2.COLOR_RGB2BGR)
        
        adjustments["brightness"] = brightness
        adjustments["contrast"] = contrast
        adjustments["sharpen"] = sharpen
        adjustments["hue"] = hue
        adjustments["saturation"] = saturation
        
        result_bgr = process_image(img_bgr, adjustments)
        
        result_pil = Image.fromarray(cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB))
        
        cleanup_temp_files(unique_id, td)
        
        return (pil2tensor(result_pil),)

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return time.time()

NODE_CLASS_MAPPINGS = {"RS_ImageAdjustments": RS_ImageAdjustments}
NODE_DISPLAY_NAME_MAPPINGS = {"RS_ImageAdjustments": "🦊 RS Image Adjustments"}