import os
import tempfile
import gradio as gr
from PIL import Image
from core import Colorizer

# Initialize global colorizer
colorizer = Colorizer()

def process_image(
    img_path: str,
    brightness: float,
    contrast: float,
    edge_enhance: bool,
    output_format: str,
    quality: str,
    progress=gr.Progress()
):
    if img_path is None:
        return None, None

    progress(0, desc="Loading image...")
    # Load input
    try:
        img = Image.open(img_path).convert("RGB")
    except Exception as e:
        print(f"Error loading image: {e}")
        return None, None

    # Map quality to resolution
    quality_map = {
        "Fast (256px)": 256,
        "Balanced (512px)": 512,
        "High (1080px)": 1080,
        "Original": 0
    }
    res = quality_map.get(quality, 512)

    progress(0.3, desc="Colorizing & Enhancing...")
    # Process using Core Logic (In-Memory)
    enhanced_img = colorizer.process(
        img,
        brightness=brightness,
        contrast=contrast,
        edge_enhance=edge_enhance,
        adaptive_resolution=res
    )

    progress(0.9, desc="Saving outputs...")
    # Save outputs for Gradio
    # 1. Enhanced image for gallery
    temp_dir = tempfile.mkdtemp()
    enhanced_path = os.path.join(temp_dir, "enhanced.png")
    enhanced_img.save(enhanced_path)

    # 2. Downloadable file
    filename = f"colorized_image.{output_format.lower()}"
    output_path = os.path.join(temp_dir, filename)
    enhanced_img.save(output_path, format=output_format.upper())

    progress(1.0, desc="Done!")
    # Return side-by-side (original, enhanced) and the downloadable file
    return ([img_path, enhanced_path], output_path)

# CSS to give a modern, centered layout with a colored header and clean panels
custom_css = """
/* Overall background */
body { 
    background-color: #f0f2f5; 
}

/* Center the Gradio container and give it a max width */
.gradio-container {
    max-width: 900px !important;
    margin: auto !important;
}

/* Header styling */
#header {
    background-color: #4CAF50;
    padding: 20px;
    border-radius: 8px;
    text-align: center;
    margin-bottom: 20px;
}
#header h2 {
    color: white;
    margin: 0;
    font-size: 2rem;
}
#header p {
    color: white;
    margin: 5px 0 0 0;
    font-size: 1rem;
}

/* White panel for controls */
#control-panel {
    background-color: white;
    padding: 20px;
    border-radius: 8px;
    box-shadow: 0px 2px 8px rgba(0,0,0,0.1);
    margin-bottom: 20px;
}

/* Style the “Colorize” button */
#submit-btn {
    background-color: #4CAF50 !important;
    color: white !important;
    border-radius: 8px !important;
    font-weight: bold;
    padding: 10px 20px !important;
    margin-top: 10px !important;
}

/* Add some spacing around sliders and checkbox */
#control-panel .gr-row {
    gap: 15px;
}
.gr-slider, .gr-checkbox, .gr-dropdown {
    margin-top: 10px;
}

/* Gallery panel styling */
#comparison_gallery {
    background-color: white;
    padding: 10px;
    border-radius: 8px;
    box-shadow: 0px 2px 8px rgba(0,0,0,0.1);
}

/* Download button spacing */
#download-btn {
    margin-top: 15px !important;
}
"""

TITLE = "🌈 Color Restorization Model"
DESCRIPTION = "Bring your old black & white photos back to life—upload, adjust, and download in vivid color."

with gr.Blocks(title=TITLE) as app:
    # Header section
    gr.HTML(
        """
        <div id="header">
            <h2>🌈 Color Restorization Model</h2>
            <p>Bring your old black & white photos back to life—upload, adjust, and download in vivid color.</p>
        </div>
        """
    )

    # Main control panel: white box with rounded corners
    with gr.Column(elem_id="control-panel"):
        with gr.Row():
            # Left column: inputs and controls
            with gr.Column():
                input_image = gr.Image(
                    type="filepath",
                    label="Upload B&W Image",
                    interactive=True
                )
                brightness_slider = gr.Slider(
                    minimum=0.5, maximum=2.0, value=1.0,
                    label="Brightness"
                )
                contrast_slider = gr.Slider(
                    minimum=0.5, maximum=2.0, value=1.0,
                    label="Contrast"
                )
                edge_enhance_checkbox = gr.Checkbox(
                    label="Apply Edge Enhancement"
                )
                quality_dropdown = gr.Dropdown(
                    choices=["Fast (256px)", "Balanced (512px)", "High (1080px)", "Original"],
                    value="Balanced (512px)",
                    label="Processing Quality (Resolution)"
                )
                output_format_dropdown = gr.Dropdown(
                    choices=["PNG", "JPEG", "TIFF"],
                    value="PNG",
                    label="Output Format"
                )
                submit_btn = gr.Button(
                    "Colorize", 
                    elem_id="submit-btn"
                )

            # Right column: results gallery & download
            with gr.Column():
                comparison_gallery = gr.Gallery(
                    label="Original vs. Colorized",
                    columns=2,
                    elem_id="comparison_gallery",
                    height="auto"
                )
                download_btn = gr.File(
                    label="Download Colorized Image",
                    elem_id="download-btn"
                )

    submit_btn.click(
        fn=process_image,
        inputs=[
            input_image,
            brightness_slider,
            contrast_slider,
            edge_enhance_checkbox,
            output_format_dropdown,
            quality_dropdown
        ],
        outputs=[comparison_gallery, download_btn]
    )

# “Production” launch: bind to 0.0.0.0 and use PORT env var if provided
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.queue().launch(server_name="0.0.0.0", server_port=port)