"""
Trouve mon papillon 🦋  /  Find my Butterfly 🦋
Photobooth de l'Insectarium — Espace pour la vie Montréal.

Le visiteur prend un selfie. On compare son visage aux 14 papillons vedettes
de l'expo "Lepidoptera — Éloge aux papillons" et on l'insère dans le cadre du
papillon qui lui ressemble le plus.

Les noms (français / anglais / latin) et la photo du papillon sont déjà
intégrés dans chaque cadre PNG du dossier Cadres_photobooth/.
"""

import base64
import csv
import io
import os
import tempfile

import gradio as gr
import numpy as np
import torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

# --------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------
MODEL_CKPT = "sasha/autotrain-butterfly-similarity-2490576840"
CSV_PATH = "Sp_et_photos.csv"
PHOTOS_DIR = "Find_butterfly_2026_Best"   # photos de référence (1 par espèce)
FRAMES_DIR = "Cadres_photobooth"          # cadres transparents (1 par espèce)

# Fenêtre transparente (le cercle) où apparaît le selfie, dans l'espace 2000x1545.
# Mesurée sur les cadres (bbox des pixels transparents) — à mettre à jour si les
# templates changent : voir scripts/alpha bbox.
WINDOW = (476, 117, 1064, 705)  # left, top, right, bottom

device = "cuda" if torch.cuda.is_available() else "cpu"


# --------------------------------------------------------------------------
# Chargement du modèle + des papillons vedettes (une seule fois au démarrage)
# --------------------------------------------------------------------------
processor = AutoImageProcessor.from_pretrained(MODEL_CKPT)
model = AutoModel.from_pretrained(MODEL_CKPT).to(device).eval()


def embed(image: Image.Image) -> np.ndarray:
    """Calcule l'empreinte (embedding L2-normalisée) d'une image."""
    inputs = processor(images=image.convert("RGB"), return_tensors="pt").to(device)
    with torch.no_grad():
        out = model(**inputs)
    vec = out.pooler_output[0].cpu().numpy().astype("float32")
    norm = np.linalg.norm(vec)
    return vec / norm if norm > 0 else vec


def _norm(s: str) -> str:
    """Normalise un nom : minuscules, sans espaces/underscores/ponctuation."""
    return "".join(ch for ch in s.lower() if ch.isalnum())


def resolve_frame(latin: str):
    """Trouve le cadre d'une espèce, peu importe le format du nom de fichier.

    Tolère espaces, underscores, casse et concaténation
    (ex. 'Troides rhadamantus' -> 'troidesrhadamantus.png').
    """
    target = _norm(latin)
    for fn in os.listdir(FRAMES_DIR):
        stem, ext = os.path.splitext(fn)
        if ext.lower() == ".png" and _norm(stem) == target:
            return os.path.join(FRAMES_DIR, fn)
    return None


def load_butterflies():
    """Lit le CSV et calcule l'empreinte de la photo de référence de chaque espèce."""
    names, frames, embeddings = [], [], []
    with open(CSV_PATH, encoding="utf-8") as f:
        rows = list(csv.reader(f))
    for row in rows[1:]:  # saute l'en-tête
        if not row or not row[0].strip():
            continue
        latin, fr_name, canonical_img = row[0].strip(), row[1].strip(), row[3].strip()
        photo_path = os.path.join(PHOTOS_DIR, latin, canonical_img)
        frame_path = resolve_frame(latin)
        if not (os.path.exists(photo_path) and frame_path):
            print(f"⚠️  Fichier manquant pour {latin} (photo ou cadre) — espèce ignorée.")
            continue
        embeddings.append(embed(Image.open(photo_path)))
        names.append(fr_name)
        frames.append(frame_path)
    print(f"✅ {len(names)} papillons vedettes chargés.")
    return names, frames, np.stack(embeddings)


NAMES, FRAMES, EMB = load_butterflies()


# --------------------------------------------------------------------------
# Composition du selfie dans le cadre
# --------------------------------------------------------------------------
def crop_to_aspect(img: Image.Image, target_w: int, target_h: int) -> Image.Image:
    """Recadre l'image au centre selon le ratio voulu, puis redimensionne."""
    w, h = img.size
    target_ratio = target_w / target_h
    if w / h > target_ratio:
        new_w = int(h * target_ratio)
        left = (w - new_w) // 2
        img = img.crop((left, 0, left + new_w, h))
    else:
        new_h = int(w / target_ratio)
        top = (h - new_h) // 2
        img = img.crop((0, top, w, top + new_h))
    return img.resize((target_w, target_h), Image.LANCZOS)


