# ============================================
# Imports
# ============================================

import os
import re
import hashlib
from pathlib import Path
from typing import List, Dict, Optional
from dataclasses import dataclass
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

print("✅ Imports de base chargés")

# ============================================
# Modèle de document
# ============================================

@dataclass
class Document:
    """Représente un document ingéré."""
    id: str                    # Hash unique
    source: str                # Chemin ou URL
    content: str               # Texte extrait
    metadata: Dict             # Infos supplémentaires
    
    @classmethod
    def from_file(cls, path: str, content: str) -> 'Document':
        """Crée un document depuis un fichier."""
        doc_id = hashlib.md5(content.encode()).hexdigest()[:12]
        return cls(
            id=doc_id,
            source=path,
            content=content,
            metadata={
                'filename': Path(path).name,
                'extension': Path(path).suffix,
                'size_chars': len(content),
                'ingested_at': datetime.now().isoformat()
            }
        )

print("✅ Modèle Document défini")

# ============================================
# Parsers pour différents formats
# ============================================

def parse_text(path: str) -> str:
    """Parse un fichier texte (.txt, .md)."""
    with open(path, 'r', encoding='utf-8') as f:
        return f.read()

def parse_pdf(path: str) -> str:
    """Parse un fichier PDF avec PyMuPDF."""
    try:
        import fitz  # PyMuPDF
        doc = fitz.open(path)
        text = ""
        for page in doc:
            text += page.get_text()
        doc.close()
        return text
    except ImportError:
        print("⚠️ PyMuPDF non installé: pip install pymupdf")
        return ""

def parse_html(path: str) -> str:
    """Parse un fichier HTML avec BeautifulSoup."""
    try:
        from bs4 import BeautifulSoup
        with open(path, 'r', encoding='utf-8') as f:
            soup = BeautifulSoup(f.read(), 'html.parser')
        # Supprimer scripts et styles
        for tag in soup(['script', 'style', 'nav', 'footer']):
            tag.decompose()
        return soup.get_text(separator=' ', strip=True)
    except ImportError:
        print("⚠️ BeautifulSoup non installé: pip install beautifulsoup4")
        return ""

def parse_document(path: str) -> Optional[Document]:
    """Parse un document selon son extension."""
    ext = Path(path).suffix.lower()
    
    parsers = {
        '.txt': parse_text,
        '.md': parse_text,
        '.pdf': parse_pdf,
        '.html': parse_html,
        '.htm': parse_html,
    }
    
    if ext not in parsers:
        print(f"⚠️ Format non supporté: {ext}")
        return None
    
    try:
        content = parsers[ext](path)
        if not content.strip():
            print(f"⚠️ Document vide: {path}")
            return None
        return Document.from_file(path, content)
    except Exception as e:
        print(f"❌ Erreur parsing {path}: {e}")
        return None

print("✅ Parsers définis")

# ============================================
# Ingestion d'un dossier
# ============================================

def ingest_folder(folder_path: str) -> List[Document]:
    """
    Ingère tous les documents d'un dossier.
    
    Args:
        folder_path: Chemin du dossier
    
    Returns:
        Liste de Documents
    """
    documents = []
    supported_extensions = {'.txt', '.md', '.pdf', '.html', '.htm'}
    
    folder = Path(folder_path)
    if not folder.exists():
        print(f"❌ Dossier non trouvé: {folder_path}")
        return documents
    
    print(f"📂 Ingestion de {folder_path}...")
    
    for file_path in folder.rglob('*'):
        if file_path.suffix.lower() in supported_extensions:
            doc = parse_document(str(file_path))
            if doc:
                documents.append(doc)
                print(f"   ✅ {file_path.name} ({doc.metadata['size_chars']} chars)")
    
    print(f"\n📊 {len(documents)} documents ingérés")
    return documents

print("✅ Fonction d'ingestion définie")

# ============================================
# Créer des documents de test
# ============================================

os.makedirs('docs', exist_ok=True)

