import gradio as gr
import numpy as np
from PIL import Image
import time
import os
import sys

# Ensure the pixelfixer package can be imported
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from pixelfixer import detect
from pixelfixer.cli import _overlay, _extract
import spaces

import tempfile

@spaces.GPU
def process_image(input_image_path, color_limit, export_scales, force_step_val):
    if input_image_path is None:
        return None, None, "Please upload an image.", None, None
        
    # Parse settings
    k_colors = 0
    if color_limit != "Unlimited":
        k_colors = int(color_limit)
        
    force_step = None
    if force_step_val > 0:
        force_step = float(force_step_val)
        
    original_name = os.path.basename(input_image_path)
    base, _ = os.path.splitext(original_name)
    if base.startswith("image") and len(base) > 5:
        base = "image"
        
    from pixelfixer.api import process
    
    with open(input_image_path, "rb") as f:
        img_bytes = f.read()
        
    try:
        result = process(
            img_bytes, 
            mode="full", 
            two_stage=True, 
            k_colors=k_colors, 
            force_step=force_step,
            return_png=False
        )
    except Exception as e:
        return None, None, f"Error processing image: {str(e)}", None, None
        
    tmp_dir = tempfile.mkdtemp()
    low_res_array = result["array"]
    
    if k_colors > 0:
        from pixelfixer.quantize import kmeans_quantize
        low_res_array, _, _ = kmeans_quantize(low_res_array, k=k_colors)
        
    base_img = Image.fromarray(low_res_array)
    
    # 1. Create a Crisp Preview Image (Upscaled heavily for UI viewing)
    preview_scale = max(1, 512 // max(result["cols"], result["rows"]))
    preview_w = result["cols"] * preview_scale
    preview_h = result["rows"] * preview_scale
    preview_img = base_img.resize((preview_w, preview_h), Image.NEAREST)
    preview_path = os.path.join(tmp_dir, f"{base}-preview.png")
    preview_img.save(preview_path, format="PNG")
    
    # 2. Generate Export Files based on user selection
    export_paths = []
    if not export_scales:
        export_scales = ["1x"] # Fallback if nothing selected
        
    for scale_str in export_scales:
        scale_mult = int(scale_str.replace("x", ""))
        target_w = result["cols"] * scale_mult
        target_h = result["rows"] * scale_mult
        export_img = base_img.resize((target_w, target_h), Image.NEAREST)
        export_path = os.path.join(tmp_dir, f"{base}-{scale_str}.png")
        export_img.save(export_path, format="PNG")
        export_paths.append(export_path)
    
    # 3. Generate Overlay
    try:
        input_img = Image.open(input_image_path).convert("RGBA")
        rgba = np.array(input_img)
        overlay_img = Image.fromarray(_overlay(rgba, result))
        overlay_path = os.path.join(tmp_dir, f"{base}-overlay.png")
        overlay_img.save(overlay_path, format="PNG")
    except Exception:
        overlay_path = None
    
    # 4. Info String
    info = f"Grid detected: {result['cols']}x{result['rows']}\n"
    info += f"Step size: {result['step_x']}x{result['step_y']}\n"
    info += f"Confidence: {result['confidence'].upper()} ({result['consensus']})\n"
    if k_colors > 0:
        info += f"Colors: {k_colors}\n"
    info += f"Detect: {result['detect_s']}s | Recon: {result['recon_s']}s"
    
    return preview_path, export_paths, info, overlay_path

with gr.Blocks(title="Pixel Art Fixer", css=".pixelated-image { image-rendering: pixelated; }") as app:
    gr.Markdown("# Pixel Art Fixer")
    gr.Markdown("Convert AI-generated 'fake' pixel art into real, grid-aligned pixel art.")
    
    with gr.Row():
        input_image = gr.Image(type="filepath", label="Input Image", elem_classes="pixelated-image")
        preview_image = gr.Image(type="filepath", label="Crisp Preview", elem_classes="pixelated-image")
            
    with gr.Row():
        with gr.Column():
            color_limit = gr.Radio(choices=["Unlimited", "8", "12", "16", "32"], value="Unlimited", label="Color Palette Limit")
            export_scales = gr.CheckboxGroup(choices=["1x", "2x", "5x", "10x"], value=["1x", "5x"], label="Export Scales (Select Multiple)")
            force_step = gr.Number(label="Force Block Size (px) - Set to 0 for Auto", value=0, minimum=0, precision=1)
            process_btn = gr.Button("Process Image", variant="primary")
            
        with gr.Column():
            export_files = gr.File(label="Download Extracted Pixel Art", file_count="multiple")
            info_text = gr.Textbox(label="Detection Info", lines=5)
            
        with gr.Column():
            overlay_image = gr.Image(type="filepath", label="Grid Overlay")
            
    process_btn.click(
        fn=process_image,
        inputs=[input_image, color_limit, export_scales, force_step],
        outputs=[preview_image, export_files, info_text, overlay_image]
    )

if __name__ == "__main__":
    app.launch()