src/dataset_browser.py

"""
dataset_browser.py — Dataset Explorer for the OpenItaLaw Streamlit app.

Provides an interactive browser showing:
- All compressed archives and their internal files
- Metadata ↔ law linkage (which metadata corresponds to which law)
- FAISS index cross-references
- Local and dynamic (HF API) dataset browsing
"""

from __future__ import annotations

import json
import logging
import time
from pathlib import Path
from typing import Any

log = logging.getLogger(__name__)

# ─── HuggingFace Dataset Catalog ─────────────────────────────────────────────

# Hardcoded from https://huggingface.co/datasets/diatribe00/ItalianLawEngine/tree/main
# This is the concrete mapping the user requested — every file linked.

DATASET_FILES: list[dict] = [
    # ── Compressed Archives (source data) ──
    {
        "filename": "normattiva.tar.zst",
        "path": "normattiva.tar.zst",
        "size_bytes": 553_648_128,
        "size_human": "528 MB",
        "type": "archive",
        "source_type": "normattiva_akn",
        "category": "Normattiva — Legislazione Nazionale",
        "description": "Testi completi di legge italiani in formato Akoma Ntoso (XML). Fonte primaria per l'ordinamento vigente.",
        "icon": "📜",
        "content_type": "AKN XML",
        "est_documents": 450_000,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/normattiva.tar.zst",
    },
    {
        "filename": "corte_costituzionale.tar.zst",
        "path": "corte_costituzionale.tar.zst",
        "size_bytes": 137_363_456,
        "size_human": "131 MB",
        "type": "archive",
        "source_type": "cc_decision",
        "category": "Corte Costituzionale",
        "description": "Sentenze, ordinanze e decisioni della Corte Costituzionale della Repubblica Italiana.",
        "icon": "⚖️",
        "content_type": "Legal Decisions",
        "est_documents": 65_000,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/corte_costituzionale.tar.zst",
    },
    {
        "filename": "ANAntiCorruzioneOpen.tar.zst",
        "path": "ANAntiCorruzioneOpen.tar.zst",
        "size_bytes": 320_864_256,
        "size_human": "306 MB",
        "type": "archive",
        "source_type": "anac_document",
        "category": "ANAC — Anticorruzione",
        "description": "Documenti dell'Autorità Nazionale Anticorruzione: delibere, linee guida, pareri e segnalazioni.",
        "icon": "🔍",
        "content_type": "ANAC Open Data",
        "est_documents": 120_000,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/ANAntiCorruzioneOpen.tar.zst",
    },
    {
        "filename": "AffariInterniETerritorialiOpen.tar.zst",
        "path": "AffariInterniETerritorialiOpen.tar.zst",
        "size_bytes": 4_433_510,
        "size_human": "4.23 MB",
        "type": "archive",
        "source_type": "ministry_interior",
        "category": "Min. Affari Interni e Territoriali",
        "description": "Atti e documenti del Ministero dell'Interno: circolari, direttive, ordinanze.",
        "icon": "🏛️",
        "content_type": "Ministry Documents",
        "est_documents": 2_000,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/AffariInterniETerritorialiOpen.tar.zst",
    },
    {
        "filename": "MinInfrTraspOpen.tar.zst",
        "path": "MinInfrTraspOpen.tar.zst",
        "size_bytes": 94_164_992,
        "size_human": "89.8 MB",
        "type": "archive",
        "source_type": "ministry_infra",
        "category": "Min. Infrastrutture e Trasporti",
        "description": "Atti del Ministero delle Infrastrutture e dei Trasporti: regolamenti, piani, norme tecniche.",
        "icon": "🚆",
        "content_type": "Ministry Documents",
        "est_documents": 35_000,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/MinInfrTraspOpen.tar.zst",
    },
    {
        "filename": "senate.tar.zst",
        "path": "senate.tar.zst",
        "size_bytes": 70_866_534,
        "size_human": "67.6 MB",
        "type": "archive",
        "source_type": "senate_document",
        "category": "Senato della Repubblica",
        "description": "Atti parlamentari del Senato: DDL, relazioni, resoconti e interrogazioni.",
        "icon": "🏛️",
        "content_type": "Parliamentary Documents",
        "est_documents": 25_000,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/senate.tar.zst",
    },
    {
        "filename": "siope.tar.zst",
        "path": "siope.tar.zst",
        "size_bytes": 210_763_776,
        "size_human": "201 MB",
        "type": "archive",
        "source_type": "siope_spending",
        "category": "SIOPE — Spesa Pubblica",
        "description": "Dati di spesa pubblica dal SIOPE: mandati, reversali, bilanci di enti pubblici.",
        "icon": "💰",
        "content_type": "Financial Open Data",
        "est_documents": 80_000,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/siope.tar.zst",
    },
    # ── JSONL Corpus ──
    {
        "filename": "rebuild_v2_docs.jsonl",
        "path": "rebuild_v2_docs.jsonl",
        "size_bytes": 3_178_209_280,
        "size_human": "2.96 GB",
        "type": "jsonl_corpus",
        "source_type": "full_corpus",
        "category": "Corpus JSONL Completo",
        "description": "Corpus completo in formato JSONL: 839.422 documenti giuridici indicizzati con testo, URN, metadati.",
        "icon": "📋",
        "content_type": "JSONL",
        "est_documents": 839_422,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/rebuild_v2_docs.jsonl",
    },
    # ── FAISS Index ──
    {
        "filename": "index.faiss",
        "path": "faiss_index/index.faiss",
        "size_bytes": 1_395_864_576,
        "size_human": "1.3 GB",
        "type": "faiss_index",
        "source_type": "vector_index",
        "category": "Indice FAISS",
        "description": "Indice vettoriale FAISS (IndexFlatIP, dim=384) per ricerca semantica sul corpus legislativo corrente.",
        "icon": "🔎",
        "content_type": "FAISS Binary",
        "est_documents": 839_422,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/faiss_index/index.faiss",
    },
    {
        "filename": "doc_metadata.jsonl",
        "path": "faiss_index/doc_metadata.jsonl",
        "size_bytes": 1_610_612_736,
        "size_human": "1.5 GB",
        "type": "metadata_jsonl",
        "source_type": "vector_metadata",
        "category": "Metadati FAISS",
        "description": "Metadati per-vettore: una riga per ogni vettore FAISS. Contiene URN, titolo, testo, source_archive.",
        "icon": "📄",
        "content_type": "JSONL",
        "est_documents": 839_422,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/faiss_index/doc_metadata.jsonl",
    },
    {
        "filename": "build_manifest.json",
        "path": "faiss_index/build_manifest.json",
        "size_bytes": 215,
        "size_human": "215 B",
        "type": "manifest",
        "source_type": "build_tracking",
        "category": "Build Manifest",
        "description": "Tracciamento dell'ultima build dell'indice: data, documenti aggiunti, stato.",
        "icon": "📊",
        "content_type": "JSON",
        "est_documents": 1,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/faiss_index/build_manifest.json",
    },
    {
        "filename": "metadata.json",
        "path": "faiss_index/metadata.json",
        "size_bytes": 1_044,
        "size_human": "1.02 KB",
        "type": "metadata_json",
        "source_type": "index_metadata",
        "category": "Index Metadata",
        "description": "Metadata dell'indice: modello embedding, dimensione, tipo, distribuzione domini.",
        "icon": "ℹ️",
        "content_type": "JSON",
        "est_documents": 1,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/faiss_index/metadata.json",
    },
    # ── Archived Legacy Full Legal FAISS Index ──
    {
        "filename": "index.faiss (legacy)",
        "path": "full_legal_faiss_index/index.faiss",
        "size_bytes": 320_864_256,
        "size_human": "306 MB",
        "type": "faiss_index",
        "source_type": "legacy_index",
        "category": "Archived Legacy FAISS Index",
        "description": "Artefatto archiviato del RAG originale. Non usato dalla Space corrente; mantenuto solo per recupero storico.",
        "icon": "📦",
        "content_type": "FAISS Binary",
        "est_documents": 0,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/full_legal_faiss_index/index.faiss",
    },
    {
        "filename": "index.pkl",
        "path": "full_legal_faiss_index/index.pkl",
        "size_bytes": 134_217_728,
        "size_human": "128 MB",
        "type": "pickle",
        "source_type": "legacy_pickle",
        "category": "Legacy Pickle Index",
        "description": "Indice pickle LangChain del sistema storico, conservato come artefatto d'archivio.",
        "icon": "📦",
        "content_type": "Pickle",
        "est_documents": 0,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/full_legal_faiss_index/index.pkl",
    },
    {
        "filename": "legal_rag_final_system_checkpoint.pkl",
        "path": "full_legal_faiss_index/legal_rag_final_system_checkpoint.pkl",
        "size_bytes": 89_231_974,
        "size_human": "85.1 MB",
        "type": "pickle",
        "source_type": "legacy_checkpoint",
        "category": "Legacy System Checkpoint",
        "description": "Checkpoint storico del sistema RAG originale, non usato dall'app attuale.",
        "icon": "📦",
        "content_type": "Pickle",
        "est_documents": 0,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/full_legal_faiss_index/legal_rag_final_system_checkpoint.pkl",
    },
    # ── Documentation ──
    {
        "filename": "README.md",
        "path": "README.md",
        "size_bytes": 33,
        "size_human": "33 B",
        "type": "documentation",
        "source_type": "readme",
        "category": "Documentazione",
        "description": "README del dataset.",
        "icon": "📖",
        "content_type": "Markdown",
        "est_documents": 0,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/README.md",
    },
    {
        "filename": ".gitattributes",
        "path": ".gitattributes",
        "size_bytes": 2_826,
        "size_human": "2.76 KB",
        "type": "config",
        "source_type": "git_config",
        "category": "Configurazione",
        "description": "Configurazione Git LFS per file di grandi dimensioni.",
        "icon": "⚙️",
        "content_type": "Git Config",
        "est_documents": 0,
        "download_url": "https://huggingface.co/datasets/diatribe00/ItalianLawEngine/resolve/main/.gitattributes",
    },
]