# Document 1: Guide de démarrage
doc1 = '''# Guide de démarrage - ACME Corp

## Introduction

Bienvenue chez ACME Corp ! Ce guide vous aidera à démarrer.

## Configuration de votre poste

1. Installez les outils requis : VS Code, Git, Docker
2. Clonez le repository principal : `git clone https://github.com/acme/main`
3. Configurez vos variables d'environnement

## Procédures importantes

### Demande de congés
Pour demander des congés, utilisez le formulaire HR-01 disponible sur l'intranet.
Délai minimum : 2 semaines avant la date souhaitée.

### Remboursement de frais
Remplissez le formulaire RMB-01 avec les justificatifs.
Délai de traitement : 30 jours ouvrés.
Montant maximum sans validation manager : 100€.
'''

# Document 2: FAQ technique
doc2 = '''# FAQ Technique

## Comment accéder au VPN ?

1. Téléchargez le client VPN depuis l'intranet
2. Utilisez vos identifiants Windows
3. Connectez-vous au serveur vpn.acme.internal

## Mot de passe oublié ?

Contactez le support IT : support@acme.corp ou extension 1234.
Le support est disponible de 9h à 18h.

## Comment déployer en production ?

1. Créez une Pull Request vers la branche main
2. Obtenez 2 approbations minimum
3. Le déploiement est automatique via CI/CD
4. Vérifiez les logs dans Datadog
'''

# Document 3: Architecture
doc3 = '''# Architecture Technique

## Vue d'ensemble

Notre stack technique comprend :
- Frontend : React, TypeScript
- Backend : Python, FastAPI
- Base de données : PostgreSQL, Redis
- Infrastructure : Kubernetes sur AWS

## Microservices

### Service Utilisateurs
Gère l'authentification et les profils.
Port : 8001
Repository : github.com/acme/user-service

### Service Commandes
Gère les commandes et le panier.
Port : 8002
Repository : github.com/acme/order-service

### Service Notifications
Envoie les emails et notifications push.
Port : 8003
Repository : github.com/acme/notification-service
'''

# Sauvegarder
with open('docs/guide_demarrage.md', 'w') as f:
    f.write(doc1)
with open('docs/faq_technique.md', 'w') as f:
    f.write(doc2)
with open('docs/architecture.md', 'w') as f:
    f.write(doc3)

print("✅ 3 documents de test créés dans docs/")

# ============================================
# Tester l'ingestion
# ============================================

documents = ingest_folder('docs')

# Afficher un aperçu
if documents:
    print("\n📄 Aperçu du premier document:")
    print(f"   ID: {documents[0].id}")
    print(f"   Source: {documents[0].source}")
    print(f"   Contenu: {documents[0].content[:200]}...")

# ============================================
# Modèle de Chunk
# ============================================

@dataclass
class Chunk:
    """Représente un morceau de document."""
    id: str              # Hash unique
    doc_id: str          # ID du document parent
    content: str         # Texte du chunk
    metadata: Dict       # Source, position, etc.
    
    @classmethod
    def from_text(cls, text: str, doc: Document, index: int) -> 'Chunk':
        """Crée un chunk depuis du texte."""
        chunk_id = hashlib.md5(f"{doc.id}-{index}-{text[:50]}".encode()).hexdigest()[:12]
        return cls(
            id=chunk_id,
            doc_id=doc.id,
            content=text,
            metadata={
                'source': doc.source,
                'filename': doc.metadata['filename'],
                'chunk_index': index,
                'size_chars': len(text)
            }
        )

print("✅ Modèle Chunk défini")

# ============================================
# Chunker récursif
# ============================================

