import os
import re
import math
import tempfile
import zipfile
from pathlib import Path
from typing import List

import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageOps, ImageFilter
from diffusers import (
    ControlNetModel,
    StableDiffusionControlNetImg2ImgPipeline,
    UniPCMultistepScheduler,
)

BASE_MODEL = os.getenv("BASE_MODEL", "runwayml/stable-diffusion-v1-5")
CONTROLNET_MODEL = os.getenv("CONTROLNET_MODEL", "lllyasviel/control_v11p_sd15_canny")

DEFAULT_PROMPT = (
    "literal hand-drawn graphite animation frame on warm off-white sketch paper, "
    "rough 2B and 4B pencil, visible construction marks, uneven pencil pressure, "
    "cross-hatching, eraser ghosts, paper tooth, imperfect hand-drawn contours, "
    "rough graphite anime storyboard, preserve the exact composition, pose, camera, "
    "character identity and guitar geometry, preserve existing red accents only, "
    "no new colors, no digital painting, no smooth vector edges"
)

DEFAULT_NEGATIVE = (
    "photorealistic, glossy, smooth, digital painting, watercolor, oil paint, airbrush, "
    "3d render, vector art, clean ink, extra fingers, extra limbs, deformed hands, "
    "mutated anatomy, different scene, different outfit, different guitar, text, watermark, logo"
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

_pipe = None


def natural_key(text: str):
    return [int(c) if c.isdigit() else c.lower() for c in re.split(r"(\d+)", text)]


def load_pipeline():
    global _pipe
    if _pipe is not None:
        return _pipe

    controlnet = ControlNetModel.from_pretrained(
        CONTROLNET_MODEL,
        torch_dtype=DTYPE,
    )

    pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
        BASE_MODEL,
        controlnet=controlnet,
        torch_dtype=DTYPE,
        safety_checker=None,
        requires_safety_checker=False,
    )
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
    pipe.enable_attention_slicing()

    if DEVICE == "cuda":
        try:
            pipe.enable_model_cpu_offload()
        except Exception:
            pipe = pipe.to("cuda")
    else:
        pipe = pipe.to("cpu")

    _pipe = pipe
    return _pipe


def resize_multiple_of_8(img: Image.Image, max_side: int):
    img = ImageOps.exif_transpose(img).convert("RGB")
    w, h = img.size
    scale = min(1.0, float(max_side) / float(max(w, h)))
    nw = max(64, int(round((w * scale) / 8) * 8))
    nh = max(64, int(round((h * scale) / 8) * 8))
    return img.resize((nw, nh), Image.Resampling.LANCZOS)


def canny_control(img: Image.Image):
    # Dependency-light edge map: PIL + NumPy only.
    gray = ImageOps.grayscale(img)
    edges = gray.filter(ImageFilter.FIND_EDGES)
    arr = np.asarray(edges, dtype=np.uint8)
    # suppress weak paper/grain edges so composition dominates
    arr = np.where(arr > 28, 255, 0).astype(np.uint8)
    edge = Image.fromarray(arr, mode="L")
    return Image.merge("RGB", (edge, edge, edge))


def extract_images(zip_path: str, workdir: str) -> List[Path]:
    src = Path(workdir) / "input"
    src.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path, "r") as zf:
        zf.extractall(src)

    exts = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
    imgs = [p for p in src.rglob("*") if p.suffix.lower() in exts]
    imgs.sort(key=lambda p: natural_key(str(p.relative_to(src))))
    if not imgs:
        raise gr.Error("No PNG/JPG/WebP/BMP images found inside the ZIP.")
    return imgs


