import os
import tempfile
from pathlib import Path

import gradio as gr
import cv2
import numpy as np
from PIL import Image, ImageCms

# =========================================================
# Color profile assets (Gradio Space safe)
# =========================================================
BASE_DIR = Path(__file__).resolve().parent
ICC_DIR = BASE_DIR / "icc"

ADOBE_ICC_PATH = ICC_DIR / "AdobeRGB1998.icc"

# sRGB profile is always available via Pillow (no file needed)
SRGB_PROFILE = ImageCms.createProfile("sRGB")
SRGB_ICC_BYTES = ImageCms.ImageCmsProfile(SRGB_PROFILE).tobytes()

def safe_open_profile(path: Path):
    """
    Safely open ICC profile.
    AdobeRGB1998 can legitimately be small (e.g., ~560 bytes), so do NOT size-gate.
    """
    try:
        if not path.exists():
            print(f"⚠️ ICC missing: {path}")
            return None
        if not path.is_file():
            print(f"⚠️ ICC is not a file: {path}")
            return None

        size = path.stat().st_size
        print(f"ℹ️ ICC found: {path} ({size} bytes)")

        prof = ImageCms.getOpenProfile(str(path))
        return prof
    except Exception as e:
        print(f"⚠️ ICC open failed: {path} → {e}")
        return None

ADOBE_PROFILE = safe_open_profile(ADOBE_ICC_PATH)

def convert_and_tag(img: Image.Image, target_color_space: str):
    """
    Converts pixel values to the selected target profile AND returns ICC bytes
    for embedding into the exported file.
    Incoming images are treated as sRGB by default (practical for AI/web images).
    """
    if img.mode != "RGB":
        img = img.convert("RGB")

    if target_color_space == "sRGB (Stock Standard)":
        return img, SRGB_ICC_BYTES

    if target_color_space == "Adobe RGB (1998)":
        if ADOBE_PROFILE is None:
            print("⚠️ AdobeRGB requested but profile not loaded; falling back to sRGB tagging.")
            return img, SRGB_ICC_BYTES

        try:
            xform = ImageCms.buildTransformFromOpenProfiles(
                SRGB_PROFILE,
                ADOBE_PROFILE,
                "RGB",
                "RGB",
                renderingIntent=0  # perceptual
            )
            out = ImageCms.applyTransform(img, xform)
            return out, ADOBE_PROFILE.tobytes()
        except Exception as e:
            print(f"⚠️ AdobeRGB transform failed; falling back to sRGB. Error: {e}")
            return img, SRGB_ICC_BYTES

    return img, SRGB_ICC_BYTES


# =========================================================
# 1) THE PROCESSING ENGINE
# =========================================================
def process_image(image, output_format, color_space, fix_lighting, add_grain):
    if image is None:
        return None, None

    # Convert PIL -> numpy
    img_array = np.array(image)

    # Stock safety: drop alpha if present
    if img_array.ndim == 3 and img_array.shape[2] == 4:
        img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)

    # RGB -> BGR for OpenCV
    img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

    # A) LIGHTING FIX (CLAHE) - gentle
    if fix_lighting:
        lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=1.5, tileGridSize=(8, 8))
        l_fixed = clahe.apply(l)
        lab_fixed = cv2.merge((l_fixed, a, b))
        img_bgr = cv2.cvtColor(lab_fixed, cv2.COLOR_LAB2BGR)

    # B) TEXTURE FIX (Monochromatic Film Grain)
    if add_grain:
        h, w = img_bgr.shape[:2]
        # IMPORTANT: int16 prevents negative noise from wrapping under uint8
        noise = np.random.normal(loc=0.0, scale=2.5, size=(h, w)).astype(np.int16)
        img_i16 = img_bgr.astype(np.int16)
        img_i16[:, :, 0] += noise
        img_i16[:, :, 1] += noise
        img_i16[:, :, 2] += noise
        img_bgr = np.clip(img_i16, 0, 255).astype(np.uint8)

    # Convert back to PIL RGB
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    final_pil = Image.fromarray(img_rgb).convert("RGB")

    # C) COLOR SPACE: convert + embed ICC
    final_pil, icc_bytes = convert_and_tag(final_pil, color_space)

    # D) EXPORT: create a real file for download
    suffix = ".jpg" if output_format == "JPEG" else ".png"
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    tmp_path = tmp.name
    tmp.close()

    if output_format == "JPEG":
        final_pil.save(
            tmp_path,
            format="JPEG",
            quality=100,
            subsampling=0,
            optimize=True,
            icc_profile=icc_bytes,
        )
    else:
        final_pil.save(
            tmp_path,
            format="PNG",
            compress_level=6,
            optimize=True,
            icc_profile=icc_bytes,
        )

    return final_pil, tmp_path


# =========================================================
# 2) THE UI (Gradio)
# =========================================================
css = """
#run-btn {background-color: #ff7c00 !important; color: white !important;}
"""

with gr.Blocks(title="StockFix AI", css=css) as app:
    gr.Markdown("## StockFix AI: De-Plasticizer")

    # Friendly status line
    if ADOBE_PROFILE is None:
        gr.Markdown(
            "⚠️ **Adobe RGB (1998) ICC not loaded.** "
            "Upload `icc/AdobeRGB1998.icc` to enable true Adobe RGB exports."
        )
    else:
        gr.Markdown("✅ Adobe RGB (1998) profile loaded.")

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input AI Image")

            with gr.Group():
                gr.Markdown("### 1. Fixes")
                chk_light = gr.Checkbox(label="Fix Flat Lighting (CLAHE)", value=True)
                chk_grain = gr.Checkbox(label="Add Film Grain (Monochromatic)", value=True)

            with gr.Group():
                gr.Markdown("### 2. Color Profile")
                radio_color = gr.Radio(
                    ["sRGB (Stock Standard)", "Adobe RGB (1998)"],
                    label="Target Color Space",
                    value="sRGB (Stock Standard)",
                )

            with gr.Group():
                gr.Markdown("### 3. Format")
                radio_fmt = gr.Radio(["JPEG", "PNG"], label="Output Format", value="JPEG")

            btn_run = gr.Button("Fix & Export", elem_id="run-btn")

        with gr.Column():
            output_img = gr.Image(label="Preview", type="pil")
            output_file = gr.File(label="Download (PNG/JPEG)")

    btn_run.click(
        fn=process_image,
        inputs=[input_img, radio_fmt, radio_color, chk_light, chk_grain],
        outputs=[output_img, output_file],
    )

if __name__ == "__main__":
    app.launch()