class RecursiveChunker:
    """
    Découpe les documents de manière récursive.
    
    Stratégie:
    1. Essaye de découper par sections (##)
    2. Si trop grand, découpe par paragraphes
    3. Si encore trop grand, découpe par phrases
    4. En dernier recours, découpe par caractères
    """
    
    def __init__(
        self,
        chunk_size: int = 500,      # Taille cible en caractères
        chunk_overlap: int = 50     # Chevauchement entre chunks
    ):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        
        # Séparateurs par ordre de priorité
        self.separators = [
            "\n## ",    # Titres niveau 2
            "\n### ",   # Titres niveau 3
            "\n\n",     # Paragraphes
            "\n",       # Lignes
            ". ",       # Phrases
            " ",        # Mots
        ]
    
    def _split_text(self, text: str, separator: str) -> List[str]:
        """Découpe le texte par un séparateur."""
        if separator:
            parts = text.split(separator)
            # Réajouter le séparateur au début de chaque partie (sauf la première)
            return [parts[0]] + [separator + p for p in parts[1:]]
        return [text]
    
    def _recursive_split(self, text: str, separators: List[str]) -> List[str]:
        """Découpe récursivement avec les séparateurs."""
        if not separators:
            # Plus de séparateurs, retourner tel quel
            return [text]
        
        sep = separators[0]
        remaining_seps = separators[1:]
        
        if sep not in text:
            # Ce séparateur n'existe pas, essayer le suivant
            return self._recursive_split(text, remaining_seps)
        
        parts = self._split_text(text, sep)
        
        result = []
        for part in parts:
            if len(part) <= self.chunk_size:
                result.append(part)
            else:
                # Trop grand, découper plus finement
                result.extend(self._recursive_split(part, remaining_seps))
        
        return result
    
    def _merge_small_chunks(self, chunks: List[str]) -> List[str]:
        """Fusionne les petits chunks avec overlap."""
        merged = []
        current = ""
        
        for chunk in chunks:
            chunk = chunk.strip()
            if not chunk:
                continue
                
            if len(current) + len(chunk) <= self.chunk_size:
                current += (" " if current else "") + chunk
            else:
                if current:
                    merged.append(current)
                # Ajouter overlap du chunk précédent
                if merged and self.chunk_overlap > 0:
                    overlap = merged[-1][-self.chunk_overlap:]
                    current = overlap + " " + chunk
                else:
                    current = chunk
        
        if current:
            merged.append(current)
        
        return merged
    
    def chunk_document(self, doc: Document) -> List[Chunk]:
        """Découpe un document en chunks."""
        # 1. Découpage récursif
        raw_chunks = self._recursive_split(doc.content, self.separators)
        
        # 2. Fusion des petits chunks
        merged = self._merge_small_chunks(raw_chunks)
        
        # 3. Créer les objets Chunk
        chunks = []
        for i, text in enumerate(merged):
            if text.strip():
                chunks.append(Chunk.from_text(text.strip(), doc, i))
        
        return chunks

print("✅ RecursiveChunker défini")

# ============================================
# Comparaison : Fixed-size vs Recursive
# ============================================

def chunk_fixed_size(text: str, size: int = 500) -> List[str]:
    """Découpage naïf par taille fixe."""
    return [text[i:i+size] for i in range(0, len(text), size)]

# Test sur le même document
test_doc = documents[0]
print(f"📄 Document: {test_doc.metadata['filename']}")
print(f"   Taille: {len(test_doc.content)} caractères\n")

# Fixed-size
fixed_chunks = chunk_fixed_size(test_doc.content, size=500)
print(f"🔲 FIXED-SIZE (500 chars):")
print(f"   {len(fixed_chunks)} chunks")
print(f"   Exemple chunk 1: '{fixed_chunks[0][:80]}...'")
print(f"   ⚠️ Coupe au milieu des mots/phrases\n")

# Recursive
recursive_chunks = chunker.chunk_document(test_doc)
print(f"🔄 RECURSIVE (500 chars, respect structure):")
print(f"   {len(recursive_chunks)} chunks")
print(f"   Exemple chunk 1: '{recursive_chunks[0].content[:80]}...'")
print(f"   ✅ Respecte les sections et paragraphes")

# ============================================
# Tester le chunking
# ============================================

chunker = RecursiveChunker(chunk_size=500, chunk_overlap=50)

all_chunks = []
for doc in documents:
    chunks = chunker.chunk_document(doc)
    all_chunks.extend(chunks)
    print(f"📄 {doc.metadata['filename']}: {len(chunks)} chunks")

print(f"\n📊 Total: {len(all_chunks)} chunks")

# Aperçu
if all_chunks:
    print(f"\n📝 Exemple de chunk:")
    print(f"   ID: {all_chunks[0].id}")
    print(f"   Taille: {all_chunks[0].metadata['size_chars']} chars")
    print(f"   Contenu: {all_chunks[0].content[:200]}...")