def compose(selfie: Image.Image, frame_path: str) -> Image.Image:
    """Place le selfie derrière le cadre : il apparaît dans le cercle transparent."""
    frame = Image.open(frame_path).convert("RGBA")
    win_w, win_h = WINDOW[2] - WINDOW[0], WINDOW[3] - WINDOW[1]
    selfie = crop_to_aspect(selfie.convert("RGBA"), win_w, win_h)
    canvas = Image.new("RGBA", frame.size, (255, 255, 255, 255))
    canvas.paste(selfie, (WINDOW[0], WINDOW[1]))
    canvas.alpha_composite(frame)
    return canvas.convert("RGB")


# --------------------------------------------------------------------------
# Logique de matching
# --------------------------------------------------------------------------
def match(selfie: Image.Image) -> Image.Image:
    """Trouve le papillon le plus proche et renvoie la composition finale."""
    query = embed(selfie)
    sims = EMB @ query  # similarité cosinus (vecteurs normalisés)
    best = int(np.argmax(sims))
    return compose(selfie, FRAMES[best])


# --------------------------------------------------------------------------
# Parcours photobooth (3 étapes)
# --------------------------------------------------------------------------
COUNTDOWN_SECONDS = 5


def show_camera():
    """Étape 2 (serveur) : cache l'accueil et un éventuel résultat précédent.

    L'ouverture réelle de la caméra et le décompte se font côté navigateur
    (voir CAPTURE_JS). Sorties : [start_btn, intro_md, result,
    download_btn, retake_btn].
    """
    return (
        gr.update(visible=False),  # start_btn
        gr.update(visible=False),  # intro_md
        gr.update(visible=False),  # result
        gr.update(visible=False),  # download_btn
        gr.update(visible=False),  # retake_btn
    )


def process_photo(data_url):
    """Étape 3 (serveur) : décode le selfie capturé, compose, affiche le résultat.

    Sorties : [start_btn, intro_md, result, download_btn, retake_btn].
    """
    if not data_url:
        return (gr.skip(),) * 5

    if data_url == "ERROR":
        # Caméra refusée / indisponible : on revient à l'accueil.
        return (
            gr.update(visible=True),   # start_btn
            gr.update(visible=True),   # intro_md
            gr.update(visible=False),  # result
            gr.update(visible=False),  # download_btn
            gr.update(visible=False),  # retake_btn
        )

    # data_url = "data:image/jpeg;base64,...."
    b64 = data_url.split(",", 1)[1]
    selfie = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
    result = match(selfie)

    # Sauvegarde dans un dossier temporaire unique (nom propre pour le téléchargement).
    out_dir = tempfile.mkdtemp()
    path = os.path.join(out_dir, "mon_papillon.jpg")
    result.save(path, quality=92)

    return (
        gr.update(visible=False),               # start_btn
        gr.update(visible=False),               # intro_md
        gr.update(value=result, visible=True),  # result
        gr.update(value=path, visible=True),    # download_btn
        gr.update(visible=True),                # retake_btn
    )


# JS exécuté dans le navigateur : construit un overlay plein écran avec l'aperçu
# caméra et le décompte, demande la caméra, capture l'image à la fin du décompte.
# Retourne l'image en base64 (ou "ERROR" si la caméra est refusée/indisponible).
CAPTURE_JS = f"""
async () => {{
    const overlay = document.createElement('div');
    overlay.style.cssText = 'position:fixed; inset:0; z-index:9999;'
        + 'background:rgba(0,0,0,0.92); display:flex;'
        + 'align-items:center; justify-content:center;';
    const video = document.createElement('video');
    video.autoplay = true; video.muted = true;
    video.setAttribute('playsinline', '');
    video.style.cssText = 'max-width:92vw; max-height:80vh; border-radius:12px;';
    const count = document.createElement('div');
    count.style.cssText = 'position:absolute; font-size:22vw; font-weight:800;'
        + 'color:#fff; text-shadow:0 2px 24px rgba(0,0,0,0.8);';
    overlay.appendChild(video);
    overlay.appendChild(count);
    document.body.appendChild(overlay);

    let stream;
    try {{
        stream = await navigator.mediaDevices.getUserMedia(
            {{ video: {{ facingMode: 'user' }}, audio: false }}
        );
    }} catch (e) {{
        overlay.remove();
        return 'ERROR';
    }}

    video.srcObject = stream;
    try {{ await video.play(); }} catch (e) {{}}

    // Décompte (ne démarre qu'une fois la caméra autorisée et active).
    for (let i = {COUNTDOWN_SECONDS}; i >= 1; i--) {{
        count.textContent = i;
        await new Promise(r => setTimeout(r, 1000));
    }}
    count.textContent = '📸';
    await new Promise(r => setTimeout(r, 250));

    // Capture l'image courante.
    const w = video.videoWidth, h = video.videoHeight;
    const canvas = document.createElement('canvas');
    canvas.width = w; canvas.height = h;
    canvas.getContext('2d').drawImage(video, 0, 0, w, h);
    const dataUrl = canvas.toDataURL('image/jpeg', 0.92);

    stream.getTracks().forEach(t => t.stop());
    overlay.remove();
    return dataUrl;
}}
"""


