# ================================================================
# 1. ZeroGPU REQUIRES `spaces` TO BE THE FIRST IMPORT
# ================================================================
import spaces  # MUST be first

import os
import gc
import sys
import math
import time
import types
import cv2
import urllib.request
import shutil
import tempfile
import warnings
import torch
import numpy as np
import gradio as gr
from PIL import Image, ImageFilter, ImageDraw, ImageEnhance

# Ignore non-critical warnings
warnings.filterwarnings("ignore", category=ResourceWarning)
warnings.filterwarnings("ignore", category=UserWarning)

# Torchvision shim for older dependency paths
try:
    import torchvision.transforms.functional_tensor  # noqa: F401
except ModuleNotFoundError:
    import torchvision.transforms.functional as _F
    _shim = types.ModuleType("torchvision.transforms.functional_tensor")
    _shim.rgb_to_grayscale = _F.rgb_to_grayscale
    sys.modules["torchvision.transforms.functional_tensor"] = _shim

from torchvision.transforms.functional import normalize
from diffusers import FluxFillPipeline, FluxTransformer2DModel, FluxKontextPipeline
from transformers import T5EncoderModel, CLIPModel, CLIPProcessor

from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from gfpgan import GFPGANer

# Optional quality metrics imports
try:
    from piq import brisque
except ImportError:
    brisque = None


# ================================================================
# 2. LOAD FLUX.1-Fill-dev NF4
# ================================================================
try:
    from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer

    def _check_quantized_param_shape(self, param_name, current_param, loaded_param):
        n = math.prod(tuple(current_param.shape))
        inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)
        if tuple(loaded_param.shape) != tuple(inferred_shape):
            raise ValueError(
                f"Expected flattened shape of ({param_name}) "
                f"to be {tuple(loaded_param.shape)} but got {tuple(inferred_shape)}."
            )
        return True

    BnB4BitDiffusersQuantizer.check_quantized_param_shape = _check_quantized_param_shape
    print("✅ bnb shape-check shim applied.")
except Exception as _exc:
    print(f"[startup] bnb shape-check shim not applied: {_exc!r}")

print("Loading FLUX.1-Fill-dev NF4...")
base_id = "black-forest-labs/FLUX.1-Fill-dev"
nf4_id  = "sayakpaul/FLUX.1-Fill-dev-nf4"

transformer = FluxTransformer2DModel.from_pretrained(
    nf4_id, subfolder="transformer", torch_dtype=torch.bfloat16
)
text_encoder_2 = T5EncoderModel.from_pretrained(
    nf4_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16
)
pipe = FluxFillPipeline.from_pretrained(
    base_id,
    transformer=transformer,
    text_encoder_2=text_encoder_2,
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")
pipe.vae.enable_slicing()
pipe.vae.enable_tiling()
print("FLUX.1-Fill loaded.")


# ================================================================
# 3. LOAD GFPGAN & REAL-ESRGAN (Photo & Face Restoration)
# ================================================================
print("Setting up GFPGAN & Real-ESRGAN...")
weights_dir = "weights"
os.makedirs(weights_dir, exist_ok=True)

# Download into a temporary file and atomically publish only a complete weight file.
def ensure_weight_file(path, url, min_bytes=1_000_000):
    if os.path.isfile(path) and os.path.getsize(path) >= min_bytes:
        return

    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    fd, temp_path = tempfile.mkstemp(prefix=".download-", dir=os.path.dirname(path) or ".")
    os.close(fd)
    try:
        with urllib.request.urlopen(url, timeout=60) as response, open(temp_path, "wb") as out:
            shutil.copyfileobj(response, out)
        if os.path.getsize(temp_path) < min_bytes:
            raise RuntimeError(f"Downloaded weight file is unexpectedly small: {path}")
        os.replace(temp_path, path)
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)


# 1. Real-ESRGAN Weights Setup
esrgan_path = f"{weights_dir}/RealESRGAN_x4plus.pth"
ensure_weight_file(
    esrgan_path,
    "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
)

