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 into memory once at startup
img_colorization = pipeline(
    Tasks.image_colorization,
    model="iic/cv_ddcolor_image-colorization"
)

def colorize_image(img_path: str) -> str:
    """
    Reads a B&W image from disk, runs the colorization model, 
    writes the colorized result to a temp file, and returns its path.
    """
    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:
    """
    Opens a colorized image from disk, applies brightness, contrast, 
    and optional edge enhancement, saves to a temp file, and returns its path.
    """
    image = Image.open(img_path)

    # Adjust brightness
    image = ImageEnhance.Brightness(image).enhance(brightness)
    # Adjust contrast
    image = ImageEnhance.Contrast(image).enhance(contrast)
    # Optionally apply an edge enhancement filter
    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
):
    """
    1) Colorizes the uploaded B&W image.
    2) Applies the chosen brightness/contrast/edge-enhancement.
    3) Re‐saves in the user’s chosen format (PNG/JPEG/TIFF).
    Returns:
      - A list [original_path, final_path] for side-by-side display.
      - The final image’s file path for download.
    """
    # Step 1: colorize
    colorized_path = colorize_image(img_path)
    # Step 2: enhancement
    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 ([original, enhanced], download_path)
    return ([img_path, enhanced_path], output_path)

# Title and description shown at the top of the interface
TITLE = "🌈 Color Restorization Model"
DESCRIPTION = "Upload a black & white photo to restore it in color using a deep learning model."

# Build the Gradio Blocks interface
with gr.Blocks(title=TITLE) as app:
    gr.Markdown(f"## {TITLE}")
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                type="filepath",
                label="Upload B&W Image",
                tool="editor"       # Enables zoom/pan on the uploaded image
            )
            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")

        with gr.Column():
            comparison_gallery = gr.Gallery(
                label="Original vs Colorized",
                columns=2,      # two images side by side
                height="auto"
            )
            download_btn = gr.File(label="Download Colorized Image")

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

# Launch in “production” mode: bind to 0.0.0.0 on configurable port
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.queue().launch(server_name="0.0.0.0", server_port=port)