"""
Hide My Data — Desidentificación visual de imágenes médicas.
Interfaz web (Gradio) para Hugging Face Spaces.

Reutiliza el pipeline del proyecto (config.py + pipeline/):
    detección YOLO (CLAHE) → OCR → anonimización (blur / LaMa / SSL)
y añade cifrado reversible de los metadatos extraídos (Fernet) + descifrado.

Requisitos en el Space: ver requirements.txt.
Pesos necesarios: models/yolo/clahe_canny/weights/best.pt  (y best.pt del SSL si se usa).
"""

import os
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")   # sin entorno gráfico en el Space

import io
import zipfile
import tempfile
from datetime import datetime
from pathlib import Path

import cv2
import numpy as np
from cryptography.fernet import Fernet
import gradio as gr

# ── Parche defensivo para el bug "argument of type 'bool' is not iterable" de
# gradio_client.utils (presente en gradio 4.44; inocuo en gradio 5). Se aplica
# solo si las funciones existen, así no estorba en versiones donde ya está resuelto.
try:
    import gradio_client.utils as _gcu
    if hasattr(_gcu, "get_type"):
        _orig_get_type = _gcu.get_type
        def _safe_get_type(schema):
            if not isinstance(schema, dict):
                return "Any"
            return _orig_get_type(schema)
        _gcu.get_type = _safe_get_type
    if hasattr(_gcu, "_json_schema_to_python_type"):
        _orig_js2pt = _gcu._json_schema_to_python_type
        def _safe_js2pt(schema, defs=None):
            if isinstance(schema, bool):
                return "bool"
            try:
                return _orig_js2pt(schema, defs)
            except Exception:
                return "Any"
        _gcu._json_schema_to_python_type = _safe_js2pt
except Exception:
    pass

import config
from pipeline.preprocess import clahe_canny
from pipeline.anonymizer import anonymize
from pipeline.ocr        import run_ocr

CLASS_NAMES  = config.CLASS_NAMES
CLASS_COLORS = config.CLASS_COLORS            # BGR
CLASS_LABELS = {"name": "Nombre", "id": "ID", "age": "Edad",
                "date": "Fecha", "time": "Hora"}
METHOD_MAP   = {
    "Borrado inteligente (LaMa)": "lama",
    "Parche difuminado (blur)":   "blur",
}

# ── Modelos (carga perezosa, una sola vez) ───────────────────────────────────────
_yolo = None


def get_yolo():
    global _yolo
    if _yolo is None:
        from ultralytics import YOLO
        _yolo = YOLO(str(config.YOLO_CLAHE_W))
    return _yolo


# ── Helpers de visión ─────────────────────────────────────────────────────────────

def detect(bgr):
    """Detecta sobre la imagen realzada con CLAHE (entrada real del modelo)."""
    inp = clahe_canny(bgr)
    return get_yolo().predict(inp, conf=config.CONF, verbose=False)[0].boxes


def draw_boxes(bgr, boxes, selected_ids):
    out = bgr.copy()
    for b in boxes:
        cid = int(b.cls.item())
        if cid not in selected_ids:
            continue
        x1, y1, x2, y2 = map(int, b.xyxy[0].tolist())
        color = CLASS_COLORS.get(cid, (255, 255, 255))
        cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)
        cv2.putText(out, CLASS_NAMES[cid], (x1, max(12, y1 - 4)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)
    return out


def _bgr2rgb(bgr):
    return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)


# ── Pestaña 1: Anonimizar ─────────────────────────────────────────────────────────