# ============================================
# Embedder avec Sentence-Transformers
# ============================================

from sentence_transformers import SentenceTransformer
import numpy as np

class TextEmbedder:
    """
    Génère des embeddings pour du texte.
    
    Modèles recommandés:
    - paraphrase-multilingual-MiniLM-L12-v2 : Multilingue (FR) ✅
    - dangvantuan/sentence-camembert-large : Français spécifique
    - all-MiniLM-L6-v2 : Anglais seulement
    """
    
    def __init__(self, model_name: str = "paraphrase-multilingual-MiniLM-L12-v2"):
        print(f"📥 Chargement du modèle {model_name}...")
        self.model = SentenceTransformer(model_name)
        self.model_name = model_name
        self.dimension = self.model.get_sentence_embedding_dimension()
        print(f"✅ Modèle chargé (dimension: {self.dimension})")
    
    def embed(self, text: str) -> np.ndarray:
        """Génère l'embedding d'un texte."""
        return self.model.encode(text, normalize_embeddings=True)
    
    def embed_batch(self, texts: List[str], batch_size: int = 32) -> np.ndarray:
        """Génère les embeddings pour une liste de textes."""
        return self.model.encode(
            texts,
            batch_size=batch_size,
            normalize_embeddings=True,
            show_progress_bar=True
        )

# Initialiser avec le modèle MULTILINGUE (supporte le français)
embedder = TextEmbedder("paraphrase-multilingual-MiniLM-L12-v2")

# ============================================
# Tester les embeddings
# ============================================

# Test simple
test_embedding = embedder.embed("Comment configurer le VPN ?")
print(f"📊 Embedding de test:")
print(f"   Shape: {test_embedding.shape}")
print(f"   Norme: {np.linalg.norm(test_embedding):.4f}")

# Test de similarité
queries = [
    "Comment se connecter au VPN ?",
    "Procédure de remboursement",
    "Architecture des microservices"
]

print(f"\n🔍 Test de similarité:")
q1_emb = embedder.embed(queries[0])
q2_emb = embedder.embed(queries[1])
q3_emb = embedder.embed(queries[2])

sim_12 = np.dot(q1_emb, q2_emb)
sim_13 = np.dot(q1_emb, q3_emb)

print(f"   '{queries[0]}' vs '{queries[1]}': {sim_12:.3f}")
print(f"   '{queries[0]}' vs '{queries[2]}': {sim_13:.3f}")

# ============================================
# Configurer Chroma pour RAG
# ============================================

import chromadb

# Client persistant
chroma_client = chromadb.PersistentClient(path="./chroma_rag_db")

# Collection pour le knowledge base
COLLECTION_NAME = "knowledge_base"

# Supprimer si existe (pour repartir de zéro)
try:
    chroma_client.delete_collection(COLLECTION_NAME)
except:
    pass

# Créer la collection
collection = chroma_client.create_collection(
    name=COLLECTION_NAME,
    metadata={"description": "Knowledge Base pour RAG - Documentation ACME"}
)

print(f"✅ Collection Chroma créée: {COLLECTION_NAME}")

# ============================================
# Indexer les chunks
# ============================================

def index_chunks(chunks: List[Chunk], embedder: TextEmbedder, collection) -> int:
    """
    Indexe les chunks dans Chroma.
    
    Args:
        chunks: Liste de Chunks
        embedder: Modèle d'embedding
        collection: Collection Chroma
    
    Returns:
        Nombre de chunks indexés
    """
    if not chunks:
        return 0
    
    print(f"📤 Indexation de {len(chunks)} chunks...")
    
    # Extraire les contenus
    texts = [c.content for c in chunks]
    
    # Générer les embeddings
    embeddings = embedder.embed_batch(texts)
    
    # Préparer les données pour Chroma
    ids = [c.id for c in chunks]
    metadatas = [c.metadata for c in chunks]
    
    # Indexer
    collection.add(
        ids=ids,
        embeddings=embeddings.tolist(),
        documents=texts,
        metadatas=metadatas
    )
    
    print(f"✅ {len(chunks)} chunks indexés")
    return len(chunks)

# Indexer nos chunks
n_indexed = index_chunks(all_chunks, embedder, collection)

