import os
import cv2
import numpy as np
from PIL import Image
import onnxruntime as ort
from huggingface_hub import hf_hub_download
import time
from watermark import apply_watermark

# Model paths and repositories
MODELS = {
    "rmbg": {
        "repo": "briaai/RMBG-1.4",
        "filename": "onnx/model.onnx", 
        "session": None
    },
    "depth": {
        "repo": "onnx-community/depth-anything-v2-small",
        "filename": "onnx/model.onnx",
        "session": None
    }
}

def init_models():
    """Download (if needed) and load ONNX models into memory."""
    print("Initializing models...")
    # Configure ONNX Runtime to use CPU and limit threads to save RAM and avoid crashes
    opts = ort.SessionOptions()
    opts.inter_op_num_threads = 1
    opts.intra_op_num_threads = 2
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    
    for name, config in MODELS.items():
        try:
            print(f"Loading {name} model from {config['repo']}...")
            model_path = hf_hub_download(repo_id=config['repo'], filename=config['filename'])
            config['session'] = ort.InferenceSession(model_path, sess_options=opts, providers=['CPUExecutionProvider'])
            print(f"{name} model loaded successfully.")
        except Exception as e:
            print(f"Error loading {name} model: {e}")
            raise e

def preprocess_rmbg(img: np.ndarray) -> np.ndarray:
    """Preprocess image for RMBG-1.4 model (1024x1024)."""
    target_size = (1024, 1024)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img_rgb, target_size, interpolation=cv2.INTER_LINEAR)
    
    # Normalize: (img / 255.0 - [0.5, 0.5, 0.5]) / [1.0, 1.0, 1.0]
    img_normalized = (img_resized.astype(np.float32) / 255.0 - 0.5)
    
    # Convert to NCHW
    img_tensor = np.transpose(img_normalized, (2, 0, 1))
    img_tensor = np.expand_dims(img_tensor, axis=0)
    return img_tensor

def postprocess_rmbg(mask: np.ndarray, original_shape: tuple) -> np.ndarray:
    """Postprocess mask from RMBG-1.4 to original size."""
    mask = mask.squeeze()
    mask = (mask - mask.min()) / (mask.max() - mask.min())
    mask = (mask * 255).astype(np.uint8)
    mask = cv2.resize(mask, (original_shape[1], original_shape[0]), interpolation=cv2.INTER_LINEAR)
    return mask

def preprocess_depth(img: np.ndarray) -> np.ndarray:
    """Preprocess image for Depth Anything V2 Small (518x518)."""
    target_size = (518, 518)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img_rgb, target_size, interpolation=cv2.INTER_LINEAR)
    
    # ImageNet normalization
    mean = np.array([0.485, 0.456, 0.406]).reshape(1, 1, 3)
    std = np.array([0.229, 0.224, 0.225]).reshape(1, 1, 3)
    img_normalized = (img_resized.astype(np.float32) / 255.0 - mean) / std
    
    # NCHW
    img_tensor = np.transpose(img_normalized, (2, 0, 1))
    img_tensor = np.expand_dims(img_tensor, axis=0)
    return img_tensor

def apply_cinematic_color_grading(img: np.ndarray) -> np.ndarray:
    """Apply Teal and Orange cinematic color grading."""
    # Convert to LAB space
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(np.float32)
    l, a, b = cv2.split(lab)
    
    # Enhance contrast in L channel (HDR feel)
    l = cv2.normalize(l, None, 0, 255, cv2.NORM_MINMAX)
    
    # Teal & Orange effect
    b_shift = np.where(l < 128, b - 10, b + 15)
    b = np.clip(b_shift, 0, 255)
    
    a_shift = np.where(l > 128, a + 5, a)
    a = np.clip(a_shift, 0, 255)
    
    lab_merged = cv2.merge((l, a, b))
    result = cv2.cvtColor(lab_merged.astype(np.uint8), cv2.COLOR_LAB2BGR)
    
    # Subtle vignette
    rows, cols = result.shape[:2]
    X_result = cv2.getGaussianKernel(cols, cols/2)
    Y_result = cv2.getGaussianKernel(rows, rows/2)
    kernel = Y_result * X_result.T
    mask = 255 * kernel / np.linalg.norm(kernel)
    mask = cv2.resize(mask, (cols, rows))
    
    for i in range(3):
        result[:,:,i] = result[:,:,i] * (mask / mask.max()) * 0.8 + result[:,:,i] * 0.2
        
    return result.astype(np.uint8)