def run_anonymize(files, selected_labels, method_label, progress=gr.Progress()):
    # gr.File(file_count="multiple") devuelve una lista; por robustez admitimos
    # también un único path o None.
    if files is None:
        files = []
    elif isinstance(files, (str, Path)):
        files = [files]
    if not files:
        return [], [], None, "⚠️ Sube al menos una imagen."
    selected_ids = {i for i, n in enumerate(CLASS_NAMES)
                    if CLASS_LABELS[n] in selected_labels}
    if not selected_ids:
        return [], [], None, "⚠️ Selecciona al menos una categoría a anonimizar."

    method        = METHOD_MAP[method_label]
    inpaint_model = None

    ts       = datetime.now().strftime("%Y%m%d_%H%M%S")
    cls_tag  = "_".join(sorted(CLASS_NAMES[i] for i in selected_ids))
    work     = Path(tempfile.mkdtemp())
    root_out = work / f"hidemydata__{ts}__{cls_tag}"
    d_img, d_plain, d_enc, d_key = (root_out / s for s in
                                    ("images", "metadata_plain", "metadata_encrypted", "keys"))
    for d in (d_img, d_plain, d_enc, d_key):
        d.mkdir(parents=True, exist_ok=True)

    fernet_key = Fernet.generate_key()
    fernet     = Fernet(fernet_key)
    (d_key / "decryption_key.txt").write_bytes(fernet_key)

    det_gallery, anon_gallery, log = [], [], []
    n_total = len(files)
    for idx, fpath in enumerate(files, 1):
        progress((idx - 1) / n_total, desc=f"Procesando {idx}/{n_total}…")
        try:
            bgr = cv2.imread(str(fpath))
            if bgr is None:
                log.append(f"[omitida] no se pudo leer {Path(fpath).name}")
                continue
            img_id   = f"{idx:05d}"
            boxes    = detect(bgr)
            filtered = [b for b in boxes if int(b.cls.item()) in selected_ids]
            anon     = anonymize(method, bgr, filtered, inpaint_model)

            cv2.imwrite(str(d_img / f"{img_id}.png"), anon)
            det_gallery.append((_bgr2rgb(draw_boxes(bgr, boxes, selected_ids)), Path(fpath).name))
            anon_gallery.append((_bgr2rgb(anon), Path(fpath).name))

            # OCR → metadatos en claro + cifrados
            fields = []
            try:
                for det in run_ocr(bgr, filtered)["detections"]:
                    t = det.get("text", "").strip()
                    if t:
                        fields.append(f"{CLASS_LABELS.get(det['class_name'], det['class_name'])}: {t}")
            except Exception as e:
                log.append(f"[aviso OCR] {Path(fpath).name}: {e}")
            plain = "\n".join(fields) or "(sin texto detectado)"
            (d_plain / f"{img_id}.txt").write_text(plain, encoding="utf-8")
            (d_enc / f"{img_id}.enc").write_bytes(fernet.encrypt(plain.encode("utf-8")))
            log.append(f"[{idx}/{n_total}] {Path(fpath).name} → {img_id}.png  ({len(filtered)} cajas)")
        except Exception as e:
            log.append(f"[error] {Path(fpath).name}: {e}")
    progress(1.0, desc="Empaquetando…")

    # ZIP de salida (imágenes + metadatos claros + cifrados + clave)
    zip_path = str(work / f"hidemydata__{ts}.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for p in root_out.rglob("*"):
            if p.is_file():
                zf.write(p, p.relative_to(root_out.parent))

    log.append("\n✓ Listo. El ZIP incluye las imágenes anonimizadas, los metadatos "
               "(en claro y cifrados) y la clave de descifrado (keys/decryption_key.txt).")
    return det_gallery, anon_gallery, zip_path, "\n".join(log)


# ── Pestaña 2: Descifrar ───────────────────────────────────────────────────────────

def run_decrypt(enc_files, key_text, key_file):
    key = (key_text or "").strip().encode() if key_text else None
    if not key and key_file:
        key = Path(key_file).read_bytes().strip()
    if not key:
        return None, "⚠️ Pega la clave o sube el fichero decryption_key.txt."
    if not enc_files:
        return None, "⚠️ Sube al menos un fichero .enc."

    try:
        fernet = Fernet(key)
    except Exception as e:
        return None, f"⚠️ Clave inválida: {e}"

    work    = Path(tempfile.mkdtemp())
    out_dir = work / "descifrado"
    out_dir.mkdir(parents=True, exist_ok=True)

    lines, ok = [], 0
    for f in enc_files:
        name = Path(f).name
        try:
            plain = fernet.decrypt(Path(f).read_bytes()).decode("utf-8")
            (out_dir / (Path(f).stem + ".txt")).write_text(plain, encoding="utf-8")
            lines.append(f"── {name} ──\n{plain}\n")
            ok += 1
        except Exception as e:
            lines.append(f"✗ {name}: {e}\n")

    zip_path = str(work / "descifrado.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for p in out_dir.glob("*.txt"):
            zf.write(p, p.name)

    header = f"✓ {ok}/{len(enc_files)} ficheros descifrados.\n\n"
    return zip_path, header + "\n".join(lines)


# ── Interfaz ────────────────────────────────────────────────────────────────────

# Paleta tomada de la app de escritorio (PyQt6): navy + teal.
NAVY, PANEL, INPUT, BORDER = "#0D1B2A", "#111F30", "#0A1520", "#1E3048"
TEAL, INK, MUTED           = "#00C9A7", "#F0F4F8", "#6B829E"
CLASS_HEX = {"name": "#FF6B6B", "id": "#FFD93D", "age": "#6BCB77",
             "date": "#4D96FF", "time": "#C77DFF"}

THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.teal,
    neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("Inter"), "Segoe UI", "sans-serif"],
).set(
    body_background_fill=NAVY,
    body_text_color=INK,
    background_fill_primary=PANEL,
    background_fill_secondary=NAVY,
    block_background_fill=PANEL,
    block_border_color=BORDER,
    block_border_width="1px",
    block_radius="10px",
    block_label_text_color=TEAL,
    block_title_text_color=TEAL,
    input_background_fill=INPUT,
    input_border_color=BORDER,
    border_color_primary=BORDER,
    button_primary_background_fill=TEAL,
    button_primary_background_fill_hover="#00DDB0",
    button_primary_text_color="#06231D",
    button_secondary_background_fill="#1A3A5C",
    button_secondary_text_color=INK,
    color_accent_soft="#1E3048",
)