esrgan_model = RRDBNet(
    num_in_ch=3, num_out_ch=3, num_feat=64,
    num_block=23, num_grow_ch=32, scale=4
)

restoration_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
enhancer = RealESRGANer(
    scale=4, model_path=esrgan_path, model=esrgan_model,
    tile=512, tile_pad=10, pre_pad=0,
    half=(restoration_device.type == "cuda"),
    device=restoration_device,
)

# 2. GFPGAN Weights Setup
gfpgan_path = f"{weights_dir}/GFPGANv1.3.pth"
ensure_weight_file(
    gfpgan_path,
    "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
)

gfpgan_restorer = GFPGANer(
    model_path=gfpgan_path,
    upscale=1,  # Set base upscale to 1, as Real-ESRGAN handles base canvas scaling
    arch="clean",
    channel_multiplier=2,
    bg_upsampler=enhancer,
    device=restoration_device
)

print("GFPGAN & Real-ESRGAN ready.")


# ================================================================
# 4. FLUX.1-Kontext + ColorizeTruer LoRA
# ================================================================
print("Loading FLUX.1-Kontext NF4...")
kontext_pipe = FluxKontextPipeline.from_pretrained(
    "aniketppanchal/flux.1-kontext-dev-nf4-pkg",
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)
kontext_pipe.load_lora_weights(
    "AlekseyCalvin/ColorizeTruer_Merge5wRef_KontexFlux_LoRA_bySAP"
)
print("✅ FLUX.1-Kontext NF4 loaded.")


# ================================================================
# 5. LOAD CLIP
# ================================================================
print("Loading CLIP for metrics...")
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to("cpu")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
print("All models loaded successfully.")


# ================================================================
# 6. HELPER FUNCTIONS
# ================================================================
def is_greyscale_or_sepia(image):
    rgb = np.asarray(image.convert("RGB"), dtype=np.uint8)
    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
    mean_sat = float(np.mean(hsv[:, :, 1]))
    is_greyscale = mean_sat < 25

    # A simple warm-channel heuristic; intentionally avoid classifying grayscale
    # images as sepia. This is a prompt hint, not a strict image classifier.
    channel_means = rgb.mean(axis=(0, 1))
    red_mean, green_mean, blue_mean = channel_means
    is_sepia = (
        not is_greyscale
        and red_mean > green_mean + 8
        and green_mean > blue_mean + 5
    )
    return is_greyscale, bool(is_sepia)


def calculate_padding(ow, oh, direction, pct):
    extend = pct / 100.0
    h_ext, v_ext = int(ow * extend), int(oh * extend)
    left = top = right = bottom = 0
    if direction == "Left": left = h_ext
    elif direction == "Right": right = h_ext
    elif direction == "Top": top = v_ext
    elif direction == "Bottom": bottom = v_ext
    elif direction == "Horizontal": left = right = h_ext
    elif direction == "Vertical": top = bottom = v_ext
    elif direction in ["Both", "All Sides"]:
        left = right = h_ext
        top = bottom = v_ext
    return left, top, right, bottom


def blur_mask(mask_pil, radius=16):
    return mask_pil.filter(ImageFilter.GaussianBlur(radius=radius))


def cap_image(image, max_side=1024):
    if max(image.size) <= max_side:
        return image
    ratio = max_side / max(image.size)
    return image.resize((int(image.width * ratio), int(image.height * ratio)), Image.LANCZOS)


def _colorfulness(image):
    rgb = np.array(image.convert("RGB"))
    R, G, B = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    rg = np.absolute(R - G)
    yb = np.absolute(0.5 * (R + G) - B)
    std_rg, mean_rg = np.std(rg), np.mean(rg)
    std_yb, mean_yb = np.std(yb), np.mean(yb)
    return np.sqrt(std_rg**2 + std_yb**2) + 0.3 * np.sqrt(mean_rg**2 + mean_yb**2)