def upscale_image(img: np.ndarray) -> np.ndarray:
    """
    Simulated 2x upscaling using OpenCV due to extreme memory constraints.
    We use Lanczos4 interpolation + unsharp masking for crispness.
    """
    height, width = img.shape[:2]
    # Resize 2x
    upscaled = cv2.resize(img, (width * 2, height * 2), interpolation=cv2.INTER_LANCZOS4)
    
    # Unsharp masking for crispness
    gaussian = cv2.GaussianBlur(upscaled, (0, 0), 2.0)
    sharpened = cv2.addWeighted(upscaled, 1.5, gaussian, -0.5, 0)
    
    return sharpened

async def process_image_pipeline(job_id: str, image_bytes: bytes, update_progress_callback, mode: str = "Portrait HD") -> np.ndarray:
    """
    Main image processing pipeline.
    """
    await update_progress_callback(job_id, 5, "Decoding image...")
    # Load image
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    original_shape = img.shape
    
    await update_progress_callback(job_id, 20, "Image Loaded. Starting Person Segmentation...")
    
    # Stage 1: Person Masking (RMBG)
    rmbg_session = MODELS["rmbg"]["session"]
    input_name_rmbg = rmbg_session.get_inputs()[0].name
    rmbg_tensor = preprocess_rmbg(img)
    rmbg_out = rmbg_session.run(None, {input_name_rmbg: rmbg_tensor})[0]
    person_mask = postprocess_rmbg(rmbg_out, original_shape)
    
    await update_progress_callback(job_id, 40, "Subject Masked. Generating Depth Map...")
    
    # Stage 2: Depth Map & Background Blur
    depth_session = MODELS["depth"]["session"]
    input_name_depth = depth_session.get_inputs()[0].name
    depth_tensor = preprocess_depth(img)
    depth_out = depth_session.run(None, {input_name_depth: depth_tensor})[0]
    
    depth_map = depth_out.squeeze()
    depth_map = cv2.resize(depth_map, (original_shape[1], original_shape[0]), interpolation=cv2.INTER_LINEAR)
    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
    
    await update_progress_callback(job_id, 50, "Applying Portrait Depth Blur...")
    
    blurred_bg = cv2.GaussianBlur(img, (51, 51), 0)
    mask_3c = cv2.cvtColor(person_mask, cv2.COLOR_GRAY2BGR).astype(float) / 255.0
    
    portrait_img = (img * mask_3c + blurred_bg * (1 - mask_3c)).astype(np.uint8)
    
    await update_progress_callback(job_id, 60, "Depth Blur Complete. Upscaling and Sharpening...")
    
    # Stage 3: Upscaling
    if mode in ["Ultra Sharp", "Cinematic HD", "Portrait HD"]:
        upscaled_img = upscale_image(portrait_img)
    else:
        upscaled_img = portrait_img
        
    await update_progress_callback(job_id, 80, "Applying Cinematic Color Grading...")
    
    # Stage 4: Color Grading
    if mode == "Cinematic HD":
        final_img = apply_cinematic_color_grading(upscaled_img)
    else:
        final_img = upscaled_img
        
    await update_progress_callback(job_id, 90, "Applying Watermark...")
    
    # Stage 5: Watermark
    final_img = apply_watermark(final_img)
    
    await update_progress_callback(job_id, 100, "Processing Complete!")
    return final_img