backend/hf_storage.py

"""
Module de sauvegarde des annotations sur HuggingFace Dataset
Permet une persistance permanente même si le Space redémarre
"""

import json
import os
from datetime import datetime
from pathlib import Path
import tempfile
import streamlit as st


def get_hf_config():
    """
    Récupère la configuration HuggingFace depuis les secrets ou variables d'env.

    Returns:
        (hf_token, dataset_repo) ou (None, None) si non configuré
    """
    hf_token = None
    dataset_repo = None

    # Essayer d'abord st.secrets (HF Spaces)
    try:
        hf_token = st.secrets.get("HF_TOKEN")
        dataset_repo = st.secrets.get("HF_DATASET_REPO")
    except (FileNotFoundError, KeyError):
        pass

    # Fallback sur variables d'environnement
    if not hf_token:
        hf_token = os.getenv("HF_TOKEN")
    if not dataset_repo:
        dataset_repo = os.getenv("HF_DATASET_REPO")

    return hf_token, dataset_repo


def is_hf_storage_enabled():
    """
    Vérifie si la sauvegarde HF est configurée et disponible.

    Returns:
        bool: True si configuré, False sinon
    """
    hf_token, dataset_repo = get_hf_config()
    return hf_token is not None and dataset_repo is not None


def save_annotations_to_hf(
    annotator_id,
    annotator_name,
    feedback_scores,
    feedback_comments,
    dataset_metadata=None
):
    """
    Sauvegarde les annotations sur un Dataset HuggingFace.

    Args:
        annotator_id: ID de l'annotateur
        annotator_name: Nom de l'annotateur
        feedback_scores: Dict des scores {item_id: score}
        feedback_comments: Dict des commentaires {item_id: comment}
        dataset_metadata: Metadata du dataset annoté (optionnel)

    Returns:
        bool: True si succès, False sinon
    """
    try:
        from huggingface_hub import HfApi

        hf_token, dataset_repo = get_hf_config()

        if not hf_token or not dataset_repo:
            st.error("❌ Configuration HuggingFace manquante (HF_TOKEN ou HF_DATASET_REPO)")
            return False

        # Préparer les données d'annotation
        timestamp = datetime.now().isoformat()
        annotation_data = {
            "annotator_id": annotator_id,
            "annotator_name": annotator_name,
            "timestamp": timestamp,
            "scores": feedback_scores,
            "comments": feedback_comments,
            "num_scored": len(feedback_scores),
            "metadata": dataset_metadata or {}
        }

        # Créer un fichier temporaire
        filename = f"annotation_{annotator_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"

        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, encoding='utf-8') as f:
            json.dump(annotation_data, f, indent=2, ensure_ascii=False)
            temp_path = f.name

        try:
            # Upload vers le Dataset HF
            api = HfApi(token=hf_token)

            # Créer le repo s'il n'existe pas
            try:
                api.create_repo(
                    repo_id=dataset_repo,
                    repo_type="dataset",
                    private=True,
                    exist_ok=True
                )
            except Exception as e:
                # Le repo existe déjà, c'est OK
                pass

            # Upload le fichier
            api.upload_file(
                path_or_fileobj=temp_path,
                path_in_repo=f"annotations/{filename}",
                repo_id=dataset_repo,
                repo_type="dataset",
                commit_message=f"Add annotations from {annotator_name} ({annotator_id})"
            )

            return True

        finally:
            # Nettoyer le fichier temporaire
            try:
                os.unlink(temp_path)
            except:
                pass

    except ImportError:
        st.error("❌ Package 'huggingface_hub' non installé")
        return False
    except Exception as e:
        st.error(f"❌ Erreur lors de la sauvegarde HF: {str(e)}")
        return False


def load_annotations_from_hf(annotator_id):
    """
    Charge les annotations d'un annotateur depuis le Dataset HF.

    Args:
        annotator_id: ID de l'annotateur

    Returns:
        dict ou None: Données d'annotation si trouvées, None sinon
    """
    try:
        from huggingface_hub import HfApi, hf_hub_download

        hf_token, dataset_repo = get_hf_config()

        if not hf_token or not dataset_repo:
            return None

        api = HfApi(token=hf_token)

        # Lister les fichiers dans le repo
        try:
            files = api.list_repo_files(repo_id=dataset_repo, repo_type="dataset")
        except Exception:
            # Le repo n'existe pas encore
            return None

        # Chercher les fichiers de cet annotateur
        annotation_files = [
            f for f in files
            if f.startswith(f"annotations/annotation_{annotator_id}_")
        ]

        if not annotation_files:
            return None

        # Prendre le plus récent
        latest_file = sorted(annotation_files)[-1]

        # Télécharger et charger
        local_path = hf_hub_download(
            repo_id=dataset_repo,
            filename=latest_file,
            repo_type="dataset",
            token=hf_token
        )

        with open(local_path, 'r', encoding='utf-8') as f:
            return json.load(f)

    except ImportError:
        return None
    except Exception as e:
        st.warning(f"⚠️ Impossible de charger les annotations HF: {str(e)}")
        return None


def get_all_annotator_files(annotator_id=None):
    """
    Liste tous les fichiers d'annotations (optionnellement pour un annotateur).

    Args:
        annotator_id: ID de l'annotateur (None = tous)

    Returns:
        list: Liste des chemins de fichiers
    """
    try:
        from huggingface_hub import HfApi

        hf_token, dataset_repo = get_hf_config()

        if not hf_token or not dataset_repo:
            return []

        api = HfApi(token=hf_token)

        try:
            files = api.list_repo_files(repo_id=dataset_repo, repo_type="dataset")
        except Exception:
            return []

        # Filtrer les fichiers d'annotation
        annotation_files = [f for f in files if f.startswith("annotations/")]

        if annotator_id:
            annotation_files = [
                f for f in annotation_files
                if f.startswith(f"annotations/annotation_{annotator_id}_")
            ]

        return sorted(annotation_files)

    except ImportError:
        return []
    except Exception:
        return []