# ================================================================
# 7. ENHANCEMENT FUNCTION (PURE RESTORATION & UPSCALING)
# ================================================================
@spaces.GPU(duration=60)
def enhance_image(image, scale_factor, mode="Pure Super-Resolution (Safe for Group Photos)", fidelity_weight=0.15, compute_metrics=False):
    if image is None:
        raise gr.Error("Please upload an image first.")

    gc.collect()
    torch.cuda.empty_cache()

    if max(image.size) > 1200:
        image = cap_image(image, 1200)

    rgb_np = np.array(image.convert("RGB"), dtype=np.uint8)
    input_bgr = cv2.cvtColor(rgb_np, cv2.COLOR_RGB2BGR)

    target_scale = 4 if scale_factor == "4x" else 2

    try:
        if "Pure Super-Resolution" in mode:
            # Pure Real-ESRGAN Upscaling & Deblurring (preserves exact facial features and group structures)
            output_bgr, _ = enhancer.enhance(input_bgr, outscale=target_scale)
        else:
            # GFPGAN Face Deblurring + Real-ESRGAN Canvas Upscaling
            gfpgan_restorer.upscale = target_scale
            _, _, output_bgr = gfpgan_restorer.enhance(
                input_bgr,
                has_aligned=False,
                only_center_face=False,
                paste_back=True,
                weight=float(fidelity_weight)
            )

        output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
        output_image = Image.fromarray(output_rgb)

    except Exception as e:
        raise gr.Error(f"Restoration failed: {e}")
    finally:
        gc.collect()
        torch.cuda.empty_cache()

    info = f"Restored {output_image.width}x{output_image.height}px using {mode}"

    if compute_metrics:
        try:
            info += f"\n🎨 Colorfulness: {_colorfulness(output_image):.2f}"
            if brisque is not None:
                from torchvision import transforms
                t = transforms.ToTensor()(output_image).unsqueeze(0).to("cuda")
                info += f"\n📉 BRISQUE: {brisque(t, data_range=1.0).item():.2f} (lower better)"
        except Exception as err:
            info += f"\n⚠️ Metrics error: {err}"

    return output_image, info


# ================================================================
# 8. COLORIZATION 
# ================================================================
@spaces.GPU(duration=58)
def colour_image(image, strength, compute_metrics=False):
    if image is None:
        raise gr.Error("Please upload an image first.")
        
    image = cap_image(image, 1024)  
    original_image = image.convert("RGB")
    ow, oh = original_image.size

    target_w = int(16 * math.ceil(ow / 16))
    target_h = int(16 * math.ceil(oh / 16))

    resized_input = original_image.resize((target_w, target_h), Image.LANCZOS)
    prompt = "colorize this black and white photograph with rich, natural, vivid color tones."

    try:
        with torch.inference_mode():
            output = kontext_pipe(
                image=resized_input,
                prompt=prompt,
                height=target_h,
                width=target_w,
                num_inference_steps=20,
                guidance_scale=3.5,
                max_sequence_length=512,
            ).images[0]

        output_aligned = output.resize((ow, oh), Image.LANCZOS)
        output_rgb = np.array(output_aligned)

        if strength < 1.0:
            grey = np.array(original_image.convert("L"))
            grey_3ch = np.stack([grey, grey, grey], axis=-1)
            output_rgb = (
                strength * output_rgb.astype(float) +
                (1.0 - strength) * grey_3ch.astype(float)
            ).astype(np.uint8)

        out_img = Image.fromarray(output_rgb)

        info = ""
        if compute_metrics:
            try:
                cf_in = _colorfulness(original_image)
                cf_out = _colorfulness(out_img)
                info += f"🎨 Colorfulness: {cf_in:.2f} → {cf_out:.2f} (Δ {cf_out-cf_in:+.2f})\n"
                if brisque is not None:
                    from torchvision import transforms
                    t = transforms.ToTensor()(out_img).unsqueeze(0).to("cuda")
                    info += f"📉 BRISQUE: {brisque(t, data_range=1.0).item():.2f} (lower better)"
            except Exception as e:
                info += f"⚠️ Metrics error: {e}"

        return out_img, info
    finally:
        gc.collect()
        torch.cuda.empty_cache()


