import os
import warnings
import logging

# ----------------------------
# 1. Warning & logging setup
# ----------------------------
# Suppress FutureWarning from timm internals
warnings.filterwarnings(
    "ignore",
    category=FutureWarning,
    module="timm.models.layers"
)
# Suppress UserWarning from modelscope (e.g. missing preprocessor config)
warnings.filterwarnings(
    "ignore",
    category=UserWarning,
    module="modelscope"
)
# Only show ERROR+ logs from modelscope
logging.getLogger("modelscope").setLevel(logging.ERROR)

# ----------------------------
# 2. Standard imports
# ----------------------------
import cv2
import tempfile
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

# ----------------------------
# 3. Load your colorization model
# ----------------------------
img_colorization = pipeline(
    Tasks.image_colorization,
    model="iic/cv_ddcolor_image-colorization",
    model_revision="v1.02",          # explicitly specify revision
)

# ----------------------------
# 4. Image processing fns
# ----------------------------
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)
    image = ImageEnhance.Brightness(image).enhance(brightness)
    image = ImageEnhance.Contrast(image).enhance(contrast)
    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
):
    # Colorize → Enhance → Re‑save in chosen format
    colorized_path = colorize_image(img_path)
    enhanced_path = enhance_image(colorized_path, brightness, contrast, edge_enhance)

    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 gallery and downloadable file
    return ([img_path, enhanced_path], output_path)

# ----------------------------
# 5. Gradio UI + custom CSS
# ----------------------------
custom_css = """
body { background-color: #f0f2f5; }
.gradio-container { max-width: 900px !important; margin: auto !important; }
#header { background-color: #4CAF50; padding: 20px; border-radius: 8px;
         text-align: center; margin-bottom: 20px; }
#header h2, #header p { color: white; margin: 0; }
#header p { margin-top: 5px; font-size: 1rem; }
#control-panel { background: white; padding: 20px; border-radius: 8px;
                 box-shadow: 0 2px 8px rgba(0,0,0,0.1); margin-bottom: 20px; }
#submit-btn { background-color: #4CAF50 !important; color: white !important;
              border-radius: 8px !important; font-weight: bold;
              padding: 10px 20px !important; margin-top: 10px !important; }
#control-panel .gr-row { gap: 15px; }
.gr-slider, .gr-checkbox, .gr-dropdown { margin-top: 10px; }
#comparison_gallery { background: white; padding: 10px;
                      border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.1); }
#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
    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>
        """
    )

    # Controls & results
    with gr.Column(elem_id="control-panel"):
        with gr.Row():
            # Inputs
            with gr.Column():
                input_image = gr.Image(type="filepath", label="Upload B&W Image", interactive=True)
                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(["PNG", "JPEG", "TIFF"], value="PNG", label="Output Format")
                submit_btn = gr.Button("Colorize", elem_id="submit-btn")

            # Outputs
            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")

    # Wire up
    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
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.queue().launch(server_name="0.0.0.0", server_port=port)