from __future__ import annotations

import json
import logging
import re
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

import numpy as np

logger = logging.getLogger(__name__)

_GLOBAL_RAG: Optional["RAGRetriever"] = None


def set_global_rag(retriever: Optional["RAGRetriever"]) -> None:
    """Imposta l'istanza RAG globale condivisa dalla pipeline."""
    global _GLOBAL_RAG
    _GLOBAL_RAG = retriever


def get_global_rag() -> Optional["RAGRetriever"]:
    """Restituisce l'istanza RAG globale (None se non abilitata)."""
    return _GLOBAL_RAG


@dataclass
class Chunk:
    """Un singolo chunk indicizzato."""

    text: str
    source: str
    chunk_index: int
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class SearchResult:
    """Risultato di una query RAG."""

    chunk: Chunk
    score: float


def _split_text_recursive(
    text: str,
    chunk_size: int = 512,
    chunk_overlap: int = 64,
) -> list[str]:
    """Semplice splitter ricorsivo: paragrafi → frasi → caratteri.

    Non dipende da langchain per mantenere zero dipendenze extra oltre a
    sentence-transformers.
    """
    if chunk_size <= 0:
        raise ValueError("chunk_size must be > 0")
    if chunk_overlap >= chunk_size:
        raise ValueError("chunk_overlap must be < chunk_size")
    if len(text) <= chunk_size:
        return [text.strip()] if text.strip() else []

    separators = [
        ("\n\n", True),
        ("\n", True),
        (". ", True),
        ("? ", True),
        ("! ", True),
        ("; ", True),
        (", ", True),
        (" ", False),
        ("", False),
    ]

    def _split_with(sep: str, keep_sep: bool) -> list[str]:
        if not sep:
            return list(text)
        if keep_sep:
            parts = re.split(rf"(?<={re.escape(sep)})", text)
        else:
            parts = text.split(sep)
        return [p for p in parts if p]

    chosen_parts: Optional[list[str]] = None
    for sep, keep in separators:
        parts = _split_with(sep, keep)
        if len(parts) > 1:
            chosen_parts = parts
            break
    if chosen_parts is None:
        chosen_parts = list(text)

    chunks: list[str] = []
    current = ""
    for part in chosen_parts:
        candidate = (current + part).lstrip()
        if len(candidate) <= chunk_size:
            current = candidate
        else:
            if current:
                chunks.append(current.rstrip())
            # fallback: se il singolo pezzo è più grande di chunk_size, tronchiamo
            if len(part) > chunk_size:
                for i in range(0, len(part), chunk_size - chunk_overlap):
                    window = part[i : i + chunk_size].strip()
                    if window:
                        chunks.append(window)
                current = ""
            else:
                # inizi un nuovo chunk con overlap
                tail = ""
                if chunks and chunk_overlap > 0:
                    tail = chunks[-1][-chunk_overlap:]
                current = (tail + part).lstrip()
    if current.strip():
        chunks.append(current.rstrip())

    # merges finale per non avere chunk troppo piccoli (ultimo < 30% chunk_size)
    merged: list[str] = []
    min_len = max(20, chunk_size // 3)
    for c in chunks:
        if merged and len(c) < min_len:
            merged[-1] = (merged[-1] + " " + c).strip()
        else:
            merged.append(c)
    return merged


def _load_text_files(kb_path: Path) -> list[tuple[str, str]]:
    """Carica tutti i file .md e .txt dalla cartella (ricorsivo).

    Restituisce lista di (source_rel_path, full_text).
    """
    docs: list[tuple[str, str]] = []
    for ext in ("*.md", "*.txt"):
        for fp in sorted(kb_path.rglob(ext)):
            if fp.is_file():
                try:
                    text = fp.read_text(encoding="utf-8")
                except UnicodeDecodeError:
                    text = fp.read_text(encoding="latin-1", errors="replace")
                rel = str(fp.relative_to(kb_path))
                docs.append((rel, text))
    return docs


def _load_jsonl(kb_path: Path) -> list[tuple[str, str, dict[str, Any]]]:
    """Carica file *.jsonl (opzionale) già pre-chunkizzati.

    Riga attesa: {"text": "...", "source": "...", "chunk_index": 0 (opzionale), ...altri metadata}
    """
    prechunked: list[tuple[str, str, dict[str, Any]]] = []
    for fp in sorted(kb_path.rglob("*.jsonl")):
        if not fp.is_file():
            continue
        rel = str(fp.relative_to(kb_path))
        with fp.open("r", encoding="utf-8") as f:
            for lineno, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                except json.JSONDecodeError as e:
                    logger.warning("RAG: riga %s non valida in %s: %s", lineno, rel, e)
                    continue
                text = obj.get("text") or obj.get("content") or ""
                if not text:
                    continue
                source = obj.get("source") or f"{rel}#riga{lineno}"
                meta = {k: v for k, v in obj.items() if k not in {"text", "content", "source", "chunk_index"}}
                ci = obj.get("chunk_index", lineno - 1)
                prechunked.append((source, text, {**meta, "chunk_index": ci, "from_jsonl": True}))
    return prechunked


class RAGRetriever:
    """Retriever RAG locale basato su sentence-transformers + numpy.

    - Indicizza una cartella di file .md/.txt o un JSONL pre-chunkizzato.
    - Salva/carica l'indice NPZ per evitare re-embedding ad ogni avvio.
    - Ricerca coseno-similarità top-k con soglia minima.
    """

    def __init__(
        self,
        kb_path: str | Path,
        *,
        embedding_model: str = "paraphrase-multilingual-MiniLM-L12-v2",
        device: str = "auto",
        top_k: int = 3,
        threshold: float = 0.25,
        chunk_size: int = 512,
        chunk_overlap: int = 64,
        embedding_batch_size: int = 32,
        language: str = "es",
        inject_as: str = "system",
    ) -> None:
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError as e:
            raise RuntimeError(
                "sentence-transformers non installato. Installa le dipendenze RAG con:\n"
                "  uv pip install -e .[rag]\n"
                f"Errore originale: {e}"
            ) from e

        self.kb_path = Path(kb_path).expanduser().resolve()
        self.embedding_model_name = embedding_model
        self.device = device
        self.top_k = max(1, int(top_k))
        self.threshold = float(threshold)
        self.chunk_size = max(64, int(chunk_size))
        self.chunk_overlap = max(0, int(chunk_overlap))
        self.embedding_batch_size = max(1, int(embedding_batch_size))
        self.language = language if language in {"es", "it", "en"} else "es"
        self.inject_as = inject_as if inject_as in {"system", "user"} else "system"

        self.index_path = self.kb_path / "_index.npz"
        self.meta_path = self.kb_path / "_chunks.jsonl"
        self.dynamic_jsonl_path = self.kb_path / "_dynamic.jsonl"

        self.chunks: list[Chunk] = []
        self.embeddings: Optional[np.ndarray] = None
        self._lock = threading.RLock()
        # File dinamici sono chunk aggiunti via API durante runtime, senza riavvio (vengono
        # persistere sul disk in kb/_dynamic.jsonl tra un riavvio).
        self._dynamic_sources: set[str] = set()
        self._persist_dirty = False

        import torch

        resolved = device
        if resolved == "auto":
            if torch.cuda.is_available():
                resolved = "cuda"
            else:
                resolved = "cpu"
        logger.info(
            "RAG: Inizializzazione modello embedding=%s device=%s su kb_path=%s",
            self.embedding_model_name,
            resolved,
            self.kb_path,
        )
        self._model = SentenceTransformer(self.embedding_model_name, device=resolved)

    # ── Indexing / persistenza ───────────────────────────────────────────────

    def _index_exists(self) -> bool:
        return self.index_path.is_file() and self.meta_path.is_file()

    def _save_index(self) -> None:
        assert self.embeddings is not None
        self.kb_path.mkdir(parents=True, exist_ok=True)
        np.savez(self.index_path, embeddings=self.embeddings.astype(np.float32))
        with self.meta_path.open("w", encoding="utf-8") as f:
            for ch in self.chunks:
                f.write(
                    json.dumps(
                        {
                            "text": ch.text,
                            "source": ch.source,
                            "chunk_index": ch.chunk_index,
                            "metadata": ch.metadata,
                        },
                        ensure_ascii=False,
                    )
                    + "\n"
                )
        logger.info(
            "RAG: Indice salvato: %d chunk, shape embeddings=%s",
            len(self.chunks),
            self.embeddings.shape,
        )

    def _load_index(self) -> bool:
        if not self._index_exists():
            return False
        try:
            data = np.load(self.index_path)
            embeddings = data["embeddings"]
            chunks: list[Chunk] = []
            with self.meta_path.open("r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    obj = json.loads(line)
                    chunks.append(
                        Chunk(
                            text=obj["text"],
                            source=obj["source"],
                            chunk_index=int(obj.get("chunk_index", 0)),
                            metadata=obj.get("metadata") or {},
                        )
                    )
            if len(chunks) != embeddings.shape[0]:
                logger.warning(
                    "RAG: mismatch chunks(%d) vs embeddings(%s); ricostruisco l'indice.",
                    len(chunks),
                    embeddings.shape,
                )
                return False
            self.chunks = chunks
            self.embeddings = embeddings.astype(np.float32)
            logger.info(
                "RAG: Indice caricato da disco: %d chunk (shape=%s).",
                len(self.chunks),
                self.embeddings.shape,
            )
            return True
        except Exception as e:
            logger.warning("RAG: caricamento indice fallito (%s); ricostruisco.", e)
            return False

    def build_index(self, *, force_rebuild: bool = False) -> None:
        """Popola chunks+embeddings. Carica da NPZ se possibile, altrimenti indicizza."""
        if not force_rebuild and self._load_index():
            return

        if not self.kb_path.is_dir():
            logger.info("RAG: kb_path=%s non esiste, lo creo vuoto.", self.kb_path)
            self.kb_path.mkdir(parents=True, exist_ok=True)

        raw_docs = _load_text_files(self.kb_path)
        prechunked = _load_jsonl(self.kb_path)

        self.chunks: list[Chunk] = []
        texts_to_embed: list[str] = []

        for source, text in raw_docs:
            pieces = _split_text_recursive(
                text,
                chunk_size=self.chunk_size,
                chunk_overlap=self.chunk_overlap,
            )
            for i, piece in enumerate(pieces):
                ch = Chunk(text=piece, source=source, chunk_index=i, metadata={"splitted": True})
                self.chunks.append(ch)
                texts_to_embed.append(piece)

        for source, text, meta in prechunked:
            ci = int(meta.pop("chunk_index", 0))
            ch = Chunk(text=text, source=source, chunk_index=ci, metadata=meta)
            self.chunks.append(ch)
            texts_to_embed.append(text)

        if not self.chunks:
            logger.warning(
                "RAG: Nessun documento trovato in %s. Indice vuoto (retrieval non inietterà niente).",
                self.kb_path,
            )
            self.embeddings = np.zeros((0, 1), dtype=np.float32)
        else:
            logger.info(
                "RAG: Embedding %d chunk (batch=%d)...",
                len(texts_to_embed),
                self.embedding_batch_size,
            )
            arr = self._model.encode(
                texts_to_embed,
                batch_size=self.embedding_batch_size,
                show_progress_bar=False,
                convert_to_numpy=True,
                normalize_embeddings=True,
            )
            self.embeddings = np.asarray(arr, dtype=np.float32)
            logger.info("RAG: Embedding completato, shape=%s", self.embeddings.shape)
            try:
                self._save_index()
            except Exception as e:
                logger.warning("RAG: salvataggio indice NPZ fallito: %s", e)

        # Appendi chunk dinamici persistiti (kb/_dynamic.jsonl) dall'avvio precedente
        self._load_dynamic_chunks()

    # ── Search ───────────────────────────────────────────────────────────────

    def search(
        self,
        query: str,
        *,
        top_k: Optional[int] = None,
        threshold: Optional[float] = None,
    ) -> list[SearchResult]:
        """Cerca i chunk più rilevanti per *query*.

        Thread-safe: prende lock in lettura.
        Restituisce lista ordinata per score decrescente, filtrata per soglia.
        Lista vuota se l'indice è vuoto o nessun chunk supera la soglia.
        """
        with self._lock:
            if self.embeddings is None or self.embeddings.shape[0] == 0:
                return []
            if not query or not query.strip():
                return []

            k = self.top_k if top_k is None else max(1, int(top_k))
            thr = self.threshold if threshold is None else float(threshold)
            k = min(k, self.embeddings.shape[0])

            q_vec = self._model.encode(
                [query],
                batch_size=1,
                show_progress_bar=False,
                convert_to_numpy=True,
                normalize_embeddings=True,
            )
            q_vec = np.asarray(q_vec, dtype=np.float32).reshape(1, -1)

            scores = (self.embeddings @ q_vec.T).reshape(-1)
            if k == self.embeddings.shape[0]:
                top_idx = np.argsort(-scores)
            else:
                top_idx = np.argpartition(-scores, k - 1)[:k]
                top_idx = top_idx[np.argsort(-scores[top_idx])]

            out: list[SearchResult] = []
            for idx in top_idx:
                sc = float(scores[idx])
                if sc < thr:
                    break
                out.append(SearchResult(chunk=self.chunks[int(idx)], score=sc))
            return out

    # ── Dynamic updates (thread-safe) ────────────────────────────────────────

    def status(self) -> dict[str, Any]:
        """Restituisce stato interno (per endpoint HTTP /status)."""
        with self._lock:
            n = self.embeddings.shape[0] if self.embeddings is not None else 0
            dim = int(self.embeddings.shape[1]) if self.embeddings is not None and n > 0 else 0
            return {
                "kb_path": str(self.kb_path),
                "index_path": str(self.index_path),
                "dynamic_path": str(self.dynamic_jsonl_path),
                "embedding_model": self.embedding_model_name,
                "device": self.device,
                "num_chunks": n,
                "embedding_dim": dim,
                "top_k": self.top_k,
                "threshold": self.threshold,
                "language": self.language,
                "inject_as": self.inject_as,
                "dynamic_sources_count": len(self._dynamic_sources),
                "sources": sorted({c.source for c in self.chunks}),
            }

    def add_chunks(
        self,
        items: list[dict[str, Any]],
        *,
        persist: bool = True,
        dynamic: bool = True,
    ) -> dict[str, Any]:
        """Aggiunge chunk pre-costruiti in modo incrementale.

        Parameters
        ----------
        items : list[dict]
            Lista di dict con chiavi:
              - ``text`` (str, obbligatoria)
              - ``source`` (str, opzionale, default='dynamic#N')
              - ``chunk_index`` (int, opzionale)
              - ``metadata`` (dict, opzionale)
        persist : bool
            Se True, scrive i nuovi chunk anche in kb/_dynamic.jsonl per il prossimo
            avvio.
        dynamic : bool
            Se True, marca questi chunk come 'dinamici' (verranno ignorati da
            remove_by_source solo se match esplicito).

        Returns
        -------
        dict con:
            ``added`` : int — numero chunk aggiunti
            ``sources`` : list[str] — sorgenti uniche aggiunte
            ``total_after`` : int — dimensione totale dell'indice dopo l'add
        """
        new_chunks: list[Chunk] = []
        new_texts: list[str] = []
        added_sources: set[str] = set()
        with self._lock:
            existing_count = len(self.chunks)
            next_auto = existing_count + 1000000
            for raw in items:
                text = (raw.get("text") or "").strip()
                if not text:
                    continue
                source = (raw.get("source") or "").strip() or f"dynamic#{existing_count + len(new_chunks)}"
                ci = int(raw.get("chunk_index") or (next_auto + len(new_chunks)))
                meta = dict(raw.get("metadata") or {})
                if dynamic:
                    meta["_dynamic"] = True
                ch = Chunk(text=text, source=source, chunk_index=ci, metadata=meta)
                new_chunks.append(ch)
                new_texts.append(text)
                added_sources.add(source)
                if dynamic:
                    self._dynamic_sources.add(source)
            if not new_chunks:
                return {"added": 0, "sources": [], "total_after": len(self.chunks)}

            logger.info(
                "RAG add_chunks: aggiungo %d chunk (batch=%d)...",
                len(new_chunks),
                self.embedding_batch_size,
            )
            new_embs = self._model.encode(
                new_texts,
                batch_size=self.embedding_batch_size,
                show_progress_bar=False,
                convert_to_numpy=True,
                normalize_embeddings=True,
            )
            new_embs = np.asarray(new_embs, dtype=np.float32)
            # Concatenazione safe: gestisce caso embeddings vuoto (0, 1)
            if self.embeddings is None or self.embeddings.size == 0 or self.embeddings.shape[1] == 1:
                self.embeddings = new_embs
            else:
                if new_embs.shape[1] != self.embeddings.shape[1]:
                    raise ValueError(
                        f"Dimensione embedding mismatch: nuovi {new_embs.shape[1]} vs esistenti {self.embeddings.shape[1]}"
                    )
                self.embeddings = np.vstack([self.embeddings, new_embs])
            self.chunks.extend(new_chunks)
            total_after = len(self.chunks)

        # Persisti fuori dal lock (non blocca retrieval)
        if persist:
            self._append_dynamic_chunks_to_disk(new_chunks)
            try:
                with self._lock:
                    self._save_index()
            except Exception as exc:
                logger.warning("RAG add_chunks: salvataggio indice fallito: %s", exc)

        logger.info(
            "RAG add_chunks OK: added=%d total_after=%d sources=%s",
            len(new_chunks),
            total_after,
            sorted(added_sources),
        )
        return {
            "added": len(new_chunks),
            "sources": sorted(added_sources),
            "total_after": total_after,
        }

    def add_document(
        self,
        text: str,
        *,
        source: str,
        metadata: Optional[dict[str, Any]] = None,
        persist: bool = True,
        dynamic: bool = True,
    ) -> dict[str, Any]:
        """Splitta un documento in chunk e li aggiunge in modo incrementale.

        Equivalente a:
          1. ``text`` → chunking ricorsivo (stesso algoritmo di build_index)
          2. ``add_chunks()`` di tutti i pezzi

        Parameters
        ----------
        text : str
          Testo completo documento.
        source : str
          Nome identificativo (es. ``crm/cliente_123_nota_20260826``).
        metadata : dict, opzionale
          Mappa passata a tutti i chunk derivati.
        persist / dynamic : vedi ``add_chunks``.
        """
        text = (text or "").strip()
        if not text:
            n = 0
            with self._lock:
                if self.embeddings is not None:
                    n = self.embeddings.shape[0]
            return {"added": 0, "sources": [], "total_after": n}
        pieces = _split_text_recursive(
            text,
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
        )
        meta = dict(metadata or {})
        items = [
            {"text": p, "source": source, "chunk_index": i, "metadata": {**meta}}
            for i, p in enumerate(pieces)
        ]
        return self.add_chunks(items, persist=persist, dynamic=dynamic)

    def remove_by_source(
        self,
        source_prefix: str,
        *,
        exact: bool = False,
        persist: bool = True,
    ) -> dict[str, Any]:
        """Rimuove chunk selezionati in base al nome sorgente.

        Parameters
        ----------
        source_prefix : str
          Prefisso o nome esterno (vedi ``exact``).
        exact : bool
          - ``False`` (default): rimuove TUTTI i chunk la cui ``source``
            **inizia** con ``source_prefix``.
          - ``True``: rimuove **solo** i chunk la cui ``source`` è
            **esattamente uguale** a ``source_prefix``.
        persist : bool
          Se True, dopo la rimozione riscrive ``_dynamic.jsonl`` e
          ``_index.npz``.

        Casi d'uso:
          - ``remove_by_source('crm/cliente_123')`` → tutte le note di un cliente
          - ``remove_by_source('dynamic#')`` → tutti i chunk temporanei auto-numerati
          - ``remove_by_source('crm/cliente_456_nota_20260826', exact=True)`` → solo quel documento
        """
        removed_indices: list[int] = []
        removed_sources: set[str] = set()
        with self._lock:
            new_chunks: list[Chunk] = []
            kept_indices: list[int] = []
            for i, c in enumerate(self.chunks):
                match = (c.source == source_prefix) if exact else c.source.startswith(source_prefix)
                if match:
                    removed_indices.append(i)
                    removed_sources.add(c.source)
                    if c.source in self._dynamic_sources:
                        self._dynamic_sources.discard(c.source)
                else:
                    new_chunks.append(c)
                    kept_indices.append(i)
            if not removed_indices:
                return {
                    "removed": 0,
                    "sources": [],
                    "total_after": len(self.chunks),
                }
            if self.embeddings is not None and self.embeddings.shape[0] > 0 and not (
                self.embeddings.shape[0] == 1 and self.embeddings.shape[1] == 1
            ):
                kept = np.asarray(kept_indices, dtype=np.int64)
                self.embeddings = self.embeddings.take(kept, axis=0)
            self.chunks = new_chunks
            total_after = len(self.chunks)

        if persist:
            try:
                with self._lock:
                    self._save_index()
                    self._rewrite_dynamic_chunks_on_disk()
            except Exception as exc:
                logger.warning("RAG remove: salvataggio dopo rimozione fallito: %s", exc)

        logger.info(
            "RAG remove_by_source(prefix=%r, exact=%s): rimossi %d chunk, %d sorgenti uniche",
            source_prefix,
            exact,
            len(removed_indices),
            len(removed_sources),
        )
        return {
            "removed": len(removed_indices),
            "sources": sorted(removed_sources),
            "total_after": total_after,
            "mode": "exact" if exact else "prefix",
        }

    def list_sources(self) -> list[dict[str, Any]]:
        """Lista di tutte le sorgenti uniche nell'indice con il conteggio chunk.

        Restituisce una lista di ``[{source, chunk_count, has_dynamic}]``
        ordinata per ``source`` alfabeticamente.
        """
        with self._lock:
            agg: dict[str, dict[str, Any]] = {}
            for c in self.chunks:
                entry = agg.get(c.source)
                if entry is None:
                    entry = {"source": c.source, "chunk_count": 0, "has_dynamic": False}
                    agg[c.source] = entry
                entry["chunk_count"] += 1
                if c.metadata.get("_dynamic", False):
                    entry["has_dynamic"] = True
            return sorted(agg.values(), key=lambda d: d["source"])

    def list_chunks(
        self,
        *,
        source_prefix: Optional[str] = None,
        source_exact: Optional[str] = None,
        query: Optional[str] = None,
        min_score: Optional[float] = None,
        offset: int = 0,
        limit: int = 100,
        include_text: bool = True,
        include_embedding: bool = False,
    ) -> dict[str, Any]:
        """Elenca chunk con filtri e paginazione.

        Parameters
        ----------
        source_prefix / source_exact : str, opzionale
          Filtro per nome sorgente (solo uno dei due, ``exact`` vince se entrambi).
        query : str, opzionale
          Se fornito, ordina i risultati per score di rilevanza coseno rispetto a
          questa frase (stesso modello embedding di ``search()``).
        min_score : float, opzionale
          Solo per modalità ``query``: filtra chunk con score >= soglia.
        offset / limit : int
          Paginazione. ``limit`` massimo forzato a 1000.
        include_text : bool
          Includi ``chunk.text`` nella risposta (default True).
        include_embedding : bool
          Includi ``chunk.embedding`` (lista float). Default False (risparmia banda).

        Returns
        -------
        dict con chiavi ``total`` (conteggio prima della paginazione), ``offset``,
        ``limit``, ``items`` (lista chunk dict).
        """
        with self._lock:
            # 1. Filtro sorgente
            candidates: list[tuple[int, Chunk]] = []
            for i, c in enumerate(self.chunks):
                if source_exact is not None:
                    if c.source != source_exact:
                        continue
                elif source_prefix is not None:
                    if not c.source.startswith(source_prefix):
                        continue
                candidates.append((i, c))

            # 2. Query / ordinamento per rilevanza
            if query and query.strip():
                q_vec = self._model.encode(
                    [query],
                    batch_size=1,
                    show_progress_bar=False,
                    convert_to_numpy=True,
                    normalize_embeddings=True,
                )
                q_vec = np.asarray(q_vec, dtype=np.float32).reshape(1, -1)
                if self.embeddings is not None and self.embeddings.shape[0] > 0 and self.embeddings.shape[1] > 1:
                    all_scores = (self.embeddings @ q_vec.T).reshape(-1)
                    scored: list[tuple[int, Chunk, float]] = []
                    for i, c in candidates:
                        sc = float(all_scores[i])
                        if min_score is None or sc >= min_score:
                            scored.append((i, c, sc))
                    scored.sort(key=lambda t: t[2], reverse=True)
                else:
                    scored = [(i, c, 0.0) for i, c in candidates]
            else:
                scored = [(i, c, None) for i, c in candidates]
                # Ordine stabile per source + chunk_index
                scored.sort(key=lambda t: (t[1].source, t[1].chunk_index))

            total = len(scored)
            offset = max(0, int(offset))
            limit = max(1, min(1000, int(limit)))
            page = scored[offset : offset + limit]

            items: list[dict[str, Any]] = []
            for i, c, sc in page:
                entry: dict[str, Any] = {
                    "index": i,
                    "source": c.source,
                    "chunk_index": c.chunk_index,
                    "score": sc,
                    "metadata": {k: v for k, v in c.metadata.items() if k != "_dynamic"} if c.metadata else {},
                    "is_dynamic": bool(c.metadata.get("_dynamic", False)),
                }
                if include_text:
                    entry["text"] = c.text
                if include_embedding and self.embeddings is not None and i < self.embeddings.shape[0]:
                    entry["embedding"] = [float(x) for x in self.embeddings[i].tolist()]
                items.append(entry)
            return {
                "total": total,
                "offset": offset,
                "limit": limit,
                "query": query,
                "source_prefix": source_prefix,
                "source_exact": source_exact,
                "items": items,
            }

    def update_chunk(
        self,
        *,
        index: Optional[int] = None,
        source: Optional[str] = None,
        chunk_index: Optional[int] = None,
        new_text: Optional[str] = None,
        new_metadata: Optional[dict[str, Any]] = None,
        new_source: Optional[str] = None,
        persist: bool = True,
    ) -> dict[str, Any]:
        """Aggiorna un singolo chunk (testo, metadata, sorgente).

        Il chunk viene identificato tramite **una sola** di queste strategie
        (la prima che matcha in ordine):

        1. ``index`` — posizione intera nell'array ``self.chunks`` (solo per
           sistemi che conoscono l'indice da ``list_chunks()``).
        2. ``(source, chunk_index)`` — chiave business: sorgente + indice chunk interno.
        3. ``source`` da sola — solo se esiste UN solo chunk con quella source
           (altrimenti errore ``AMBUGUOUS_SOURCE``).

        Returns
        -------
        dict con ``updated`` (0/1), ``chunk`` (dati aggiornati senza embedding),
        ``error`` (se presente).
        """
        # Trova indice
        found_idx: Optional[int] = None
        with self._lock:
            if index is not None:
                if 0 <= int(index) < len(self.chunks):
                    found_idx = int(index)
            elif source is not None and chunk_index is not None:
                for i, c in enumerate(self.chunks):
                    if c.source == source and c.chunk_index == int(chunk_index):
                        found_idx = i
                        break
            elif source is not None:
                matches = [i for i, c in enumerate(self.chunks) if c.source == source]
                if len(matches) == 1:
                    found_idx = matches[0]
                elif len(matches) > 1:
                    return {
                        "updated": 0,
                        "error": "AMBIGUOUS_SOURCE",
                        "message": f"La source {source!r} ha {len(matches)} chunk — specifica anche chunk_index o usa index.",
                        "matches": [self.chunks[i].chunk_index for i in matches],
                    }
            if found_idx is None:
                return {"updated": 0, "error": "CHUNK_NOT_FOUND", "message": "Nessun chunk individuato con i parametri forniti."}

            old = self.chunks[found_idx]
            new_text_clean = (new_text or "").strip() or None

            # Calcola nuova source/chunk_index
            resulting_source = new_source if (new_source is not None) else old.source
            resulting_chunk_index = old.chunk_index
            resulting_metadata = dict(new_metadata) if new_metadata is not None else dict(old.metadata)
            # Mantieni flag _dynamic se era dinamico
            if old.metadata.get("_dynamic", False):
                resulting_metadata["_dynamic"] = True
            elif resulting_source in self._dynamic_sources:
                resulting_metadata["_dynamic"] = True

            # Aggiornamento testo → re-embedding
            if new_text_clean is not None and new_text_clean != old.text:
                vec = self._model.encode(
                    [new_text_clean],
                    batch_size=1,
                    show_progress_bar=False,
                    convert_to_numpy=True,
                    normalize_embeddings=True,
                )
                vec = np.asarray(vec, dtype=np.float32).reshape(1, -1)
                if self.embeddings is not None and self.embeddings.shape[0] > 0 and self.embeddings.shape[1] == vec.shape[1]:
                    self.embeddings[found_idx] = vec[0]
                else:
                    # fallback: ricalcola tutto se shape cambiato (caso estremo)
                    logger.warning("RAG update_chunk: shape embedding incompatibile, ricostruisco la riga")
                    if self.embeddings is None or self.embeddings.size == 0:
                        self.embeddings = np.zeros((len(self.chunks), vec.shape[1]), dtype=np.float32)
                    else:
                        new_mat = np.zeros((len(self.chunks), vec.shape[1]), dtype=np.float32)
                        n = min(self.embeddings.shape[0], len(self.chunks))
                        m = min(self.embeddings.shape[1], vec.shape[1])
                        new_mat[:n, :m] = self.embeddings[:n, :m]
                        self.embeddings = new_mat
                    self.embeddings[found_idx] = vec[0]

                self.chunks[found_idx] = Chunk(
                    text=new_text_clean,
                    source=resulting_source,
                    chunk_index=resulting_chunk_index,
                    metadata=resulting_metadata,
                )
            else:
                # Solo metadata/source (nessun re-embedding)
                self.chunks[found_idx] = Chunk(
                    text=old.text,
                    source=resulting_source,
                    chunk_index=resulting_chunk_index,
                    metadata=resulting_metadata,
                )
            updated_chunk = self.chunks[found_idx]
            if resulting_source in self._dynamic_sources or updated_chunk.metadata.get("_dynamic", False):
                self._dynamic_sources.add(resulting_source)
                if new_source is not None and new_source != old.source and old.source in self._dynamic_sources:
                    # vecchia source non viene rimossa da _dynamic_sources se
                    # ci sono altri chunk con quella source (verificato dopo persist)
                    pass

            saved = {
                "index": found_idx,
                "source": updated_chunk.source,
                "chunk_index": updated_chunk.chunk_index,
                "text": updated_chunk.text,
                "metadata": {k: v for k, v in updated_chunk.metadata.items() if k != "_dynamic"},
                "is_dynamic": bool(updated_chunk.metadata.get("_dynamic", False)),
            }

        if persist:
            try:
                with self._lock:
                    self._save_index()
                    self._rewrite_dynamic_chunks_on_disk()
            except Exception as exc:
                logger.warning("RAG update_chunk: salvataggio fallito: %s", exc)

        logger.info(
            "RAG update_chunk OK: idx=%d source=%s text_changed=%s",
            found_idx,
            updated_chunk.source,
            new_text_clean is not None,
        )
        return {"updated": 1, "chunk": saved}

    def upsert_document(
        self,
        text: str,
        *,
        source: str,
        metadata: Optional[dict[str, Any]] = None,
        persist: bool = True,
        dynamic: bool = True,
    ) -> dict[str, Any]:
        """Aggiorna in modo **atomico** un intero documento identificato da ``source``.

        Operazioni (tutte dentro lo stesso lock):
          1. Rimuovi tutti i chunk esistenti con ``source`` esatto
          2. Splitta ``text`` con lo stesso chunker ricorsivo
          3. Aggiungi i nuovi chunk con embedding e metadata
        """
        text_clean = (text or "").strip()
        source = (source or "").strip()
        if not source:
            return {"error": "SOURCE_EMPTY", "message": "source è obbligatoria."}
        with self._lock:
            # ── 1. Rimuovi per source exact (SENZA persist, facciamo tutto alla fine)
            removed_count = 0
            kept_indices: list[int] = []
            for i, c in enumerate(self.chunks):
                if c.source == source:
                    removed_count += 1
                else:
                    kept_indices.append(i)
            if removed_count > 0:
                if self.embeddings is not None and self.embeddings.shape[0] > 0 and not (
                    self.embeddings.shape[0] == 1 and self.embeddings.shape[1] == 1
                ):
                    kept = np.asarray(kept_indices, dtype=np.int64)
                    self.embeddings = self.embeddings.take(kept, axis=0)
                self.chunks = [self.chunks[i] for i in kept_indices]

            if not text_clean:
                # Testo vuoto = delete atomico
                if persist:
                    try:
                        self._save_index()
                        self._rewrite_dynamic_chunks_on_disk()
                    except Exception as exc:
                        logger.warning("RAG upsert(delete): salvataggio fallito: %s", exc)
                logger.info("RAG upsert_document %r → cancellato (testo vuoto). rimossi=%d", source, removed_count)
                return {"removed": removed_count, "added": 0, "sources": [source], "total_after": len(self.chunks), "mode": "delete"}

            # ── 2. + 3. Split + add (tutto dentro lo stesso lock così atomico)
            pieces = _split_text_recursive(
                text_clean,
                chunk_size=self.chunk_size,
                chunk_overlap=self.chunk_overlap,
            )
            new_chunks: list[Chunk] = []
            new_texts: list[str] = []
            meta = dict(metadata or {})
            if dynamic:
                meta["_dynamic"] = True
                self._dynamic_sources.add(source)
            for i, p in enumerate(pieces):
                ch = Chunk(text=p, source=source, chunk_index=i, metadata={**meta})
                new_chunks.append(ch)
                new_texts.append(p)

            logger.info(
                "RAG upsert_document %r: rimossi=%d, aggiungo=%d chunk...",
                source,
                removed_count,
                len(new_chunks),
            )
            new_embs = self._model.encode(
                new_texts,
                batch_size=self.embedding_batch_size,
                show_progress_bar=False,
                convert_to_numpy=True,
                normalize_embeddings=True,
            )
            new_embs = np.asarray(new_embs, dtype=np.float32)
            if self.embeddings is None or self.embeddings.size == 0 or self.embeddings.shape[1] == 1:
                self.embeddings = new_embs
            else:
                if new_embs.shape[1] != self.embeddings.shape[1]:
                    raise ValueError(
                        f"Dimensione embedding mismatch: nuovi {new_embs.shape[1]} vs esistenti {self.embeddings.shape[1]}"
                    )
                self.embeddings = np.vstack([self.embeddings, new_embs])
            self.chunks.extend(new_chunks)
            total_after = len(self.chunks)

        if persist:
            try:
                with self._lock:
                    self._save_index()
                    self._rewrite_dynamic_chunks_on_disk()
            except Exception as exc:
                logger.warning("RAG upsert_document: salvataggio fallito: %s", exc)

        logger.info(
            "RAG upsert_document OK %s: removed=%d added=%d total_after=%d",
            source,
            removed_count,
            len(new_chunks),
            total_after,
        )
        return {
            "removed": removed_count,
            "added": len(new_chunks),
            "sources": [source],
            "total_after": total_after,
            "mode": "upsert",
        }

    def reload_from_disk(self, *, force_rebuild: bool = True) -> dict[str, Any]:
        """Rileggere TUTTI i file md/txt/jsonl dalla cartella kb/ e ricostruisce indice.

        Equivalente a riavviare con ``--rag_force_rebuild`` ma a runtime, senza spegnere.
        """
        logger.info("RAG reload_from_disk: ricostruzione indice (force_rebuild=%s)...", force_rebuild)
        self.build_index(force_rebuild=True)
        with self._lock:
            n = len(self.chunks)
            dim = int(self.embeddings.shape[1]) if self.embeddings is not None and n > 0 else 0
        logger.info("RAG reload_from_disk OK: %d chunk, dim=%d", n, dim)
        return {"chunks": n, "embedding_dim": dim}

    # ── persistence internals for dynamic chunks ──────────────────────────

    def _load_dynamic_chunks(self) -> None:
        """Avvio: carica kb/_dynamic.jsonl, embedding e appende all'indice esistente."""
        if not self.dynamic_jsonl_path.is_file():
            return
        try:
            items: list[dict[str, Any]] = []
            with self.dynamic_jsonl_path.open("r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        obj = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    items.append(obj)
            if not items:
                return
            res = self.add_chunks(items, persist=False, dynamic=True)
            logger.info(
                "RAG _load_dynamic_chunks: caricati %d chunk da %s (added=%d)",
                len(items),
                self.dynamic_jsonl_path.name,
                res.get("added", 0),
            )
        except Exception as exc:
            logger.warning("RAG _load_dynamic_chunks fallito: %s", exc)

    def _append_dynamic_chunks_to_disk(self, new_chunks: list[Chunk]) -> None:
        try:
            self.kb_path.mkdir(parents=True, exist_ok=True)
            with self.dynamic_jsonl_path.open("a", encoding="utf-8") as f:
                for ch in new_chunks:
                    f.write(
                        json.dumps(
                            {
                                "text": ch.text,
                                "source": ch.source,
                                "chunk_index": ch.chunk_index,
                                "metadata": ch.metadata,
                            },
                            ensure_ascii=False,
                        )
                        + "\n"
                    )
        except Exception as exc:
            logger.warning("RAG append _dynamic.jsonl fallito: %s", exc)

    def _rewrite_dynamic_chunks_on_disk(self) -> None:
        """Riscrittura completa di _dynamic.jsonl (chiamato dopo remove)."""
        try:
            with self.dynamic_jsonl_path.open("w", encoding="utf-8") as f:
                for ch in self.chunks:
                    if not ch.metadata.get("_dynamic", False):
                        continue
                    f.write(
                        json.dumps(
                            {
                                "text": ch.text,
                                "source": ch.source,
                                "chunk_index": ch.chunk_index,
                                "metadata": ch.metadata,
                            },
                            ensure_ascii=False,
                        )
                        + "\n"
                    )
        except Exception as exc:
            logger.warning("RAG rewrite _dynamic.jsonl fallito: %s", exc)

    # ── Helpers per il prompt injection ──────────────────────────────────────

    @staticmethod
    def format_results(results: list[SearchResult], *, language: str = "es") -> str:
        """Formatta i risultati retrieval in un blocco testuale per il system prompt.

        Lingua supportate per l'intestazione: 'es' (spagnolo, default), 'it', 'en'.
        """
        if not results:
            return ""
        headers = {
            "es": "Fragmentos relevantes recuperados de la base de conocimientos (utilízalos solo si responden directamente a la pregunta del usuario):",
            "it": "Frammenti rilevanti recuperati dalla knowledge base (usali solo se rispondono direttamente alla domanda dell'utente):",
            "en": "Relevant snippets retrieved from the knowledge base (use them only if they directly answer the user's question):",
        }
        hdr = headers.get(language, headers["es"])
        lines = [hdr]
        for i, r in enumerate(results, 1):
            src = r.chunk.source or "sconosciuto"
            lines.append(f"[{i}] (source={src}, score={r.score:.3f})")
            lines.append(r.chunk.text.strip())
        return "\n".join(lines)