# ================================================================
# 9. OUTPAINTING
# ================================================================
@spaces.GPU(duration=120)
def fast_single_pass_outpaint(
    image,
    direction="Horizontal",
    extend_percent=25,
    custom_prompt="",
    custom_override=False,
    seed=42,
    randomize_seed=True,
    guidance_scale=3.5,
    steps=20,
    compute_metrics=False,
    progress=gr.Progress()
):
    if image is None:
        raise gr.Error("Please upload an image first.")
    
    start_time = time.time()
    steps = int(steps)
    if not 10 <= steps <= 50:
        raise gr.Error("Inference steps must be between 10 and 50.")

    image = cap_image(image, 1024)
    ow, oh = image.size

    progress(0.05, desc="Analyzing scene color and style...")
    bw, sepia = is_greyscale_or_sepia(image)
    
    if custom_override and custom_prompt and custom_prompt.strip():
        prompt = custom_prompt.strip()
    else:
        prompt = "Seamless background environment extension, natural scenery, outdoor landscape, empty background."
        if custom_prompt and custom_prompt.strip():
            prompt += f" {custom_prompt.strip()}"
        
        if bw:
            prompt += " Black and white, vintage monochrome, matching film grain."
        elif sepia:
            prompt += " Warm sepia tones."
        else:
            prompt += " Photorealistic, vivid natural colors, matching original lighting."

    progress(0.10, desc="Preparing padded canvas...")
    left, top, right, bottom = calculate_padding(ow, oh, direction, extend_percent)

    target_w = ow + left + right
    target_h = oh + top + bottom

    target_w = int(16 * math.ceil(target_w / 16))
    target_h = int(16 * math.ceil(target_h / 16))

    # Bound worst-case FLUX memory use. Raise this only after validating the
    # selected ZeroGPU allocation can handle larger canvases reliably.
    max_outpaint_side = 2048
    max_outpaint_pixels = 4_194_304
    if (
        target_w > max_outpaint_side
        or target_h > max_outpaint_side
        or target_w * target_h > max_outpaint_pixels
    ):
        raise gr.Error(
            f"Requested canvas is too large ({target_w}x{target_h}). "
            f"Reduce the extension percentage or use a smaller input image."
        )

    canvas = Image.new("RGB", (target_w, target_h), (127, 127, 127))
    canvas.paste(image, (left, top))

    mask = Image.new("L", (target_w, target_h), 255)
    draw = ImageDraw.Draw(mask)
    draw.rectangle([left, top, left + ow - 1, top + oh - 1], fill=0)
    feathered_mask = blur_mask(mask, radius=16)

    if randomize_seed or seed is None or seed < 0:
        seed = int(np.random.randint(0, 2**32 - 1))
    else:
        seed = int(seed)
        
    generator = torch.Generator("cuda").manual_seed(seed)

    try:
        progress(0.20, desc="Generating background extension with FLUX Fill...")
        with torch.inference_mode():
            result = pipe(
                prompt=prompt,
                image=canvas,
                mask_image=feathered_mask,
                height=target_h,
                width=target_w,
                max_sequence_length=512,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                generator=generator,
            ).images[0]

        mask_inv = Image.eval(feathered_mask, lambda p: 255 - p)
        mask_inv_cropped = mask_inv.crop((left, top, left + ow, top + oh))
        result.paste(image, (left, top), mask_inv_cropped)

        final_w = ow + left + right
        final_h = oh + top + bottom
        current_img = result.crop((0, 0, final_w, final_h))

        total_time = time.time() - start_time
        final_info = (
            f"⏱️ Time: {total_time:.2f} sec\n"
            f"📐 Final: {current_img.width}x{current_img.height}\n"
            f"🔢 Seed: {seed}\n"
            f"📝 Prompt: {prompt}\n"
            f"🔁 Passes: Single Pass\n"
        )

        if compute_metrics:
            progress(0.92, desc="Computing quality metrics...")
            try:
                cf = _colorfulness(current_img)
                final_info += f"\n🎨 Colorfulness: {cf:.2f}"

                if brisque is not None:
                    from torchvision import transforms
                    t = transforms.ToTensor()(current_img).unsqueeze(0).to("cuda")
                    bs = brisque(t, data_range=1.0).item()
                    final_info += f"\n📉 BRISQUE: {bs:.2f} (lower better)"

                try:
                    clip_model.to("cuda")
                    with torch.no_grad():
                        inputs = clip_processor(
                            text=[prompt[:77]], images=current_img,
                            return_tensors="pt", padding=True, truncation=True
                        ).to("cuda")
                        clip_score = clip_model(**inputs).logits_per_image.item()
                    final_info += f"\n🤖 CLIP: {clip_score:.4f} (higher better)"
                finally:
                    # Keep this large optional metrics model off the GPU even
                    # when preprocessing or scoring fails.
                    clip_model.to("cpu")
            except Exception as e:
                final_info += f"\n⚠️ Metrics error: {e}"

        progress(1.0, desc="Done!")
        return current_img, canvas, feathered_mask, final_info
    finally:
        gc.collect()
        torch.cuda.empty_cache()