# ─── Archive → Metadata Linkage Map ─────────────────────────────────────────
# Maps archive source_type to the metadata fields that link back to it

ARCHIVE_METADATA_LINKAGE = {
    "normattiva_akn": {
        "metadata_field": "source_archive",
        "match_values": ["normattiva.tar.zst", "normattiva"],
        "urn_pattern": r"^urn:nir:",
        "linked_fields": ["urn", "title", "pub_date", "text", "validity_status"],
        "description": "URN-based linkage: ogni atto su Normattiva ha un URN univoco (urn:nir:stato:...)",
    },
    "cc_decision": {
        "metadata_field": "source_archive",
        "match_values": ["corte_costituzionale.tar.zst", "corte_cost"],
        "urn_pattern": r"corte.cost|costituzionale",
        "linked_fields": ["urn", "title", "pub_date", "text"],
        "description": "Sentenze CC linkate per numero/anno (es. sent. 1/2014)",
    },
    "anac_document": {
        "metadata_field": "source_archive",
        "match_values": ["ANAntiCorruzioneOpen.tar.zst", "anac"],
        "urn_pattern": r"anac|anticorruzione",
        "linked_fields": ["urn", "title", "text"],
        "description": "Documenti ANAC linkati per codice delibera",
    },
    "ministry_interior": {
        "metadata_field": "source_archive",
        "match_values": ["AffariInterniETerritorialiOpen.tar.zst", "min_interno"],
        "urn_pattern": r"interior|interni|territoriali",
        "linked_fields": ["urn", "title", "text"],
        "description": "Atti ministeriali linkati per protocollo",
    },
    "ministry_infra": {
        "metadata_field": "source_archive",
        "match_values": ["MinInfrTraspOpen.tar.zst", "min_infra"],
        "urn_pattern": r"infrastruttur|trasport",
        "linked_fields": ["urn", "title", "text"],
        "description": "Atti ministeriali linkati per protocollo",
    },
    "senate_document": {
        "metadata_field": "source_archive",
        "match_values": ["senate.tar.zst", "senato"],
        "urn_pattern": r"senate|senato",
        "linked_fields": ["urn", "title", "pub_date", "text"],
        "description": "Atti parlamentari linkati per numero DDL/legislatura",
    },
    "siope_spending": {
        "metadata_field": "source_archive",
        "match_values": ["siope.tar.zst", "siope"],
        "urn_pattern": r"siope|spesa",
        "linked_fields": ["urn", "title", "text"],
        "description": "Dati finanziari linkati per codice ente/mandato",
    },
}