print(f"\n📊 Collection Chroma: {collection.count()} documents")

# ============================================
# Fonction de recherche
# ============================================

def search(query: str, n_results: int = 3) -> List[Dict]:
    """
    Recherche les chunks les plus pertinents.
    
    Args:
        query: Question de l'utilisateur
        n_results: Nombre de résultats
    
    Returns:
        Liste de résultats avec score et contenu
    """
    # Embed la query
    query_embedding = embedder.embed(query)
    
    # Rechercher dans Chroma
    results = collection.query(
        query_embeddings=[query_embedding.tolist()],
        n_results=n_results,
        include=['documents', 'metadatas', 'distances']
    )
    
    # Formatter les résultats
    formatted = []
    for i in range(len(results['ids'][0])):
        formatted.append({
            'id': results['ids'][0][i],
            'content': results['documents'][0][i],
            'source': results['metadatas'][0][i].get('filename', 'Unknown'),
            'score': 1 - results['distances'][0][i]  # Convertir distance en score
        })
    
    return formatted

print("✅ Fonction de recherche définie")

# ============================================
# Tester la recherche
# ============================================

test_queries = [
    "Comment se connecter au VPN ?",
    "Quelle est la procédure de remboursement ?",
    "Quels sont les microservices ?"
]

for query in test_queries:
    print(f"\n🔍 Query: '{query}'")
    print("─" * 50)
    
    results = search(query, n_results=2)
    
    for i, r in enumerate(results, 1):
        print(f"\n   📄 Résultat {i} (score: {r['score']:.3f})")
        print(f"   Source: {r['source']}")
        print(f"   Contenu: {r['content'][:200]}...")

# ============================================
# Évaluation du retrieval
# ============================================

def evaluate_retrieval(
    test_queries: List[Dict],  # [{"query": "...", "expected_source": "..."}]
    k: int = 3
) -> Dict:
    """
    Évalue la qualité du retrieval.
    
    Args:
        test_queries: Liste de queries avec la source attendue
        k: Nombre de résultats à considérer
    
    Returns:
        Métriques d'évaluation
    """
    hits = 0
    reciprocal_ranks = []
    
    print(f"📊 Évaluation sur {len(test_queries)} queries (k={k})\n")
    
    for item in test_queries:
        query = item['query']
        expected = item['expected_source']
        
        results = search(query, n_results=k)
        sources = [r['source'] for r in results]
        
        # Hit Rate
        if expected in sources:
            hits += 1
            rank = sources.index(expected) + 1
            reciprocal_ranks.append(1.0 / rank)
            status = f"✅ Trouvé (rang {rank})"
        else:
            reciprocal_ranks.append(0.0)
            status = "❌ Non trouvé"
        
        print(f"   Query: '{query[:40]}...'")
        print(f"   Attendu: {expected} → {status}")
    
    # Calculer les métriques
    hit_rate = hits / len(test_queries)
    mrr = sum(reciprocal_ranks) / len(reciprocal_ranks)
    
    print(f"\n" + "="*50)
    print(f"📈 RÉSULTATS")
    print(f"="*50)
    print(f"   Hit Rate @{k}: {hit_rate:.1%}")
    print(f"   MRR: {mrr:.3f}")
    
    return {'hit_rate': hit_rate, 'mrr': mrr}

# Créer un jeu de test
test_queries = [
    {"query": "Comment demander des congés ?", "expected_source": "guide_demarrage.md"},
    {"query": "Procédure de remboursement de frais", "expected_source": "guide_demarrage.md"},
    {"query": "Comment se connecter au VPN ?", "expected_source": "faq_technique.md"},
    {"query": "Mot de passe oublié que faire", "expected_source": "faq_technique.md"},
    {"query": "Architecture des microservices", "expected_source": "architecture.md"},
    {"query": "Quel port pour le service utilisateurs", "expected_source": "architecture.md"},
]

# Évaluer
metrics = evaluate_retrieval(test_queries, k=3)

# ============================================
# Pipeline RAG complet
# ============================================

