t2a/src/viewer/referentiels.py

"""Gestionnaire de référentiels utilisateur pour le RAG."""

from __future__ import annotations

import json
import logging
import shutil
import uuid
from datetime import datetime
from pathlib import Path

from ..config import REFERENTIELS_DIR, ALLOWED_EXTENSIONS, UPLOAD_MAX_SIZE_MB

logger = logging.getLogger(__name__)


class ReferentielManager:
    """CRUD pour les fichiers de référentiels utilisateur.

    Stocke les fichiers dans REFERENTIELS_DIR avec un index.json
    pour les métadonnées.
    """

    def __init__(self, referentiels_dir: Path | None = None):
        self._dir = referentiels_dir or REFERENTIELS_DIR
        self._dir.mkdir(parents=True, exist_ok=True)
        self._index_path = self._dir / "index.json"
        self._index: list[dict] = self._load_index()

    def _load_index(self) -> list[dict]:
        if self._index_path.exists():
            try:
                return json.loads(self._index_path.read_text(encoding="utf-8"))
            except (json.JSONDecodeError, KeyError):
                logger.warning("Index référentiels corrompu, réinitialisé")
        return []

    def _save_index(self) -> None:
        self._index_path.write_text(
            json.dumps(self._index, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    def list_all(self) -> list[dict]:
        """Retourne la liste de tous les référentiels."""
        return list(self._index)

    def get(self, ref_id: str) -> dict | None:
        """Retourne un référentiel par son ID."""
        for ref in self._index:
            if ref["id"] == ref_id:
                return ref
        return None

    def add_file(self, filename: str, file_data: bytes) -> dict:
        """Ajoute un fichier de référentiel.

        Args:
            filename: Nom original du fichier.
            file_data: Contenu binaire du fichier.

        Returns:
            Métadonnées du référentiel créé.

        Raises:
            ValueError: Extension non autorisée ou taille dépassée.
        """
        ext = Path(filename).suffix.lower()
        if ext not in ALLOWED_EXTENSIONS:
            raise ValueError(f"Extension '{ext}' non autorisée. Extensions valides : {ALLOWED_EXTENSIONS}")

        size_mb = len(file_data) / (1024 * 1024)
        if size_mb > UPLOAD_MAX_SIZE_MB:
            raise ValueError(f"Fichier trop volumineux ({size_mb:.1f} Mo > {UPLOAD_MAX_SIZE_MB} Mo)")

        ref_id = uuid.uuid4().hex[:12]
        safe_name = f"{ref_id}_{Path(filename).stem}{ext}"
        file_path = self._dir / safe_name

        file_path.write_bytes(file_data)

        ref = {
            "id": ref_id,
            "filename": filename,
            "stored_name": safe_name,
            "extension": ext,
            "size_bytes": len(file_data),
            "date_added": datetime.now().isoformat(),
            "status": "uploaded",
            "chunks_count": 0,
        }
        self._index.append(ref)
        self._save_index()

        logger.info("Référentiel ajouté : %s (%s)", filename, ref_id)
        return ref

    def remove(self, ref_id: str) -> bool:
        """Supprime un référentiel (fichier + métadonnées).

        Returns:
            True si trouvé et supprimé, False sinon.
        """
        ref = self.get(ref_id)
        if not ref:
            return False

        file_path = self._dir / ref["stored_name"]
        if file_path.exists():
            file_path.unlink()

        self._index = [r for r in self._index if r["id"] != ref_id]
        self._save_index()

        logger.info("Référentiel supprimé : %s (%s)", ref["filename"], ref_id)
        return True

    def index_referentiel(self, ref_id: str) -> int:
        """Indexe un référentiel dans FAISS.

        Args:
            ref_id: ID du référentiel à indexer.

        Returns:
            Nombre de chunks indexés.

        Raises:
            ValueError: Référentiel introuvable.
        """
        ref = self.get(ref_id)
        if not ref:
            raise ValueError(f"Référentiel {ref_id} introuvable")

        file_path = self._dir / ref["stored_name"]
        if not file_path.exists():
            raise ValueError(f"Fichier {ref['stored_name']} introuvable")

        from ..medical.rag_index import chunk_user_file, add_chunks_to_index

        doc_name = f"ref:{ref['filename']}"
        chunks = chunk_user_file(file_path, doc_name)

        if not chunks:
            ref["status"] = "empty"
            ref["chunks_count"] = 0
            self._save_index()
            return 0

        count = add_chunks_to_index(chunks)
        ref["status"] = "indexed"
        ref["chunks_count"] = count
        self._save_index()

        logger.info("Référentiel indexé : %s → %d chunks", ref["filename"], count)
        return count