def compute_source_stats(metadata: list[dict]) -> dict[str, dict]:
    """Compute per-source statistics from loaded FAISS metadata.

    Returns: { source_type: { count, sample_urns, sample_titles, has_archive_link } }
    """
    stats: dict[str, dict] = {}
    for doc in metadata:
        if not isinstance(doc, dict):
            continue
        source = doc.get("build_source") or doc.get("source_type") or "unknown"
        if source not in stats:
            stats[source] = {
                "count": 0,
                "sample_urns": [],
                "sample_titles": [],
                "has_archive_link": 0,
                "has_urn": 0,
                "has_title": 0,
                "has_text": 0,
            }
        s = stats[source]
        s["count"] += 1
        urn = doc.get("urn", "")
        title = doc.get("title") or doc.get("source_title") or ""
        if urn:
            s["has_urn"] += 1
            if len(s["sample_urns"]) < 5:
                s["sample_urns"].append(urn)
        if title:
            s["has_title"] += 1
            if len(s["sample_titles"]) < 5:
                s["sample_titles"].append(title[:120])
        if doc.get("text"):
            s["has_text"] += 1
        if doc.get("source_archive"):
            s["has_archive_link"] += 1
    return stats


def find_metadata_for_archive(metadata: list[dict], archive_name: str) -> list[dict]:
    """Find all metadata entries that came from a specific archive."""
    results = []
    linkage = None
    for src_type, info in ARCHIVE_METADATA_LINKAGE.items():
        if archive_name in info["match_values"]:
            linkage = info
            break

    if not linkage:
        return results

    for idx, doc in enumerate(metadata):
        if not isinstance(doc, dict):
            continue
        src_archive = doc.get("source_archive", "")
        build_source = doc.get("build_source", "")
        source_type = doc.get("source_type", "")

        match = False
        for val in linkage["match_values"]:
            if val in src_archive or val in build_source or val in source_type:
                match = True
                break

        if match:
            results.append({"faiss_idx": idx, **doc})
    return results