# ================================================================
# 10. TAB HELPERS & GRADIO UI
# ================================================================

def send_to_outpaint(img):
    return (img, gr.update(selected="tab_outpaint")) if img is not None else (None, gr.update(selected="tab_outpaint"))

def send_to_enhance(img):
    return (img, gr.update(selected="tab_enhance")) if img is not None else (None, gr.update(selected="tab_enhance"))

with gr.Blocks(
    title="CanvasAI",
    theme=gr.themes.Soft(
        primary_hue=gr.themes.colors.violet,
        secondary_hue=gr.themes.colors.purple,
        neutral_hue=gr.themes.colors.slate,
        radius_size=gr.themes.sizes.radius_sm,
    )
) as demo:

    gr.Markdown("# CanvasAI\n*AI-powered image enhancement, colorization, and outpainting*")
    gr.Markdown("**Enhance** with GFPGAN & Real-ESRGAN  |  **Colorize** B&W photos  |  **Outpaint** with FLUX.1-Fill")
    gr.Markdown("---")

    with gr.Tabs() as tabs:
        
        # Colorize Tab
        with gr.Tab("🎨 Colorize", id="tab_colorize"):
            gr.Markdown("## B&W Photo Colorization\nFLUX.1-Kontext + ColorizeTruer LoRA.")
            with gr.Row(equal_height=True):
                with gr.Column(scale=1):
                    col_input = gr.Image(label="Upload B&W Image", type="pil", height=350)
                    col_strength = gr.Slider(minimum=0.5, maximum=1.0, value=0.9, step=0.05, label="Color Strength")
                    col_metrics = gr.Checkbox(label="📊 Compute quality metrics", value=False)
                    col_btn = gr.Button("Colorize", variant="primary", size="lg")
                with gr.Column(scale=1):
                    col_output = gr.Image(label="Colorized Result", type="pil", interactive=False, height=350)
                    col_info = gr.Textbox(label="Metrics", interactive=False, lines=3)
                    with gr.Row():
                        btn_col_to_outpaint = gr.Button("➡️ Send to Outpaint", variant="secondary")
                        btn_col_to_enhance  = gr.Button("✨ Send to Enhance", variant="secondary")

            col_btn.click(fn=colour_image, inputs=[col_input, col_strength, col_metrics], outputs=[col_output, col_info])

        # Outpaint Tab
        with gr.Tab("🖼 Outpaint", id="tab_outpaint"):
            gr.Markdown("## FLUX Outpainting Engine\nExpand image boundaries naturally using single-pass FLUX Fill without subject duplication.")

            with gr.Row():
                with gr.Column(scale=1):
                    out_input = gr.Image(type="pil", label="Input Image")
                    
                    direction = gr.Dropdown(
                        choices=["Horizontal", "Vertical", "Left", "Right", "Top", "Bottom", "Both"],
                        value="Horizontal",
                        label="Expansion Direction"
                    )
                    extend_percent = gr.Slider(minimum=10, maximum=100, step=5, value=25, label="Extension Percentage (%)")
                    
                    with gr.Accordion("Advanced Prompting Options", open=False):
                        custom_prompt = gr.Textbox(label="Additional Prompt Directives", placeholder="e.g. dense forest, snowy mountains...")
                        custom_override = gr.Checkbox(label="Override Auto-Prompt Completely", value=False)
                    
                    with gr.Accordion("Inference Hyperparameters", open=False):
                        guidance_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, value=3.5, label="Guidance Scale (Keep low ~3.5)")
                        steps = gr.Slider(minimum=10, maximum=50, step=1, value=20, label="Inference Steps")
                        seed = gr.Number(value=42, label="Seed", precision=0)
                        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                        compute_metrics = gr.Checkbox(label="Calculate Quality Metrics", value=False)

                    run_btn = gr.Button("Outpaint Image", variant="primary")

                with gr.Column(scale=1):
                    out_output = gr.Image(label="Outpainted Result")
                    info_output = gr.Textbox(label="Execution Info & Metrics", lines=6)
                    
                    btn_out_to_enhance = gr.Button("✨ Send to Enhance", variant="secondary")

                    with gr.Accordion("Debug Canvas & Feathered Mask", open=False):
                        debug_canvas = gr.Image(label="Padded Canvas")
                        debug_mask = gr.Image(label="Feathered Mask")

            run_btn.click(
                fn=fast_single_pass_outpaint,
                inputs=[
                    out_input, direction, extend_percent, custom_prompt, custom_override,
                    seed, randomize_seed, guidance_scale, steps, compute_metrics
                ],
                outputs=[out_output, debug_canvas, debug_mask, info_output]
            )

        # Enhance Tab UI
        with gr.Tab("✨ Enhance", id="tab_enhance"):
            gr.Markdown("## Restore old blurry photos, sharpen features, and upscale up to 4x.")
            with gr.Row(equal_height=True):
                with gr.Column(scale=1):
                    enh_input = gr.Image(label="Upload Image", type="pil", height=350)
                    
                    enh_mode = gr.Radio(
                        choices=[
                            "Pure Super-Resolution (Safe for Group Photos)", 
                            "Subtle Face Refinement (For Close-ups)"
                        ], 
                        value="Pure Super-Resolution (Safe for Group Photos)", 
                        label="Restoration Mode"
                    )
                    
                    enh_scale = gr.Radio(choices=["2x", "4x"], value="4x", label="Upscale Factor")
                    enh_weight = gr.Slider(
                        minimum=0.0, maximum=0.5, value=0.15, step=0.05, 
                        label="Face Restoration Weight (Only applies in Refinement mode; keep low <= 0.2)"
                    )
                    enh_metrics = gr.Checkbox(label="📊 Compute quality metrics", value=False)
                    enh_btn = gr.Button("Enhance Image", variant="primary", size="lg")
                with gr.Column(scale=1):
                    enh_output = gr.Image(label="Result", type="pil", interactive=False, height=350)
                    enh_info = gr.Textbox(label="Size Info + Metrics", interactive=False, lines=4)
            
            enh_btn.click(
                fn=enhance_image, 
                inputs=[enh_input, enh_scale, enh_mode, enh_weight, enh_metrics], 
                outputs=[enh_output, enh_info]
            )

    # Inter-tab Event Routing Connections
    btn_col_to_outpaint.click(fn=send_to_outpaint, inputs=[col_output], outputs=[out_input, tabs])
    btn_col_to_enhance.click(fn=send_to_enhance, inputs=[col_output], outputs=[enh_input, tabs])
    btn_out_to_enhance.click(fn=send_to_enhance, inputs=[out_output], outputs=[enh_input, tabs])

demo.queue().launch(show_error=True)