class RAGPipeline:
    """
    Pipeline complet pour construire un Knowledge Base RAG.
    
    Étapes:
    1. Ingest: Charger les documents
    2. Chunk: Découper en morceaux
    3. Embed: Vectoriser
    4. Index: Stocker dans Chroma
    """
    
    def __init__(
        self,
        collection_name: str = "rag_knowledge_base",
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        embedding_model: str = "all-MiniLM-L6-v2"
    ):
        self.collection_name = collection_name
        self.chunker = RecursiveChunker(chunk_size, chunk_overlap)
        self.embedder = TextEmbedder(embedding_model)
        
        # Créer/reset la collection
        try:
            chroma_client.delete_collection(collection_name)
        except:
            pass
        self.collection = chroma_client.create_collection(name=collection_name)
        
        self.stats = {
            'documents': 0,
            'chunks': 0,
            'indexed': 0
        }
    
    def run(self, source_folder: str) -> Dict:
        """
        Exécute le pipeline complet.
        """
        print("="*60)
        print("🚀 PIPELINE RAG - KNOWLEDGE BASE")
        print("="*60)
        
        # 1. Ingest
        print("\n📥 Étape 1: Ingestion...")
        documents = ingest_folder(source_folder)
        self.stats['documents'] = len(documents)
        
        # 2. Chunk
        print("\n✂️ Étape 2: Chunking...")
        all_chunks = []
        for doc in documents:
            chunks = self.chunker.chunk_document(doc)
            all_chunks.extend(chunks)
            print(f"   {doc.metadata['filename']}: {len(chunks)} chunks")
        self.stats['chunks'] = len(all_chunks)
        
        # 3. Embed & Index
        print("\n🧠 Étape 3: Embedding & Indexation...")
        n_indexed = index_chunks(all_chunks, self.embedder, self.collection)
        self.stats['indexed'] = n_indexed
        
        # Résumé
        print("\n" + "="*60)
        print("📊 RÉSUMÉ")
        print("="*60)
        print(f"   Documents traités: {self.stats['documents']}")
        print(f"   Chunks créés: {self.stats['chunks']}")
        print(f"   Chunks indexés: {self.stats['indexed']}")
        print(f"   Collection: {self.collection_name}")
        
        return self.stats
    
    def search(self, query: str, n_results: int = 3) -> List[Dict]:
        """Recherche dans le knowledge base."""
        query_embedding = self.embedder.embed(query)
        results = self.collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=n_results,
            include=['documents', 'metadatas', 'distances']
        )
        
        formatted = []
        for i in range(len(results['ids'][0])):
            formatted.append({
                'content': results['documents'][0][i],
                'source': results['metadatas'][0][i].get('filename'),
                'score': 1 - results['distances'][0][i]
            })
        return formatted

print("✅ RAGPipeline défini")

# ============================================
# Exécuter le pipeline
# ============================================

# Créer et exécuter
rag_pipeline = RAGPipeline(
    collection_name="acme_knowledge_base",
    chunk_size=500,
    chunk_overlap=50
)

stats = rag_pipeline.run("docs")

# Test de recherche
print("\n" + "="*60)
print("🔍 TEST DE RECHERCHE")
print("="*60)

query = "Comment demander un remboursement ?"
print(f"\nQuery: '{query}'")
results = rag_pipeline.search(query, n_results=2)

for i, r in enumerate(results, 1):
    print(f"\n📄 Résultat {i} (score: {r['score']:.3f}) - {r['source']}")
    print(f"   {r['content'][:300]}...")

# ============================================
# Setup Ollama
# ============================================

# Installation (exécuter dans le terminal):
# curl -fsSL https://ollama.com/install.sh | sh
# ollama pull mistral

# Vérifier qu'Ollama tourne
import requests

OLLAMA_URL = "http://localhost:11434"

try:
    response = requests.get(f"{OLLAMA_URL}/api/tags")
    models = [m['name'] for m in response.json().get('models', [])]
    print(f"✅ Ollama connecté")
    print(f"   Modèles disponibles: {models}")
except:
    print("⚠️ Ollama non détecté")
    print("   Installation: curl -fsSL https://ollama.com/install.sh | sh")
    print("   Puis: ollama pull mistral")

# ============================================
# RAG complet avec Ollama
# ============================================

