#!/usr/bin/env python3
"""
Classifica le immagini in agenzie_allegati assegnando il `tipo` corretto.

Pattern xlsx:
- Foglio "Copertina"   → eliminate (di solito logo)
- Foglio "Dati"        → planimetria_layout (planimetria sede)
- Foglio "Conclusione" → in ordine: planimetria_emergenza, luogo_sicuro_esterno
- Altri fogli          → foto_immobile (foto/dettagli sezione Pericoli)

Pattern docx:
- L'ordine sequenziale del documento + esclusione loghi:
  1) prime non-logo dopo "PLANIMETRIA DELLA SEDE"   → planimetria_layout
  2) immagini centrali                              → foto_immobile
  3) ultime 1-2 immagini (sezione ALLEGATI)         → planimetria_emergenza, luogo_sicuro_esterno

Le immagini con flag `excluded=1` (loghi/icone) restano escluse.

Uso:
  python3 app/scripts/classify_images.py
"""
from __future__ import annotations
import sqlite3
import sys
import warnings
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path

warnings.filterwarnings("ignore")

ROOT = Path(__file__).resolve().parents[2]
import os as _os
DB = Path(_os.environ["IMPORT_DB"]) if _os.environ.get("IMPORT_DB") else ROOT / "import-tools" / "db" / "dev.db"
DVR_DIR = ROOT / "-mat" / "DVR"

NS_R   = "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}"
NS_M   = "http://schemas.openxmlformats.org/spreadsheetml/2006/main"
NS_XDR = "http://schemas.openxmlformats.org/drawingml/2006/spreadsheetDrawing"
NS_A   = "http://schemas.openxmlformats.org/drawingml/2006/main"
NS_W   = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"


# ---------- XLSX ----------

def map_xlsx_images(path: Path) -> dict[str, tuple[str, int]]:
    """Ritorna mapping image_filename → (sheet_name, from_row).

    image_filename = "image1.jpeg" (basename del file in xl/media/).
    """
    out: dict[str, tuple[str, int]] = {}
    with zipfile.ZipFile(path) as z:
        names = set(z.namelist())

        # 1) sheet name + rId (da workbook.xml)
        wb_root = ET.fromstring(z.read("xl/workbook.xml"))
        sheets = wb_root.find(f"{{{NS_M}}}sheets")
        rid_to_sheet_name: dict[str, str] = {}
        for s in sheets:
            rid_to_sheet_name[s.get(NS_R + "id")] = s.get("name")

        # 2) rId → sheet*.xml file (da workbook.xml.rels)
        rid_to_sheet_xml: dict[str, str] = {}
        rels_root = ET.fromstring(z.read("xl/_rels/workbook.xml.rels"))
        for rel in rels_root:
            tgt = rel.get("Target", "")
            if "worksheets/" in tgt:
                rid_to_sheet_xml[rel.get("Id")] = tgt.replace("worksheets/", "")

        # 3) Per ogni sheet*.xml.rels, leggi il file drawing*.xml linkato
        sheet_to_drawing: dict[str, str] = {}
        for rid, sheet_xml_file in rid_to_sheet_xml.items():
            sheet_name = rid_to_sheet_name.get(rid)
            rels_path = f"xl/worksheets/_rels/{sheet_xml_file}.rels"
            if rels_path not in names:
                continue
            r2 = ET.fromstring(z.read(rels_path))
            for rel in r2:
                tgt = rel.get("Target", "")
                if "drawing" in tgt.lower():
                    drawing_file = tgt.replace("../", "").replace("drawings/", "")
                    sheet_to_drawing[sheet_name] = drawing_file

        # 4) Per ogni drawing, mappare embed rId → image filename
        for sheet_name, drawing_file in sheet_to_drawing.items():
            d_path = f"xl/drawings/{drawing_file}"
            d_rels_path = f"xl/drawings/_rels/{drawing_file}.rels"
            if d_path not in names or d_rels_path not in names:
                continue
            d_root = ET.fromstring(z.read(d_path))
            d_rels = ET.fromstring(z.read(d_rels_path))
            embed_to_img: dict[str, str] = {}
            for rel in d_rels:
                tgt = rel.get("Target", "").replace("../", "").replace("media/", "")
                embed_to_img[rel.get("Id")] = tgt

            # 5) Estrai gli ancoraggi
            for anchor in d_root:
                from_n = anchor.find(f"{{{NS_XDR}}}from")
                pic_n = anchor.find(f"{{{NS_XDR}}}pic")
                if from_n is None or pic_n is None:
                    continue
                row_node = from_n.find(f"{{{NS_XDR}}}row")
                row = int(row_node.text) if row_node is not None else 0
                blip = pic_n.find(f".//{{{NS_A}}}blip")
                if blip is None:
                    continue
                embed = blip.get(NS_R + "embed")
                img = embed_to_img.get(embed)
                if img:
                    out[img] = (sheet_name, row)
    return out


