import gradio as gr
import tempfile
import os
from simple_mosaic import SimpleMosaicImage
from tile_library import build_cifar10_tile_library, build_cifar100_tile_library
from performance import calculate_all_metrics, format_metrics_report

MAX_IMAGE_SIZE = 4500
# Global cache for tile libraries to avoid reloading
tile_cache = {}

def get_tile_library(tile_type, max_per_class):
    """Get cached tile library or create new one"""
    cache_key = f"{tile_type}_{max_per_class}"

    if cache_key not in tile_cache:
        if tile_type == "CIFAR-10":
            tiles, means, labels = build_cifar10_tile_library(max_per_class=max_per_class)
        elif tile_type == "CIFAR-100":
            tiles, means, labels = build_cifar100_tile_library(max_per_class=max_per_class)
        else:
            return None, None, None

        tile_cache[cache_key] = (tiles, means, labels)

    return tile_cache[cache_key]

def process_mosaic(image, start_size, min_size, threshold, tile_type, max_per_class, quantize_colors):
    """Process image to create mosaic"""
    if image is None:
        return None, "Please upload an image first."

    try:
        # Create temporary file for processing
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
            image.save(tmp_file.name)

            # Load and process image
            loader = SimpleMosaicImage(tmp_file.name)

            # Apply Resize if needed
            if max(loader.width, loader.height) > MAX_IMAGE_SIZE:
                loader.resize(MAX_IMAGE_SIZE);

            # Apply color quantization if requested
            if quantize_colors > 0:
                loader.quantize_colors(quantize_colors)

            # Smart boundary handling
            loader.crop_to_grid(2)

            # Build adaptive cells
            cells = loader.build_adaptive_cells(
                start_size=int(start_size),
                min_size=int(min_size),
                threshold=float(threshold)
            )

            # Generate mosaic based on tile type
            if tile_type == "None (Average Colors)":
                result_img = loader.mosaic_average_color_adaptive(cells)
                basic_info = f"Generated mosaic with {len(cells)} adaptive cells using average colors."
            else:
                tiles, tile_means, _ = get_tile_library(tile_type, int(max_per_class))
                if tiles is None:
                    return None, f"Failed to load {tile_type} tile library."

                result_img = loader.mosaic_with_tiles_adaptive(cells, tiles, tile_means)
                basic_info = f"Generated mosaic with {len(cells)} cells using {len(tiles)} {tile_type} tiles."

            # Calculate performance metrics
            try:
                metrics = calculate_all_metrics(image, result_img)
                metrics_report = format_metrics_report(metrics)
                info = f"{basic_info}\n\n{metrics_report}"
            except Exception as e:
                info = f"{basic_info}\n\nMetrics calculation failed: {str(e)}"

            # Clean up temporary file
            os.unlink(tmp_file.name)

            return result_img, info

    except Exception as e:
        return None, f"Error processing image: {str(e)}"

def update_max_per_class_visibility(tile_type):
    """Show/hide max_per_class slider based on tile type"""
    if tile_type in ["CIFAR-10", "CIFAR-100"]:
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)

def get_max_per_class_range(tile_type):
    """Get appropriate range for max_per_class slider"""
    if tile_type == "CIFAR-10":
        return gr.update(maximum=1000, value=500)
    elif tile_type == "CIFAR-100":
        return gr.update(maximum=400, value=200)
    else:
        return gr.update(maximum=1000, value=500)

# Simple CSS - white background, black text
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600&display=swap');

/* White background everywhere */
* {
    background: #ffffff !important;
    background-color: #ffffff !important;
    font-family: 'Inter', sans-serif !important;
}

/* Black text everywhere */
* {
    color: #000000 !important;
}

/* Button exception - keep dark */
.gr-button-primary {
    background: #000000 !important;
    background-color: #000000 !important;
    color: #ffffff !important;
    border: 1px solid #000000 !important;
    border-radius: 6px !important;
    padding: 8px 16px !important;
}

.gr-button-primary:hover {
    background: #333333 !important;
    background-color: #333333 !important;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, title="Interactive Image Mosaic Generator") as demo:
    gr.Markdown("# Interactive Image Mosaic Generator")
    gr.Markdown("Transform your images into mosaics with adaptive grid technology")

    with gr.Row():
        # Left panel - Controls
        with gr.Column(scale=1, min_width=300, elem_classes=["control-panel"]):
            gr.Markdown("### Controls")

            # Image upload
            input_image = gr.Image(
                label="Upload Image",
                type="pil",
                height=200,
                interactive=True
            )

            # Grid parameters
            gr.Markdown("**Grid Parameters**")
            start_size = gr.Slider(
                minimum=16, maximum=128, value=64, step=8,
                label="Initial Block Size",
                info="Starting grid size (larger = fewer details)"
            )

            min_size = gr.Slider(
                minimum=2, maximum=32, value=4, step=2,
                label="Minimum Block Size",
                info="Smallest allowed grid size"
            )

            threshold = gr.Slider(
                minimum=0.1, maximum=20.0, value=5.0, step=0.1,
                label="Subdivision Threshold",
                info="Lower values = more subdivision"
            )

            # Color quantization
            quantize_colors = gr.Slider(
                minimum=0, maximum=128, value=16, step=1,
                label="Color Quantization",
                info="Number of colors (0 = no quantization)"
            )

            # Tile library selection
            gr.Markdown("**Tile Library**")
            tile_type = gr.Radio(
                choices=["None (Average Colors)", "CIFAR-10", "CIFAR-100"],
                value="CIFAR-10",
                label="Tile Type",
                info="Choose tile source or use average colors"
            )

            max_per_class = gr.Slider(
                minimum=10, maximum=1000, value=500, step=10,
                label="Tiles per Class",
                info="Number of tiles per category",
                visible=True
            )

            # Process button
            process_btn = gr.Button(
                "Generate Mosaic",
                variant="primary",
                size="lg"
            )

        # Right panel - Images
        with gr.Column(scale=2, elem_classes=["image-display"]):
            gr.Markdown("### Results")

            with gr.Row():
                # Original image preview
                with gr.Column():
                    gr.Markdown("**Original**")
                    original_display = gr.Image(
                        label="Original Image",
                        height=400,
                        interactive=False
                    )

                # Result image
                with gr.Column():
                    gr.Markdown("**Mosaic Result**")
                    result_image = gr.Image(
                        label="Mosaic Image",
                        height=400,
                        interactive=False
                    )

            # Info output
            info_output = gr.Textbox(
                label="Processing Info & Performance Metrics",
                interactive=False,
                max_lines=10,
                lines=8
            )

    # Event handlers
    tile_type.change(
        fn=update_max_per_class_visibility,
        inputs=[tile_type],
        outputs=[max_per_class]
    )

    tile_type.change(
        fn=get_max_per_class_range,
        inputs=[tile_type],
        outputs=[max_per_class]
    )

    input_image.change(
        fn=lambda img: img,
        inputs=[input_image],
        outputs=[original_display]
    )

    process_btn.click(
        fn=process_mosaic,
        inputs=[
            input_image, start_size, min_size, threshold,
            tile_type, max_per_class, quantize_colors
        ],
        outputs=[result_image, info_output]
    )

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