def call_ollama(prompt: str, model: str = "mistral") -> str:
    """
    Appelle Ollama en local.
    
    Modèles recommandés:
    - mistral : Bon en français, léger (4GB)
    - llama3 : Plus puissant, plus lourd (8GB)
    - phi3 : Très léger (2GB), moins bon
    """
    try:
        response = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={
                "model": model,
                "prompt": prompt,
                "stream": False
            },
            timeout=60
        )
        return response.json()['response']
    except Exception as e:
        return f"Erreur Ollama: {e}"

def rag_query(question: str, n_context: int = 3, model: str = "mistral") -> str:
    """
    Exécute une requête RAG complète.
    
    1. Retrieve : Cherche les chunks pertinents
    2. Augment : Construit le prompt avec le contexte
    3. Generate : Appelle le LLM local (Ollama)
    """
    
    # 1. RETRIEVE
    print(f"🔍 Question: {question}")
    print(f"\n📚 Étape 1: RETRIEVE")
    
    results = search(question, n_results=n_context)
    
    context_parts = []
    for i, r in enumerate(results, 1):
        print(f"   [{i}] {r['source']} (score: {r['score']:.2f})")
        context_parts.append(f"[Source: {r['source']}]\n{r['content']}")
    
    context = "\n\n---\n\n".join(context_parts)
    
    # 2. AUGMENT
    print(f"\n📝 Étape 2: AUGMENT")
    
    prompt = f"""Tu es un assistant qui répond aux questions en te basant UNIQUEMENT sur le contexte fourni.
Si la réponse n'est pas dans le contexte, dis-le clairement.
Réponds en français.

CONTEXTE:
{context}

QUESTION: {question}

RÉPONSE:"""
    
    print(f"   Prompt: {len(prompt)} caractères")
    
    # 3. GENERATE
    print(f"\n🤖 Étape 3: GENERATE (Ollama/{model})")
    
    answer = call_ollama(prompt, model=model)
    
    return answer

print("✅ Fonction RAG définie")

# ============================================
# Test RAG complet
# ============================================

print("="*60)
print("🚀 TEST RAG COMPLET AVEC OLLAMA")
print("="*60 + "\n")

question = "Comment demander un remboursement de frais ?"

try:
    answer = rag_query(question, n_context=2, model="mistral")
    
    print("\n" + "="*60)
    print("💬 RÉPONSE DU LLM:")
    print("="*60)
    print(answer)
except Exception as e:
    print(f"\n⚠️ Erreur: {e}")
    print("\n💡 Assure-toi qu'Ollama est lancé:")
    print("   ollama serve")
    print("   ollama pull mistral")

# ============================================
# Alternative : API Cloud (si tu as une clé)
# ============================================

# Décommenter pour utiliser

'''
# ===== OPTION A: GROQ (Gratuit) =====
# pip install groq
# Clé API gratuite sur : https://console.groq.com
from groq import Groq

groq_client = Groq(api_key="ta-clé-groq")

def call_groq(prompt: str, model: str = "llama3-8b-8192") -> str:
    """
    Appelle Groq (gratuit, ultra-rapide).
    
    Modèles disponibles:
    - llama3-8b-8192  : Léger et rapide ✅
    - llama3-70b-8192 : Plus puissant
    - mixtral-8x7b    : Bon en français
    """
    response = groq_client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=500
    )
    return response.choices[0].message.content


# ===== OPTION A-bis: Mistral API (Gratuit, bon en français) =====
# pip install mistralai
# Clé API gratuite sur : https://console.mistral.ai
from mistralai import Mistral

mistral_client = Mistral(api_key="ta-clé-mistral")

def call_mistral(prompt: str, model: str = "mistral-small-latest") -> str:
    """
    Appelle Mistral API (gratuit, excellent en français).
    
    Modèles gratuits:
    - mistral-small-latest : Rapide et efficace ✅
    - open-mistral-nemo     : Très léger, gratuit
    """
    response = mistral_client.chat.complete(
        model=model,
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content


# ===== OPTION B: Anthropic (Claude) =====


import anthropic

client = anthropic.Anthropic(api_key="ta-clé-api")

def call_claude(prompt: str) -> str:
    response = client.messages.create(
        model="claude-3-haiku-20240307",
        max_tokens=500,
        messages=[{"role": "user", "content": prompt}]
    )
    return response.content[0].text


# ===== OPTION C: OpenAI =====
from openai import OpenAI

client = OpenAI(api_key="ta-clé-api")

def call_openai(prompt: str) -> str:
    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content


'''

