"""
GPU-Accelerated Color Operations Library

Uses CuPy for CUDA-accelerated color transforms.

Performance Mode:
- Set STREAMING_GPU_PERSIST=1 to keep frames on GPU between operations
- This dramatically improves performance by avoiding CPU<->GPU transfers
"""
import os
import numpy as np

# Try to import CuPy for GPU acceleration
try:
    import cupy as cp
    GPU_AVAILABLE = True
    print("[color_ops_gpu] CuPy GPU acceleration enabled")
except ImportError:
    cp = np
    GPU_AVAILABLE = False
    print("[color_ops_gpu] CuPy not available, using CPU fallback")

# GPU persistence mode - keep frames on GPU between operations
GPU_PERSIST = os.environ.get("STREAMING_GPU_PERSIST", "0") == "1"
if GPU_AVAILABLE and GPU_PERSIST:
    print("[color_ops_gpu] GPU persistence enabled - frames stay on GPU")


def _to_gpu(img):
    """Move image to GPU if available."""
    if GPU_AVAILABLE and not isinstance(img, cp.ndarray):
        return cp.asarray(img)
    return img


def _to_cpu(img):
    """Move image back to CPU (only if GPU_PERSIST is disabled)."""
    if not GPU_PERSIST and GPU_AVAILABLE and isinstance(img, cp.ndarray):
        return cp.asnumpy(img)
    return img


def prim_invert(img):
    """Invert image colors."""
    if GPU_AVAILABLE:
        img_gpu = _to_gpu(img)
        return _to_cpu(255 - img_gpu)
    return 255 - img


def prim_grayscale(img):
    """Convert to grayscale."""
    if img.ndim != 3:
        return img

    if GPU_AVAILABLE:
        img_gpu = _to_gpu(img.astype(np.float32))
        # Standard luminance weights
        gray = 0.299 * img_gpu[:, :, 0] + 0.587 * img_gpu[:, :, 1] + 0.114 * img_gpu[:, :, 2]
        gray = cp.clip(gray, 0, 255).astype(cp.uint8)
        # Stack to 3 channels
        result = cp.stack([gray, gray, gray], axis=2)
        return _to_cpu(result)

    gray = 0.299 * img[:, :, 0] + 0.587 * img[:, :, 1] + 0.114 * img[:, :, 2]
    gray = np.clip(gray, 0, 255).astype(np.uint8)
    return np.stack([gray, gray, gray], axis=2)


def prim_brightness(img, factor=1.0):
    """Adjust brightness by factor."""
    xp = cp if GPU_AVAILABLE else np
    if GPU_AVAILABLE:
        img_gpu = _to_gpu(img.astype(np.float32))
        result = xp.clip(img_gpu * factor, 0, 255).astype(xp.uint8)
        return _to_cpu(result)
    return np.clip(img.astype(np.float32) * factor, 0, 255).astype(np.uint8)


def prim_contrast(img, factor=1.0):
    """Adjust contrast around midpoint."""
    xp = cp if GPU_AVAILABLE else np
    if GPU_AVAILABLE:
        img_gpu = _to_gpu(img.astype(np.float32))
        result = xp.clip((img_gpu - 128) * factor + 128, 0, 255).astype(xp.uint8)
        return _to_cpu(result)
    return np.clip((img.astype(np.float32) - 128) * factor + 128, 0, 255).astype(np.uint8)


# CUDA kernel for HSV hue shift
if GPU_AVAILABLE:
    _hue_shift_kernel = cp.RawKernel(r'''
    extern "C" __global__
    void hue_shift(unsigned char* img, int width, int height, float shift) {
        int x = blockDim.x * blockIdx.x + threadIdx.x;
        int y = blockDim.y * blockIdx.y + threadIdx.y;

        if (x >= width || y >= height) return;

        int idx = (y * width + x) * 3;

        // Get RGB
        float r = img[idx] / 255.0f;
        float g = img[idx + 1] / 255.0f;
        float b = img[idx + 2] / 255.0f;

        // RGB to HSV
        float max_c = fmaxf(r, fmaxf(g, b));
        float min_c = fminf(r, fminf(g, b));
        float delta = max_c - min_c;

        float h = 0.0f, s = 0.0f, v = max_c;

        if (delta > 0.00001f) {
            s = delta / max_c;

            if (max_c == r) {
                h = 60.0f * fmodf((g - b) / delta, 6.0f);
            } else if (max_c == g) {
                h = 60.0f * ((b - r) / delta + 2.0f);
            } else {
                h = 60.0f * ((r - g) / delta + 4.0f);
            }

            if (h < 0) h += 360.0f;
        }

        // Shift hue
        h = fmodf(h + shift, 360.0f);
        if (h < 0) h += 360.0f;

        // HSV to RGB
        float c = v * s;
        float x_val = c * (1.0f - fabsf(fmodf(h / 60.0f, 2.0f) - 1.0f));
        float m = v - c;

        float r_out, g_out, b_out;
        if (h < 60) {
            r_out = c; g_out = x_val; b_out = 0;
        } else if (h < 120) {
            r_out = x_val; g_out = c; b_out = 0;
        } else if (h < 180) {
            r_out = 0; g_out = c; b_out = x_val;
        } else if (h < 240) {
            r_out = 0; g_out = x_val; b_out = c;
        } else if (h < 300) {
            r_out = x_val; g_out = 0; b_out = c;
        } else {
            r_out = c; g_out = 0; b_out = x_val;
        }

        img[idx] = (unsigned char)fminf(255.0f, (r_out + m) * 255.0f);
        img[idx + 1] = (unsigned char)fminf(255.0f, (g_out + m) * 255.0f);
        img[idx + 2] = (unsigned char)fminf(255.0f, (b_out + m) * 255.0f);
    }
    ''', 'hue_shift')