# --------------------------------------------------------------------------
# Interface
# --------------------------------------------------------------------------
PANTONE_7417C = "#E1523D"

CSS = f"""
#title, #subtitle, #intro {{ text-align: center; }}
#title h1 {{ margin-bottom: 0; }}
#result {{ max-width: 540px; margin: 0 auto; }}

/* Le français (1er paragraphe) plus grand que l'anglais (<small>) */
#intro p {{ margin: 0.2em 0; }}
#intro p:first-child {{ font-size: 1.2rem; line-height: 1.4; }}
#intro small {{ font-size: 0.85rem; opacity: 0.75; }}

.pantone-btn, .pantone-btn button, button.pantone-btn {{
    background: {PANTONE_7417C} !important;
    border-color: {PANTONE_7417C} !important;
    color: #ffffff !important;
}}
.pantone-btn:hover, .pantone-btn button:hover, button.pantone-btn:hover {{
    filter: brightness(0.92);
}}

/* Boutons bilingues : libellé FR (taille du bouton) + EN plus petit en dessous */
.bibtn::after {{
    display: block;
    font-size: 0.7em;
    font-weight: 400;
    font-style: normal;
    opacity: 0.85;
    margin-top: 2px;
}}
.bi-find::after {{ content: "Find my butterfly"; }}
.bi-dl::after   {{ content: "Download my photo"; }}
.bi-rt::after   {{ content: "Take another photo"; }}
"""

with gr.Blocks(title="Trouve mon papillon 🦋", theme=gr.themes.Soft(), css=CSS) as demo:
    # Toujours visibles
    gr.Markdown(
        "# Quel papillon es-tu ? 🦋\n"
        "#### What butterfly are you?",
        elem_id="title",
    )
    gr.Markdown(
        "### Lepidoptera — Éloge aux papillons\n"
        "##### *A tribute to butterflies* · Insectarium, Espace pour la vie",
        elem_id="subtitle",
    )

    # Visible seulement à l'accueil (étape 1) — FR (grand) puis EN (<small>)
    intro_md = gr.Markdown(
        "Chaque papillon arbore des couleurs uniques qui le rendent reconnaissable. "
        "Montre tes propres couleurs pour découvrir à quel papillon tu ressembles le plus !\n\n"
        "<small>Butterflies and moths show off unique colors that makes each of them "
        "recognizable. Show off your own colors and find out which butterfly you are!</small>",
        elem_id="intro",
    )

    start_btn = gr.Button(
        "Trouve mon papillon ! 🦋",
        variant="primary",
        size="lg",
        elem_classes=["pantone-btn", "bibtn", "bi-find"],
    )

    # Étape 2 : aperçu caméra + décompte gérés entièrement en JS (overlay plein écran)
    photo_data = gr.Textbox(visible=False)  # reçoit le selfie capturé (base64)

    # Étape 3 : résultat + téléchargement + reprise (cachés au départ)
    result = gr.Image(
        show_label=False,
        show_download_button=False,
        interactive=False,
        format="jpeg",
        visible=False,
        elem_id="result",
    )
    download_btn = gr.DownloadButton(
        "📥 Télécharger ma photo",
        variant="primary",
        visible=False,
        elem_classes=["pantone-btn", "bibtn", "bi-dl"],
    )
    retake_btn = gr.Button(
        "↻ Reprendre une autre photo",
        variant="secondary",
        size="sm",
        visible=False,
        elem_classes=["bibtn", "bi-rt"],
    )

    step_outputs = [start_btn, intro_md, result, download_btn, retake_btn]

    # Clic accueil / reprise : on masque l'UI (serveur) puis on lance la caméra (JS).
    start_btn.click(show_camera, outputs=step_outputs).then(
        None, inputs=None, outputs=photo_data, js=CAPTURE_JS
    )
    retake_btn.click(show_camera, outputs=step_outputs).then(
        None, inputs=None, outputs=photo_data, js=CAPTURE_JS
    )
    # Quand le selfie capturé arrive, on compose et on affiche le résultat.
    photo_data.change(process_photo, inputs=photo_data, outputs=step_outputs)


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=4)
    demo.launch()