import os, cv2, numpy as np, onnxruntime as ort, gradio as gr
from huggingface_hub import hf_hub_download

# --- MODEL DIRECTORY ---
# Friendly labels for "noobs" mapped to technical repos
MODEL_OPTIONS = {
    "💎 Pro Detail (APISR RRDB)": {
        "repo": "Xenova/2x_APISR_RRDB_GAN_generator-onnx",
        "desc": "Best for AI art and textures. High-quality reconstruction.",
        "scale": 2
    },
    "⚡ Fast & Sharp (APISR GRL)": {
        "repo": "Xenova/4x_APISR_GRL_GAN_generator-onnx",
        "desc": "Best for architecture and sharp lines. Very clean results.",
        "scale": 4
    },
    "📸 Realistic Photos (Swin2SR)": {
        "repo": "Xenova/swin2SR-realworld-sr-x4-64-bsrgan-psnr",
        "desc": "Best for real-life photography. Keeps skin and nature natural.",
        "scale": 4
    },
    "🧹 Noise Cleanup (Swin2SR Compressed)": {
        "repo": "Xenova/swin2SR-compressed-sr-x4-48",
        "desc": "Best for blurry, low-quality internet images with 'blocks'.",
        "scale": 4
    }
}

# Global dictionary to cache loaded models
loaded_sessions = {}

def get_session(model_key):
    if model_key not in loaded_sessions:
        repo = MODEL_OPTIONS[model_key]["repo"]
        path = hf_hub_download(repo_id=repo, filename="onnx/model.onnx")
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = 2
        loaded_sessions[model_key] = ort.InferenceSession(path, opts, providers=['CPUExecutionProvider'])
    return loaded_sessions[model_key]

def upscale_image_tiled(frame, model_key, tile_size=128, overlap=16):
    h, w, c = frame.shape
    scale = MODEL_OPTIONS[model_key]["scale"]
    session = get_session(model_key)
    
    output_h, output_w = h * scale, w * scale
    upscaled_img = np.zeros((output_h, output_w, c), dtype=np.uint8)
    
    stride = tile_size - (overlap * 2)
    
    for y in range(0, h, stride):
        for x in range(0, w, stride):
            y1, y2 = max(0, y - overlap), min(h, y + stride + overlap)
            x1, x2 = max(0, x - overlap), min(w, x + stride + overlap)
            
            tile = frame[y1:y2, x1:x2]
            
            # Pad to multiple of 8 for APISR/Swin2SR compatibility
            th, tw = tile.shape[:2]
            pad_h = (8 - (th % 8)) % 8
            pad_w = (8 - (tw % 8)) % 8
            if pad_h > 0 or pad_w > 0:
                tile = cv2.copyMakeBorder(tile, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT)
            
            # AI Inference
            img_input = cv2.cvtColor(tile, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
            img_input = np.transpose(img_input, (2, 0, 1))[np.newaxis, :]
            
            output = session.run(None, {session.get_inputs()[0].name: img_input})[0]
            
            # Post-process tile
            tile_out = np.clip(np.squeeze(output), 0, 1).transpose(1, 2, 0)
            tile_out = cv2.cvtColor((tile_out * 255.0).astype(np.uint8), cv2.COLOR_RGB2BGR)
            
            # Remove padding and overlap
            tile_out = tile_out[:(th*scale), :(tw*scale)]
            oy1, ox1 = (y - y1) * scale, (x - x1) * scale
            py1, py2 = y * scale, min(output_h, (y + stride) * scale)
            px1, px2 = x * scale, min(output_w, (x + stride) * scale)
            
            upscaled_img[py1:py2, px1:px2] = tile_out[oy1 : oy1 + (py2-py1), ox1 : ox1 + (px2-px1)]
            
    return upscaled_img

def run_universal(img_data, model_choice, sharpen):
    if img_data is None: return None
    img = img_data["composite"]
    if img.shape[2] == 4: img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
    
    res = upscale_image_tiled(img, model_choice)
    
    if sharpen > 0:
        blurred = cv2.GaussianBlur(res, (0, 0), 3)
        res = cv2.addWeighted(res, 1 + sharpen, blurred, -sharpen, 0)
        
    return res

# --- UI DESIGN ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("# 🌌 Universal AI Image Enhancer")
    gr.Markdown("Pick a goal for your image and let the AI handle the rest.")
    
    with gr.Row():
        with gr.Column():
            image_in = gr.ImageEditor(label="Upload Image", type="numpy")
            
            # Friendly Selector
            model_dropdown = gr.Dropdown(
                choices=list(MODEL_OPTIONS.keys()), 
                value="💎 Pro Detail (APISR RRDB)", 
                label="What is your goal?"
            )
            
            sharp_slider = gr.Slider(0, 0.5, value=0.15, label="Sharpness Boost")
            submit_btn = gr.Button("🚀 ENHANCE IMAGE", variant="primary")
            
        with gr.Column():
            image_out = gr.Image(label="Upscaled Result")

    submit_btn.click(run_universal, [image_in, model_dropdown, sharp_slider], image_out)

if __name__ == "__main__":
    demo.queue().launch()