import gradio as gr
import os
import json
import zipfile
import shutil
from PIL import Image

Image.MAX_IMAGE_PIXELS = 200_000_000

from engine.mapper import map_all

TEMPLATES = {
    "mens_tshirt": {
        "Front": [4649, 5428],
        "Back": [4655, 5427],
        "Left Sleeve": [3915, 1939],
        "Right Sleeve": [3915, 1939],
        "Collar": [3615, 315],
    },
    "loost_tshirt": {
        "Front": [3803, 4890],
        "Back": [3803, 5026],
        "Left Sleeve": [3053, 1689],
        "Right Sleeve": [3053, 1689],
    },
    "unisex_sweatshirt": {
        "Front": [4164, 4716],
        "Back": [4164, 4716],
        "Collar": [2809, 413],
        "Left Cuff": [1575, 900],
        "Right Cuff": [1575, 900],
        "Left Sleeve": [3640, 3607],
        "Right Sleeve": [3640, 3607],
        "Waistband": [7050, 900],
    },
    "bomber_jacket": {
        "Front": [9251, 11226],
        "Back": [9120, 10067],
        "Left Sleeve": [7830, 9097],
        "Right Sleeve": [7830, 9097],
    },
}

CSS = """
body, .gradio-container {
  background: #050316 !important;
  color: #7ffcff !important;
  font-family: cursive !important;
}

.gradio-container {
  max-width: 100% !important;
}

textarea, input, select {
  background: transparent !important;
  color: white !important;
  border: 1px solid #62f7ff !important;
}

button {
  background: transparent !important;
  color: #7ffcff !important;
  border: 1px solid #ff4dff !important;
}

h1, h2, h3, label, p, span {
  text-shadow:
    0 0 8px #000,
    0 0 14px #00eaff,
    0 0 28px #ff2efb !important;
}
"""

def clean_output():
    if os.path.exists("output"):
        shutil.rmtree("output")
    os.makedirs("output/production_pngs", exist_ok=True)
    os.makedirs("output/proofs", exist_ok=True)

def generate(image_path, template_name, prompt, use_optional_ai):
    if not image_path:
        return None, "Please upload artwork first."

    clean_output()

    outputs, report = map_all(
        image_path=image_path,
        template_id=template_name,
        pieces=TEMPLATES[template_name],
        prompt=prompt or "",
        use_optional_ai=use_optional_ai,
    )

    files = []

    for piece_name, img in outputs.items():
        path = f"output/production_pngs/{piece_name}.png"
        img.save(path)
        files.append(path)

    report_path = "output/warnings.json"

    with open(report_path, "w") as f:
        json.dump(report, f, indent=2)

    zip_path = "output/soul_threading_output.zip"

    with zipfile.ZipFile(zip_path, "w") as z:
        for f in files:
            z.write(f, arcname=f"production_pngs/{os.path.basename(f)}")
        z.write(report_path, arcname="warnings.json")

    return zip_path, json.dumps(report, indent=2)

with gr.Blocks(title="Soul Threading") as demo:
    gr.Markdown(
        """
# Soul Threading  
## Garment-Aware AOP Mapping Engine

Upload artwork, enter a design guidance prompt, select a garment, and export production PNG pieces.

Current engine supports:
- Prompt-guided focal region selection
- Piece-specific mapping rules
- Sleeve/environment separation
- Mask-aware shaped PNG export when masks exist
- ZIP export with warnings.json
"""
    )

    image = gr.Image(type="filepath", label="Upload Artwork")

    template = gr.Dropdown(
        choices=list(TEMPLATES.keys()),
        value="mens_tshirt",
        label="Garment Template",
    )

    prompt = gr.Textbox(
        label="Textual Design Guidance Prompt",
        placeholder="Keep the main figure centered. Use smoke/flames on sleeves. Protect text from seams. Use lower smoke for pocket.",
        lines=4,
    )

    use_ai = gr.Checkbox(
        label="Use optional Hugging Face AI processors if configured",
        value=False,
    )

    btn = gr.Button("Generate Production ZIP")
    zip_out = gr.File(label="Download Production ZIP")
    report = gr.Textbox(label="Mapping Report / Warnings", lines=24)

    btn.click(
        generate,
        inputs=[image, template, prompt, use_ai],
        outputs=[zip_out, report],
    )

demo.launch(css=CSS)