def classify_xlsx_image(sheet: str, row: int, ordine_in_sezione: int) -> tuple[str, str]:
    """Ritorna (tipo, motivazione)."""
    s = sheet.strip().lower()
    if s == "copertina":
        return ("altro", "su Copertina (sezione decorativa)")
    if s == "dati":
        # Nel foglio Dati ci sono tipicamente 2 immagini: la foto di copertina
        # dell'edificio in cima al foglio (righe 1-25, sopra la tabella
        # anagrafica) e la planimetria della sede sotto la "DESCRIZIONE
        # IMMOBILE" (riga 38+).
        if row <= 25:
            return ("foto_immobile", f"foglio Dati riga {row} (sopra anagrafica)")
        return ("planimetria_layout", f"foglio Dati riga {row}")
    if s.startswith("conclusione"):
        if ordine_in_sezione == 0:
            return ("planimetria_emergenza", "ALL.1 in foglio Conclusione")
        if ordine_in_sezione == 1:
            return ("luogo_sicuro_esterno", "ALL.2 in foglio Conclusione")
        return ("altro", f"foglio Conclusione (img {ordine_in_sezione+1})")
    if s in ("pericoli", "rischi", "incendio"):
        return ("foto_immobile", f"foglio {sheet}")
    return ("altro", f"foglio {sheet}")


# ---------- DOCX ----------

def map_docx_images_order(path: Path) -> list[str]:
    """Ritorna lista ordinata dei nomi immagine come appaiono nel documento."""
    with zipfile.ZipFile(path) as z:
        # rId → image filename
        rels_path = "word/_rels/document.xml.rels"
        if rels_path not in z.namelist():
            return []
        rels = ET.fromstring(z.read(rels_path))
        embed_to_img: dict[str, str] = {}
        for rel in rels:
            tgt = rel.get("Target", "")
            if "media/" in tgt:
                embed_to_img[rel.get("Id")] = tgt.replace("media/", "")

        # Cerca tutti i blip nell'ordine di apparizione
        doc_root = ET.fromstring(z.read("word/document.xml"))
        ordered: list[str] = []
        for blip in doc_root.iter(f"{{{NS_A}}}blip"):
            embed = blip.get(NS_R + "embed")
            img = embed_to_img.get(embed)
            if img:
                ordered.append(img)
        return ordered


def classify_docx_images(image_order: list[str], excluded_set: set[str]) -> dict[str, tuple[str, str]]:
    """Per i docx: usa ordine + esclusione loghi.

    Strategia:
    - Filtra le immagini escluse (loghi/icone).
    - Le ultime 2 → ALL.1 (planimetria_emergenza), ALL.2 (luogo_sicuro_esterno)
    - La prima rimanente → planimetria_layout
    - Le centrali → foto_immobile
    """
    classification: dict[str, tuple[str, str]] = {}
    useful = [img for img in image_order if img not in excluded_set]
    if not useful:
        return classification
    n = len(useful)
    for idx, img in enumerate(useful):
        if n >= 3 and idx == n - 2:
            classification[img] = ("planimetria_emergenza", f"penultima immagine (ALL.1) di {n}")
        elif n >= 2 and idx == n - 1:
            classification[img] = ("luogo_sicuro_esterno", f"ultima immagine (ALL.2) di {n}")
        elif idx == 0:
            classification[img] = ("planimetria_layout", "prima immagine non-logo (planimetria sede)")
        else:
            classification[img] = ("foto_immobile", f"immagine {idx+1}/{n} (centrale)")
    return classification