print("💡 Alternatives cloud disponibles (décommenter si tu as une clé API):")
print("   - Groq (gratuit, rapide) : console.groq.com")
print("   - Mistral API (gratuit, bon en FR) : console.mistral.ai
    print("   - Anthropic Claude (payant)")
print("   - OpenAI GPT (payant)")

# ============================================
# Valider le dataset de fine-tuning
# ============================================

def validate_finetuning_dataset(path: str) -> Dict:
    """
    Valide un dataset JSONL pour fine-tuning.
    
    Checks:
    - Format JSONL valide
    - Structure messages correcte
    - Rôles valides (system, user, assistant)
    - Pas de contenu vide
    """
    stats = {
        'total': 0,
        'valid': 0,
        'errors': []
    }
    
    with open(path, 'r') as f:
        for i, line in enumerate(f, 1):
            stats['total'] += 1
            
            try:
                data = json.loads(line)
                
                # Check structure
                if 'messages' not in data:
                    stats['errors'].append(f"Ligne {i}: 'messages' manquant")
                    continue
                
                # Check messages
                for msg in data['messages']:
                    if 'role' not in msg or 'content' not in msg:
                        stats['errors'].append(f"Ligne {i}: message mal formé")
                        continue
                    if msg['role'] not in ['system', 'user', 'assistant']:
                        stats['errors'].append(f"Ligne {i}: rôle invalide '{msg['role']}'")
                        continue
                    if not msg['content'].strip():
                        stats['errors'].append(f"Ligne {i}: contenu vide")
                        continue
                
                stats['valid'] += 1
                
            except json.JSONDecodeError:
                stats['errors'].append(f"Ligne {i}: JSON invalide")
    
    # Résumé
    print(f"📊 Validation de {path}")
    print(f"   Total: {stats['total']}")
    print(f"   Valides: {stats['valid']}")
    print(f"   Erreurs: {len(stats['errors'])}")
    
    if stats['errors']:
        print("\n⚠️ Erreurs:")
        for err in stats['errors'][:5]:
            print(f"   {err}")
    else:
        print("\n✅ Dataset valide !")
    
    return stats

# Valider
validate_finetuning_dataset("finetuning_dataset.jsonl")

# ============================================
# Générer le package téléchargeable
# ============================================

import zipfile
import base64
from IPython.display import HTML, display

# Créer la structure
os.makedirs('module6-rag-project/src', exist_ok=True)
os.makedirs('module6-rag-project/docs', exist_ok=True)

# requirements.txt
requirements = '''sentence-transformers>=2.2.0
chromadb>=0.4.0
pymupdf>=1.23.0
beautifulsoup4>=4.12.0
boto3>=1.28.0
'''

with open('module6-rag-project/requirements.txt', 'w') as f:
    f.write(requirements)

# Copier les docs de test
for doc_file in Path('docs').glob('*.md'):
    with open(doc_file, 'r') as src:
        with open(f'module6-rag-project/docs/{doc_file.name}', 'w') as dst:
            dst.write(src.read())

# Créer le ZIP
zip_path = 'module6-rag-project.zip'
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
    for root, dirs, files in os.walk('module6-rag-project'):
        for file in files:
            file_path = os.path.join(root, file)
            zipf.write(file_path, file_path)

# Lien de téléchargement
with open(zip_path, 'rb') as f:
    b64 = base64.b64encode(f.read()).decode()

html = f'''
<a download="module6-rag-project.zip" 
   href="data:application/zip;base64,{b64}" 
   style="display: inline-block; padding: 10px 20px; 
          background-color: #4CAF50; color: white; 
          text-decoration: none; border-radius: 5px;
          font-weight: bold;">
   📦 Télécharger module6-rag-project.zip
</a>
'''

display(HTML(html))
print(f"\n✅ Archive créée: {zip_path}")