CSS = f"""
.gradio-container {{ background: {NAVY} !important; max-width: 1180px !important; }}
footer {{ display: none !important; }}
#hmd-header {{
    display:flex; align-items:center; gap:14px;
    padding:14px 20px; margin-bottom:6px;
    background:#162032; border:1px solid {BORDER};
    border-bottom:2px solid {TEAL}55; border-radius:12px;
}}
#hmd-logo {{ color:{TEAL}; font-size:24px; font-weight:800; letter-spacing:2px; }}
#hmd-sub  {{ color:{MUTED}; font-size:13px; }}
#hmd-badge {{
    margin-left:auto; color:{TEAL}; background:{TEAL}22;
    border:1px solid {TEAL}44; border-radius:9px; padding:3px 11px;
    font-size:12px; font-weight:600;
}}
.hmd-legend {{ display:flex; gap:16px; flex-wrap:wrap; font-size:13px; color:{INK}; padding:2px 2px 8px; }}
.hmd-legend span {{ display:inline-flex; align-items:center; gap:6px; }}
.hmd-dot {{ width:12px; height:12px; border-radius:3px; display:inline-block; }}
"""

HEADER = f"""
<div id="hmd-header">
  <div id="hmd-logo">HideMyData</div>
  <div id="hmd-sub">Desidentificación visual de imágenes médicas · privacidad con utilidad clínica</div>
  <div id="hmd-badge">v1.0</div>
</div>
"""

LEGEND = ('<div class="hmd-legend">' + "".join(
    f'<span><span class="hmd-dot" style="background:{CLASS_HEX[n]}"></span>{CLASS_LABELS[n]}</span>'
    for n in CLASS_NAMES) + "</div>")


with gr.Blocks(title="HideMyData", theme=THEME, css=CSS) as demo:
    gr.HTML(HEADER)

    with gr.Tab("  Anonimizar  "):
        with gr.Row(equal_height=False):
            with gr.Column(scale=4):
                in_files = gr.File(label="Imágenes  (arrastra una o varias · .png / .jpg)",
                                   file_count="multiple", file_types=["image"],
                                   type="filepath", height=180)
                gr.HTML(LEGEND)
                cls_sel = gr.CheckboxGroup(
                    choices=[CLASS_LABELS[n] for n in CLASS_NAMES],
                    value=[CLASS_LABELS[n] for n in CLASS_NAMES],
                    label="Categorías a ocultar")
                method_sel = gr.Radio(choices=list(METHOD_MAP.keys()),
                                      value="Borrado inteligente (LaMa)",
                                      label="Método de anonimización")
                btn = gr.Button("Anonimizar", variant="primary", size="lg")
            with gr.Column(scale=6):
                with gr.Row():
                    g_det  = gr.Gallery(label="Detección", columns=2, height=300,
                                        object_fit="contain")
                    g_anon = gr.Gallery(label="Anonimizadas", columns=2, height=300,
                                        object_fit="contain")
                out_zip = gr.File(label="Descargar resultados (ZIP: imágenes + metadatos + clave)")
                log = gr.Textbox(label="Registro", lines=7, max_lines=14, show_copy_button=True)
        btn.click(run_anonymize, [in_files, cls_sel, method_sel],
                  [g_det, g_anon, out_zip, log])

    with gr.Tab("  Descifrar metadatos  "):
        gr.Markdown("Recupera los datos originales subiendo los ficheros **`.enc`** y la **clave** "
                    "(`decryption_key.txt`) que se generó al anonimizar.")
        with gr.Row(equal_height=False):
            with gr.Column(scale=4):
                enc_files = gr.File(label="Ficheros .enc", file_count="multiple",
                                    type="filepath", height=150)
                key_text  = gr.Textbox(label="Clave (pega el contenido de decryption_key.txt)",
                                       lines=1, type="password")
                key_file  = gr.File(label="…o sube decryption_key.txt", type="filepath", height=90)
                btn_dec   = gr.Button("Descifrar", variant="primary", size="lg")
            with gr.Column(scale=6):
                dec_zip = gr.File(label="Descargar textos descifrados (ZIP)")
                dec_out = gr.Textbox(label="Contenido descifrado", lines=14, show_copy_button=True)
        btn_dec.click(run_decrypt, [enc_files, key_text, key_file], [dec_zip, dec_out])


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