# ---------- MAIN ----------

def main():
    conn = sqlite3.connect(DB)
    conn.execute("PRAGMA foreign_keys = ON")
    cur = conn.cursor()

    rows = cur.execute("""
        SELECT d.id, a.id, a.codice_fisico, d.source_file, d.source_format
        FROM dvr d JOIN agenzie a ON a.id = d.agenzia_id
        ORDER BY a.codice_fisico
    """).fetchall()
    print(f"Sedi da classificare: {len(rows)}")

    stats = {"xlsx": 0, "docx": 0, "errore": 0}
    tipo_count: dict[str, int] = {}

    for i, (dvr_id, agenzia_id, cf, source_file, source_format) in enumerate(rows, 1):
        src = DVR_DIR / source_file
        if not src.exists():
            stats["errore"] += 1
            continue

        # Allegati di questa agenzia (filename → id, excluded)
        allegati = cur.execute(
            "SELECT id, filename, excluded FROM agenzie_allegati WHERE agenzia_id = ?",
            (agenzia_id,),
        ).fetchall()
        if not allegati:
            continue
        excluded_set = {row[1] for row in allegati if row[2]}
        filename_to_id = {row[1]: row[0] for row in allegati}

        try:
            if source_format == "xlsx":
                img_to_anchor = map_xlsx_images(src)
                # Conta progressivo per "Conclusione" in ordine
                conclusione_imgs = sorted(
                    [(img, anc[1]) for img, anc in img_to_anchor.items() if anc[0].lower().startswith("conclusione")],
                    key=lambda x: (x[1], x[0]),
                )
                concl_index = {img: idx for idx, (img, _) in enumerate(conclusione_imgs)}
                for img_name, alleg_id in filename_to_id.items():
                    if img_name in excluded_set:
                        continue
                    if img_name not in img_to_anchor:
                        continue
                    sheet, row = img_to_anchor[img_name]
                    ordine = concl_index.get(img_name, 0)
                    tipo, reason = classify_xlsx_image(sheet, row, ordine)
                    cur.execute(
                        "UPDATE agenzie_allegati SET tipo = ?, note = note || ' | ' || ? WHERE id = ?",
                        (tipo, f"classify: {reason}", alleg_id),
                    )
                    tipo_count[tipo] = tipo_count.get(tipo, 0) + 1
                stats["xlsx"] += 1

            elif source_format == "docx":
                order = map_docx_images_order(src)
                clf = classify_docx_images(order, excluded_set)
                for img_name, (tipo, reason) in clf.items():
                    alleg_id = filename_to_id.get(img_name)
                    if not alleg_id:
                        continue
                    cur.execute(
                        "UPDATE agenzie_allegati SET tipo = ?, note = note || ' | ' || ? WHERE id = ?",
                        (tipo, f"classify: {reason}", alleg_id),
                    )
                    tipo_count[tipo] = tipo_count.get(tipo, 0) + 1
                stats["docx"] += 1
        except Exception as e:
            stats["errore"] += 1
            print(f"  ⚠ {cf} ({source_file}): {type(e).__name__}: {e}")

        if i % 100 == 0:
            conn.commit()
            print(f"  ... {i}/{len(rows)} sedi processate")

    conn.commit()

    print()
    print("=" * 60)
    print(f"FINE. xlsx={stats['xlsx']} docx={stats['docx']} errori={stats['errore']}")
    print()
    print("=== Distribuzione finale tipo (tutte le immagini, escluse comprese) ===")
    for r in cur.execute(
        """SELECT tipo, COUNT(*), SUM(CASE WHEN excluded=1 THEN 1 ELSE 0 END) as esc
           FROM agenzie_allegati GROUP BY tipo ORDER BY 2 DESC"""
    ):
        print(f"  {r[0]:<24} {r[1]:>5}   (di cui esclusi: {r[2]})")

    conn.close()


if __name__ == "__main__":
    main()