def prim_hue_shift(img, shift=0.0):
    """Shift hue by degrees."""
    if img.ndim != 3 or img.shape[2] != 3:
        return img

    if not GPU_AVAILABLE:
        import cv2
        hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
        hsv[:, :, 0] = (hsv[:, :, 0].astype(np.float32) + shift / 2) % 180
        return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)

    h, w = img.shape[:2]
    img_gpu = _to_gpu(img.astype(np.uint8)).copy()

    block = (16, 16)
    grid = ((w + block[0] - 1) // block[0], (h + block[1] - 1) // block[1])

    _hue_shift_kernel(grid, block, (img_gpu, np.int32(w), np.int32(h), np.float32(shift)))

    return _to_cpu(img_gpu)


def prim_saturate(img, factor=1.0):
    """Adjust saturation by factor."""
    if img.ndim != 3:
        return img

    if not GPU_AVAILABLE:
        import cv2
        hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32)
        hsv[:, :, 1] = np.clip(hsv[:, :, 1] * factor, 0, 255)
        return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)

    # GPU version - simple desaturation blend
    img_gpu = _to_gpu(img.astype(np.float32))
    gray = 0.299 * img_gpu[:, :, 0] + 0.587 * img_gpu[:, :, 1] + 0.114 * img_gpu[:, :, 2]
    gray = gray[:, :, cp.newaxis]

    if factor < 1.0:
        # Desaturate: blend toward gray
        result = img_gpu * factor + gray * (1 - factor)
    else:
        # Oversaturate: extrapolate away from gray
        result = gray + (img_gpu - gray) * factor

    result = cp.clip(result, 0, 255).astype(cp.uint8)
    return _to_cpu(result)


def prim_blend(img1, img2, alpha=0.5):
    """Blend two images with alpha."""
    xp = cp if GPU_AVAILABLE else np

    if GPU_AVAILABLE:
        img1_gpu = _to_gpu(img1.astype(np.float32))
        img2_gpu = _to_gpu(img2.astype(np.float32))
        result = img1_gpu * (1 - alpha) + img2_gpu * alpha
        result = xp.clip(result, 0, 255).astype(xp.uint8)
        return _to_cpu(result)

    result = img1.astype(np.float32) * (1 - alpha) + img2.astype(np.float32) * alpha
    return np.clip(result, 0, 255).astype(np.uint8)


def prim_add(img1, img2):
    """Add two images (clamped)."""
    xp = cp if GPU_AVAILABLE else np
    if GPU_AVAILABLE:
        result = xp.clip(_to_gpu(img1).astype(np.int16) + _to_gpu(img2).astype(np.int16), 0, 255)
        return _to_cpu(result.astype(xp.uint8))
    return np.clip(img1.astype(np.int16) + img2.astype(np.int16), 0, 255).astype(np.uint8)


def prim_multiply(img1, img2):
    """Multiply two images (normalized)."""
    xp = cp if GPU_AVAILABLE else np
    if GPU_AVAILABLE:
        result = (_to_gpu(img1).astype(np.float32) * _to_gpu(img2).astype(np.float32)) / 255.0
        result = xp.clip(result, 0, 255).astype(xp.uint8)
        return _to_cpu(result)
    result = (img1.astype(np.float32) * img2.astype(np.float32)) / 255.0
    return np.clip(result, 0, 255).astype(np.uint8)


def prim_screen(img1, img2):
    """Screen blend mode."""
    xp = cp if GPU_AVAILABLE else np
    if GPU_AVAILABLE:
        i1 = _to_gpu(img1).astype(np.float32) / 255.0
        i2 = _to_gpu(img2).astype(np.float32) / 255.0
        result = 1.0 - (1.0 - i1) * (1.0 - i2)
        result = xp.clip(result * 255, 0, 255).astype(xp.uint8)
        return _to_cpu(result)
    i1 = img1.astype(np.float32) / 255.0
    i2 = img2.astype(np.float32) / 255.0
    result = 1.0 - (1.0 - i1) * (1.0 - i2)
    return np.clip(result * 255, 0, 255).astype(np.uint8)


# Import CPU primitives as fallbacks
def _get_cpu_primitives():
    """Get all primitives from CPU color_ops module as fallbacks."""
    from sexp_effects.primitive_libs import color_ops
    return color_ops.PRIMITIVES


# Export functions - start with CPU primitives, then override with GPU versions
PRIMITIVES = _get_cpu_primitives().copy()

# Override specific primitives with GPU-accelerated versions
PRIMITIVES.update({
    'invert': prim_invert,
    'grayscale': prim_grayscale,
    'brightness': prim_brightness,
    'contrast': prim_contrast,
    'hue-shift': prim_hue_shift,
    'saturate': prim_saturate,
    'blend': prim_blend,
    'add': prim_add,
    'multiply': prim_multiply,
    'screen': prim_screen,
})
