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

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

def colorize_image(img_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, brightness=1.0, contrast=1.0, edge_enhance=False):
    image = Image.open(img_path)
    # Adjust brightness
    image = ImageEnhance.Brightness(image).enhance(brightness)
    # Adjust contrast
    image = ImageEnhance.Contrast(image).enhance(contrast)
    # Optionally apply 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, brightness, contrast, edge_enhance, output_format):
    # 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) plus the single downloadable file
    return ([img_path, enhanced_path], output_path)

title = "🌈 Color Restorization Model"
description = "Upload a black & white photo to restore it in color using a deep learning model."

with gr.Blocks(title=title) as demo:
    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")
            brightness_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
            contrast_slider = gr.Slider(0.5, 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, 
                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]
    )

demo.launch(enable_queue=True)