def make_sheet(paths, out_path: Path, cols=4, thumb=(240, 180)):
    if not paths:
        return None
    canvases = []
    for p in paths[:16]:
        im = Image.open(p).convert("RGB")
        im.thumbnail(thumb, Image.Resampling.LANCZOS)
        bg = Image.new("RGB", thumb, (246, 243, 236))
        bg.paste(im, ((thumb[0]-im.width)//2, (thumb[1]-im.height)//2))
        canvases.append(bg)

    rows = math.ceil(len(canvases) / cols)
    sheet = Image.new("RGB", (cols*thumb[0], rows*thumb[1]), (250,248,242))
    for i, im in enumerate(canvases):
        sheet.paste(im, ((i % cols)*thumb[0], (i // cols)*thumb[1]))
    sheet.save(out_path)
    return str(out_path)


def process_zip(
    zip_file,
    prompt,
    negative_prompt,
    max_side,
    strength,
    control_scale,
    steps,
    guidance,
    seed,
    progress=gr.Progress(),
):
    if zip_file is None:
        raise gr.Error("Upload a ZIP first.")

    pipe = load_pipeline()
    workdir = tempfile.mkdtemp(prefix="graphite_")
    out_dir = Path(workdir) / "output"
    out_dir.mkdir(parents=True, exist_ok=True)

    images = extract_images(zip_file, workdir)
    previews = []

    for idx, path in enumerate(images, start=1):
        progress((idx - 1) / len(images), desc=f"{idx}/{len(images)} · {path.name}")

        source = resize_multiple_of_8(Image.open(path), int(max_side))
        control = canny_control(source)

        generator = torch.Generator(device="cpu").manual_seed(int(seed) + idx)

        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=source,
            control_image=control,
            strength=float(strength),
            controlnet_conditioning_scale=float(control_scale),
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            generator=generator,
        ).images[0]

        out = out_dir / f"{path.stem}.png"
        result.save(out, format="PNG")
        if len(previews) < 16:
            previews.append(out)

    result_zip = Path(workdir) / "graphite_frames.zip"
    with zipfile.ZipFile(result_zip, "w", zipfile.ZIP_DEFLATED) as zf:
        for p in sorted(out_dir.glob("*.png"), key=lambda x: natural_key(x.name)):
            zf.write(p, arcname=p.name)

    sheet = make_sheet(previews, Path(workdir) / "preview.png")
    progress(1.0, desc="Done")

    status = (
        f"Processed {len(images)} frame(s).\n"
        f"Runtime device: {DEVICE}\n"
        f"Base model: {BASE_MODEL}\n"
        f"ControlNet: {CONTROLNET_MODEL}"
    )
    return str(result_zip), sheet, status


with gr.Blocks(title="Graphite Frame Batch Restorer") as demo:
    gr.Markdown(
        "# ✏️ Graphite Frame Batch Restorer\n"
        "Upload a ZIP of individual storyboard/flip-book frames. "
        "The app redraws each frame as rough graphite while using an edge-control pass "
        "to keep composition and pose from wandering."
    )

    with gr.Row():
        with gr.Column():
            zip_in = gr.File(label="ZIP of frames", file_types=[".zip"], type="filepath")
            prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=6)
            negative = gr.Textbox(label="Negative prompt", value=DEFAULT_NEGATIVE, lines=4)

            with gr.Accordion("Advanced", open=False):
                max_side = gr.Slider(512, 1024, value=768, step=64, label="Model max side")
                strength = gr.Slider(0.15, 0.60, value=0.28, step=0.01, label="Img2Img strength")
                control_scale = gr.Slider(0.50, 1.50, value=1.00, step=0.05, label="Structure lock")
                steps = gr.Slider(10, 35, value=22, step=1, label="Steps")
                guidance = gr.Slider(3.0, 10.0, value=6.5, step=0.5, label="Guidance")
                seed = gr.Number(value=1337, precision=0, label="Seed")

            run = gr.Button("Process ZIP", variant="primary")

        with gr.Column():
            zip_out = gr.File(label="Processed ZIP")
            preview = gr.Image(label="Preview")
            status = gr.Textbox(label="Status", lines=5)

    run.click(
        process_zip,
        inputs=[zip_in, prompt, negative, max_side, strength, control_scale, steps, guidance, seed],
        outputs=[zip_out, preview, status],
    )

demo.queue(default_concurrency_limit=1).launch(ssr_mode=False)