def search_metadata_by_urn(metadata: list[dict], urn_query: str) -> list[dict]:
    """Search metadata entries by URN (partial match)."""
    results = []
    query_lower = urn_query.lower()
    for idx, doc in enumerate(metadata):
        if not isinstance(doc, dict):
            continue
        doc_urn = (doc.get("urn") or "").lower()
        if query_lower in doc_urn:
            results.append({"faiss_idx": idx, **doc})
            if len(results) >= 100:
                break
    return results


def search_metadata_by_title(metadata: list[dict], title_query: str) -> list[dict]:
    """Search metadata entries by title (partial match)."""
    results = []
    query_lower = title_query.lower()
    for idx, doc in enumerate(metadata):
        if not isinstance(doc, dict):
            continue
        doc_title = (doc.get("title") or doc.get("source_title") or "").lower()
        if query_lower in doc_title:
            results.append({"faiss_idx": idx, **doc})
            if len(results) >= 100:
                break
    return results


def get_archive_from_metadata(doc: dict) -> dict | None:
    """Given a metadata entry, find which archive file it came from."""
    source_archive = doc.get("source_archive", "")
    build_source = doc.get("build_source", "")
    source_type = doc.get("source_type", "")

    for file_info in DATASET_FILES:
        if file_info["type"] != "archive":
            continue
        fname = file_info["filename"]
        ftype = file_info["source_type"]
        if fname in source_archive or ftype == source_type or fname.replace(".tar.zst", "") in build_source:
            return file_info

    return None


