data_loader.py

"""
Module de chargement des données depuis Hugging Face
"""
import os
import pandas as pd
from datasets import load_dataset
from huggingface_hub import HfApi, hf_hub_download
from config import HF_TOKEN, DATASET_ID, REQUIRED_COLUMNS, MESSAGES


class DataLoader:
    """Classe responsable du chargement des données depuis différentes sources"""
    
    def __init__(self):
        self.df = None
    
    def load_data(self):
        """Charge les données du dataset Hugging Face avec gestion d'erreur robuste"""
        try:
            print(MESSAGES["loading"])
            print(f"📋 Dataset ID: {DATASET_ID}")
            print(f"📋 Token disponible: {'Oui' if HF_TOKEN else 'Non'}")

            self.df = None

            # 1) Tentative de chargement direct via datasets.load_dataset
            try:
                print("🔄 Tentative de chargement direct...")
                dataset = load_dataset(
                    DATASET_ID,
                    split="train",
                    token=HF_TOKEN,
                    trust_remote_code=True,
                )
                print(f"📊 Dataset chargé: {len(dataset)} exemples")

                try:
                    self.df = dataset.to_pandas()
                    print("✅ Conversion to_pandas() réussie")
                except Exception as pandas_error:
                    print(f"❌ Erreur to_pandas(): {pandas_error}")
                    print("🔄 Tentative de conversion manuelle...")
                    
                    try:
                        data_list = []
                        max_examples = min(len(dataset), 1000)  # Limiter pour éviter les problèmes de mémoire
                        
                        for i, item in enumerate(dataset):
                            if i >= max_examples:
                                break
                            data_list.append(item)
                            if i < 5:
                                print(f"📋 Exemple {i}: {list(item.keys())}")
                        
                        self.df = pd.DataFrame(data_list)
                        print(f"✅ Conversion manuelle réussie: {len(self.df)} lignes")
                    except Exception as manual_error:
                        print(f"❌ Erreur lors de la conversion manuelle: {manual_error}")
                        self.df = None
                        
            except Exception as e:
                print(f"❌ Erreur lors du chargement depuis Hugging Face: {str(e)}")
                print(f"❌ Type d'erreur: {type(e).__name__}")
                
                # 2) Fallback: récupérer directement les fichiers du repo
                try:
                    fallback_msg = self._fallback_load_from_repo_files()
                    if self.df is None and fallback_msg:
                        return f"❌ Erreur lors du chargement du dataset : {str(e)} | Fallback: {fallback_msg}"
                except Exception as fallback_error:
                    print(f"❌ Erreur dans le fallback: {fallback_error}")
                    # Continue vers le chargement local

            # Si on n'a toujours pas de dataframe, arrêter
            if self.df is None:
                print("⚠️ Aucune méthode de chargement n'a fonctionné")
                return MESSAGES["no_data"]

            print(f"📊 Données chargées: {len(self.df)} lignes")
            print(f"📊 Colonnes disponibles: {list(self.df.columns)}")

            # Nettoyage et validation
            return self._clean_and_validate_data()
            
        except Exception as e:
            print(f"❌ Erreur critique dans load_data: {e}")
            import traceback
            traceback.print_exc()
            return f"❌ Erreur critique lors du chargement: {str(e)}"
    
    def _clean_and_validate_data(self):
        """Nettoie et valide les données chargées"""
        missing_cols = [col for col in REQUIRED_COLUMNS if col not in self.df.columns]

        if missing_cols:
            print(f"❌ Colonnes manquantes: {missing_cols}")
            self.df = None
            return f"❌ Colonnes manquantes: {missing_cols}"

        # Nettoyage
        initial_len = len(self.df)
        self.df = self.df.dropna(subset=REQUIRED_COLUMNS)

        print(f"📊 Avant nettoyage: {initial_len} lignes")
        print(f"📊 Après nettoyage: {len(self.df)} lignes")
        
        return MESSAGES["success"]
    
    def _fallback_load_from_repo_files(self):
        """Fallback pour charger les données en téléchargeant directement les fichiers du repo HF."""
        try:
            print("🔄 Tentative de chargement alternatif via fichiers du dépôt Hugging Face...")
            api = HfApi()
            files = api.list_repo_files(repo_id=DATASET_ID, repo_type="dataset", token=HF_TOKEN)
            if not files:
                print("❌ Aucun fichier dans le dépôt")
                return "Aucun fichier trouvé dans le dépôt."

            data_files = [
                f for f in files if f.lower().endswith((".parquet", ".csv", ".tsv", ".json"))
            ]
            if not data_files:
                print("❌ Aucun fichier de données exploitable (csv/tsv/parquet/json)")
                return "Aucun fichier exploitable (csv/tsv/parquet/json)."

            # Priorité: parquet > csv > tsv > json
            for ext in [".parquet", ".csv", ".tsv", ".json"]:
                selected = [f for f in data_files if f.lower().endswith(ext)]
                if selected:
                    chosen_ext = ext
                    selected_files = selected
                    break

            print(f"📂 Fichiers détectés ({chosen_ext}): {selected_files[:5]}{' ...' if len(selected_files) > 5 else ''}")

            local_paths = []
            for f in selected_files:
                local_path = hf_hub_download(
                    repo_id=DATASET_ID,
                    repo_type="dataset",
                    filename=f,
                    token=HF_TOKEN,
                )
                local_paths.append(local_path)

            frames = []
            if chosen_ext == ".parquet":
                for p in local_paths:
                    frames.append(pd.read_parquet(p))
            elif chosen_ext == ".csv":
                for p in local_paths:
                    frames.append(pd.read_csv(p))
            elif chosen_ext == ".tsv":
                for p in local_paths:
                    frames.append(pd.read_csv(p, sep="\t"))
            elif chosen_ext == ".json":
                for p in local_paths:
                    try:
                        frames.append(pd.read_json(p, lines=True))
                    except Exception:
                        frames.append(pd.read_json(p))

            self.df = pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]
            print(f"✅ Fallback réussi: {len(self.df)} lignes chargées depuis les fichiers du dépôt")
            return None
        except Exception as e:
            print(f"❌ Fallback échoué: {e}")
            # Dernier recours: fichier local d'exemple
            return self._load_local_sample()
    
    def _load_local_sample(self):
        """Charge un fichier local de secours"""
        sample_path = os.path.join(os.path.dirname(__file__), "sample_data.csv")
        if os.path.exists(sample_path):
            try:
                self.df = pd.read_csv(sample_path)
                print(f"✅ Chargement du fichier local 'sample_data.csv' ({len(self.df)} lignes)")
                return "Chargement via fichier local de secours."
            except Exception as e2:
                print(f"❌ Échec du chargement du fichier local: {e2}")
        return "Aucune source de données disponible."
    
    def get_data(self):
        """Retourne les données chargées"""
        return self.df
    
    def has_data(self):
        """Vérifie si des données sont disponibles"""
        return self.df is not None and len(self.df) > 0