import os
import cv2
import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

# Load the colorization model at startup
img_colorization = pipeline(
    Tasks.image_colorization,
    model="iic/cv_ddcolor_image-colorization"
)

def colorize_image(img_path: str) -> str:
    image = cv2.imread(str(img_path))
    output = img_colorization(image[..., ::-1])
    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)

    temp_dir = tempfile.mkdtemp()
    out_path = os.path.join(temp_dir, "colorized.png")
    cv2.imwrite(out_path, result)
    return out_path

def enhance_image(
    img_path: str,
    brightness: float = 1.0,
    contrast: float = 1.0,
    edge_enhance: bool = False
) -> str:
    image = Image.open(img_path)

    # Brightness adjustment
    image = ImageEnhance.Brightness(image).enhance(brightness)
    # Contrast adjustment
    image = ImageEnhance.Contrast(image).enhance(contrast)
    # Optional edge enhancement
    if edge_enhance:
        image = image.filter(ImageFilter.EDGE_ENHANCE)

    temp_dir = tempfile.mkdtemp()
    enhanced_path = os.path.join(temp_dir, "enhanced.png")
    image.save(enhanced_path)
    return enhanced_path

def process_image(
    img_path: str,
    brightness: float,
    contrast: float,
    edge_enhance: bool,
    output_format: str
):
    # Step 1: Colorize
    colorized_path = colorize_image(img_path)
    # Step 2: Enhance (brightness/contrast/edge)
    enhanced_path = enhance_image(colorized_path, brightness, contrast, edge_enhance)
    # Step 3: Convert to chosen format
    img = Image.open(enhanced_path)
    temp_dir = tempfile.mkdtemp()
    filename = f"colorized_image.{output_format.lower()}"
    output_path = os.path.join(temp_dir, filename)
    img.save(output_path, format=output_format.upper())

    # 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, css=custom_css) 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",
                    tool="editor",          # Enables zoom/pan
                    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"
                )
                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
        ],
        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)