def try_load_dynamic_manifest(data_dir: Path) -> list[dict] | None:
    """Try to load dataset_manifest.jsonl if it was downloaded."""
    manifest_path = data_dir / "dataset_manifest.jsonl"
    if not manifest_path.exists():
        return None

    entries = []
    try:
        with open(manifest_path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    entries.append(json.loads(line))
    except Exception as exc:
        log.warning("Failed to load dataset manifest: %s", exc)
        return None

    return entries if entries else None


def render_dataset_browser(st_module, manager, data_dir: Path, hf_repo_id: str, hf_token: str):
    """Render the full dataset browser tab in Streamlit.

    Args:
        st_module: Streamlit import (st)
        manager: LawDataManager instance with loaded metadata
        data_dir: Local data directory path
        hf_repo_id: HuggingFace dataset repo ID
        hf_token: HuggingFace token
    """
    st = st_module

    st.header("🗄️ Esplora Dataset — ItalianLawEngine")
    st.caption(
        f"Dataset: [`{hf_repo_id}`](https://huggingface.co/datasets/{hf_repo_id}) · "
        f"Dimensione totale: **7.6 GB** · File: **{len(DATASET_FILES)}**"
    )

    # ── Overview metrics ──
    archives = [f for f in DATASET_FILES if f["type"] == "archive"]
    total_size = sum(f["size_bytes"] for f in DATASET_FILES)
    archive_size = sum(f["size_bytes"] for f in archives)
    total_est_docs = sum(f.get("est_documents", 0) for f in archives)

    c1, c2, c3, c4 = st.columns(4)
    c1.metric("📦 Archivi Compressi", len(archives))
    c2.metric("📄 File Totali", len(DATASET_FILES))
    c3.metric("💾 Dimensione", f"{total_size / (1024**3):.1f} GB")
    c4.metric("📊 Documenti Stimati", f"{total_est_docs:,}")

    st.divider()

    # ── Sub-tabs for different views ──
    browse_mode = st.radio(
        "Modalità di esplorazione",
        ["📦 Archivi", "🔗 Linkage Metadati↔Leggi", "🔍 Cerca nel Dataset", "📊 Source Breakdown"],
        horizontal=True,
        label_visibility="collapsed",
    )

    # ════════════════════════════════════════════════════════════════════
    # MODE 1: Archive Browser
    # ════════════════════════════════════════════════════════════════════
    if browse_mode == "📦 Archivi":
        st.subheader("📦 Archivi Compressi del Dataset")
        st.caption(
            "Ogni archivio contiene documenti giuridici di una fonte specifica. "
            "Clicca per espandere e vedere i dettagli."
        )

        for file_info in DATASET_FILES:
            if file_info["type"] != "archive":
                continue

            with st.expander(
                f"{file_info['icon']} **{file_info['filename']}** — "
                f"{file_info['size_human']} · ~{file_info['est_documents']:,} documenti",
                expanded=False,
            ):
                col_a, col_b = st.columns([2, 1])
                with col_a:
                    st.markdown(f"**Categoria**: {file_info['category']}")
                    st.markdown(f"**Descrizione**: {file_info['description']}")
                    st.markdown(f"**Tipo contenuto**: {file_info['content_type']}")
                    st.markdown(f"**Source type**: `{file_info['source_type']}`")
                with col_b:
                    st.metric("Dimensione", file_info["size_human"])
                    st.metric("Documenti stimati", f"{file_info['est_documents']:,}")

                # Show linked metadata if available
                if manager and manager.metadata:
                    linked = find_metadata_for_archive(manager.metadata, file_info["filename"])
                    if linked:
                        st.success(f"✅ **{len(linked):,}** documenti linkati nell'indice FAISS")
                        with st.expander(f"📄 Anteprima documenti linkati (primi 10)", expanded=False):
                            for doc in linked[:10]:
                                urn = doc.get("urn", "N/D")
                                title = doc.get("title") or doc.get("source_title") or "Senza titolo"
                                st.markdown(f"- **{title[:100]}** — `{urn}`")
                    else:
                        st.info(
                            "📝 Nessun documento di questo archivio risulta ancora indicizzato nel FAISS index attuale. "
                            "Esegui il builder in modalità 'full_archives' per indicizzarli."
                        )

                # Linkage info
                linkage = ARCHIVE_METADATA_LINKAGE.get(file_info["source_type"])
                if linkage:
                    st.caption(f"🔗 **Linkage**: {linkage['description']}")
                    st.caption(f"Campi linkati: {', '.join(linkage['linked_fields'])}")

                st.markdown(
                    f"[📥 Download da HuggingFace]({file_info['download_url']})",
                )

        # Non-archive files
        st.divider()
        st.subheader("📄 Altri File nel Dataset")
        non_archives = [f for f in DATASET_FILES if f["type"] != "archive"]
        for file_info in non_archives:
            st.markdown(
                f"{file_info['icon']} **{file_info['filename']}** — "
                f"{file_info['size_human']} · {file_info['category']}"
            )

    # ════════════════════════════════════════════════════════════════════
    # MODE 2: Metadata ↔ Law Linkage
    # ════════════════════════════════════════════════════════════════════
    elif browse_mode == "🔗 Linkage Metadati↔Leggi":
        st.subheader("🔗 Linkage: Metadati ↔ Leggi ↔ Archivi")
        st.caption(
            "Questa vista mostra come i metadati nell'indice FAISS sono collegati "
            "ai documenti originali negli archivi compressi."
        )

        if not manager or not manager.metadata:
            st.warning("⚠️ Indice FAISS non caricato — impossibile mostrare il linkage.")
            return

        # Linkage diagram
        st.markdown("""
        ```
        ┌─────────────────────┐     ┌──────────────────────┐     ┌───────────────────────┐
        │  doc_metadata.jsonl  │────▶│   source_archive     │────▶│  normattiva.tar.zst   │
        │  (839.422 vettori)   │     │   build_source       │     │  corte_cost.tar.zst   │
        │                      │     │   archive_internal_   │     │  ANAC.tar.zst         │
        │  Campi per vettore:  │     │   path               │     │  senate.tar.zst       │
        │  - urn               │     │                      │     │  siope.tar.zst        │
        │  - title             │     │  urn ──────────▶     │     │  ministeri.tar.zst    │
        │  - text              │     │  Normattiva API      │     │                       │
        │  - validity_status   │     │  (verifica live)     │     │  rebuild_v2_docs.jsonl │
        │  - source_type       │     │                      │     │  (corpus completo)     │
        └─────────────────────┘     └──────────────────────┘     └───────────────────────┘
        ```
        """)

        # Per-archive linkage stats
        st.markdown("#### Stato Linkage per Archivio")

        for file_info in archives:
            linked = find_metadata_for_archive(manager.metadata, file_info["filename"])
            pct = (len(linked) / len(manager.metadata) * 100) if manager.metadata else 0

            col_l, col_r = st.columns([3, 1])
            with col_l:
                st.markdown(
                    f"{file_info['icon']} **{file_info['filename']}** → "
                    f"**{len(linked):,}** documenti linkati ({pct:.1f}%)"
                )
            with col_r:
                if linked:
                    st.progress(min(pct / 100, 1.0))
                else:
                    st.caption("Non indicizzato")

        # Overall linkage quality
        st.divider()
        total = len(manager.metadata) if manager.metadata else 0
        with_archive = sum(1 for d in (manager.metadata or []) if isinstance(d, dict) and d.get("source_archive"))
        with_urn = sum(1 for d in (manager.metadata or []) if isinstance(d, dict) and d.get("urn"))
        with_title = sum(1 for d in (manager.metadata or []) if isinstance(d, dict) and (d.get("title") or d.get("source_title")))

        st.markdown("#### Qualità Complessiva del Linkage")
        q1, q2, q3, q4 = st.columns(4)
        q1.metric("Totale vettori", f"{total:,}")
        q2.metric("Con source_archive", f"{with_archive:,}", f"{with_archive/total*100:.1f}%" if total else "0%")
        q3.metric("Con URN", f"{with_urn:,}", f"{with_urn/total*100:.1f}%" if total else "0%")
        q4.metric("Con titolo", f"{with_title:,}", f"{with_title/total*100:.1f}%" if total else "0%")

    # ════════════════════════════════════════════════════════════════════
    # MODE 3: Search
    # ════════════════════════════════════════════════════════════════════
    elif browse_mode == "🔍 Cerca nel Dataset":
        st.subheader("🔍 Cerca nel Dataset")

        if not manager or not manager.metadata:
            st.warning("⚠️ Indice FAISS non caricato.")
            return

        search_type = st.radio(
            "Cerca per:", ["URN", "Titolo", "Source Archive"],
            horizontal=True,
        )
        search_query = st.text_input(
            "🔎 Query di ricerca",
            placeholder="es. urn:nir:stato:legge:2020 oppure 'codice civile'",
        )

        if search_query and len(search_query) >= 3:
            with st.spinner("Cercando..."):
                if search_type == "URN":
                    results = search_metadata_by_urn(manager.metadata, search_query)
                elif search_type == "Titolo":
                    results = search_metadata_by_title(manager.metadata, search_query)
                else:
                    results = find_metadata_for_archive(manager.metadata, search_query)

            if results:
                st.success(f"Trovati **{len(results)}** risultati")

                for doc in results[:25]:
                    urn = doc.get("urn", "N/D")
                    title = doc.get("title") or doc.get("source_title") or "Senza titolo"
                    source = doc.get("source_archive") or doc.get("build_source") or "N/D"
                    faiss_idx = doc.get("faiss_idx", "?")
                    validity = doc.get("validity_status", "N/D")

                    with st.container(border=True):
                        st.markdown(f"**{title[:150]}**")
                        c_a, c_b, c_c = st.columns(3)
                        c_a.caption(f"URN: `{urn}`")
                        c_b.caption(f"Archivio: `{source}`")
                        c_c.caption(f"FAISS idx: {faiss_idx} · Vigenza: {validity}")

                        # Show which archive this came from
                        archive_info = get_archive_from_metadata(doc)
                        if archive_info:
                            st.caption(
                                f"📦 Fonte: {archive_info['icon']} {archive_info['filename']} "
                                f"({archive_info['size_human']})"
                            )

                        # Show internal path if available
                        internal_path = doc.get("archive_internal_path")
                        if internal_path:
                            st.caption(f"📂 Path interno: `{internal_path}`")

                        # Text preview
                        text = doc.get("text", "")
                        if text:
                            st.text_area(
                                "Anteprima testo",
                                text[:500] + ("..." if len(text) > 500 else ""),
                                height=100,
                                disabled=True,
                                key=f"preview_{faiss_idx}_{hash(urn)}",
                            )
            else:
                st.info("Nessun risultato trovato.")

    # ════════════════════════════════════════════════════════════════════
    # MODE 4: Source Breakdown
    # ════════════════════════════════════════════════════════════════════
    elif browse_mode == "📊 Source Breakdown":
        st.subheader("📊 Distribuzione per Fonte")

        if not manager or not manager.metadata:
            st.warning("⚠️ Indice FAISS non caricato.")
            return

        source_stats = compute_source_stats(manager.metadata)

        # Sort by count
        sorted_sources = sorted(source_stats.items(), key=lambda x: -x[1]["count"])

        for source_name, sinfo in sorted_sources:
            pct = sinfo["count"] / len(manager.metadata) * 100 if manager.metadata else 0

            with st.expander(
                f"**{source_name}** — {sinfo['count']:,} documenti ({pct:.1f}%)",
                expanded=(sinfo["count"] > 10000),
            ):
                m1, m2, m3, m4 = st.columns(4)
                m1.metric("Documenti", f"{sinfo['count']:,}")
                m2.metric("Con URN", f"{sinfo['has_urn']:,}")
                m3.metric("Con titolo", f"{sinfo['has_title']:,}")
                m4.metric("Link archivio", f"{sinfo['has_archive_link']:,}")

                if sinfo["sample_titles"]:
                    st.markdown("**Titoli di esempio:**")
                    for t in sinfo["sample_titles"]:
                        st.caption(f"• {t}")

                if sinfo["sample_urns"]:
                    st.markdown("**URN di esempio:**")
                    for u in sinfo["sample_urns"]:
                        st.code(u)

    # ── Dynamic manifest (if available) ──
    st.divider()
    dynamic_manifest = try_load_dynamic_manifest(data_dir)
    if dynamic_manifest:
        with st.expander(f"📋 Dataset Manifest Dinamico ({len(dynamic_manifest):,} entries)", expanded=False):
            st.caption(
                "Questo manifest è stato generato dal builder e contiene il catalogo "
                "completo di ogni file in ogni archivio del dataset."
            )
            # Show summary stats
            by_archive = {}
            for entry in dynamic_manifest:
                archive = entry.get("archive_file") or "top_level"
                by_archive[archive] = by_archive.get(archive, 0) + 1
            for arch, count in sorted(by_archive.items(), key=lambda x: -x[1]):
                st.caption(f"📦 {arch}: {count:,} file")
    else:
        st.caption(
            "💡 Per un catalogo più dettagliato dei file interni agli archivi, "
            "esegui il builder in modalità **manifest** per generare dataset_manifest.jsonl."
        )