main.py

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from threading import Thread
import os
from io import BytesIO
import logging
from datetime import datetime
import json
import hashlib

# --- CONFIGURATION LOGGING ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- CONFIGURATION PDF ---
try:
    from pypdf import PdfReader as PypdfReader
    PDF_READER = "pypdf"
except ImportError:
    try:
        import PyPDF2
        from PyPDF2 import PdfReader as PypdfReader
        PDF_READER = "PyPDF2"
    except ImportError:
        raise ImportError("Installe pypdf ou PyPDF2 : pip install pypdf")

# --- OUTILS PERSONNALISÉS ---
from tools.web import web_search
from tools.todo import execute_reflection_plan
from tools.geo import get_geo_context

app = FastAPI(title="Kibali AI API", version="1.0")

# --- SERVEUR STATIQUE ---
script_dir = os.path.dirname(os.path.abspath(__file__))
static_dir = os.path.join(script_dir, "static")
os.makedirs(static_dir, exist_ok=True)
app.mount("/static", StaticFiles(directory=static_dir), name="static")

# --- CORS ---
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- CHARGEMENT DES MODÈLES (téléchargement depuis Hugging Face Hub) ---
HF_MODEL_ID = "BelikanM/kibali-final-merged"
CACHE_DIR = "/data/cache"  # Dossier persistant sur HF Spaces

os.makedirs(CACHE_DIR, exist_ok=True)

logger.info("Chargement du modèle d'embedding...")
embed_model = SentenceTransformer(
    'paraphrase-multilingual-MiniLM-L12-v2',
    cache_folder=CACHE_DIR
)

logger.info(f"Chargement du tokenizer et du modèle LLM depuis Hugging Face : {HF_MODEL_ID}")
tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, cache_dir=CACHE_DIR)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Configuration 4-bit pour réduire la consommation VRAM
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16
)

try:
    model = AutoModelForCausalLM.from_pretrained(
        HF_MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=torch.float16,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        cache_dir=CACHE_DIR
    )
    logger.info(f"Modèle chargé avec succès sur {model.device}")
except Exception as e:
    logger.error(f"Erreur lors du chargement du modèle : {e}")
    raise e

# --- BASES VECTORIELLES GLOBALES ---
dimension = 384
doc_index = faiss.IndexFlatL2(dimension)
doc_chunks: List[str] = []
doc_metadata: List[dict] = []

memory_index = faiss.IndexFlatL2(dimension)
memory_texts: List[str] = []
memory_metadata: List[dict] = []

# --- GESTION DU CONTEXTE CONVERSATIONNEL ---
class ConversationContext:
    def __init__(self):
        self.current_subject = None
        self.subject_embedding = None
        self.subject_start_time = None
        self.message_count = 0
        self.subject_keywords = []
      
    def update_subject(self, message: str, embedding: np.ndarray):
        keywords = self._extract_keywords(message)
      
        if self.subject_embedding is not None:
            similarity = np.dot(embedding.flatten(), self.subject_embedding.flatten())
            if similarity < 0.6:
                logger.info(f"Changement de sujet détecté (similarité: {similarity:.2f})")
                self._archive_current_subject()
                self.current_subject = message
                self.subject_embedding = embedding
                self.subject_start_time = datetime.now()
                self.message_count = 1
                self.subject_keywords = keywords
            else:
                self.message_count += 1
                self.subject_keywords.extend(keywords)
                self.subject_keywords = list(set(self.subject_keywords))[:10]
        else:
            self.current_subject = message
            self.subject_embedding = embedding
            self.subject_start_time = datetime.now()
            self.message_count = 1
            self.subject_keywords = keywords
   
    def _extract_keywords(self, text: str) -> List[str]:
        stopwords = {'le', 'la', 'les', 'un', 'une', 'des', 'de', 'du', 'et', 'ou',
                     'est', 'sont', 'à', 'au', 'en', 'pour', 'dans', 'sur', 'avec'}
        words = text.lower().split()
        keywords = [w for w in words if len(w) > 3 and w not in stopwords]
        return keywords[:5]
   
    def _archive_current_subject(self):
        if self.current_subject and memory_index.ntotal > 0:
            summary = {
                "subject": self.current_subject[:200],
                "keywords": self.subject_keywords,
                "message_count": self.message_count,
                "duration": (datetime.now() - self.subject_start_time).seconds,
                "archived_at": datetime.now().isoformat()
            }
            logger.info(f"Sujet archivé: {summary['keywords']}")

conversation_ctx = ConversationContext()

# --- MODÈLES PYDANTIC ---
class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    latitude: float
    longitude: float
    city: Optional[str] = "Libreville"
    thinking_mode: bool = True

class ChatResponse(BaseModel):
    response: str
    images: List[str] = []
    context_info: Optional[dict] = None

# --- UTILITAIRES ---
def extract_text_from_pdf(pdf_bytes: bytes) -> str:
    text = ""
    try:
        pdf_file = BytesIO(pdf_bytes)
        reader = PypdfReader(pdf_file)
        for page in reader.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
        return text.strip()
    except Exception as e:
        logger.error(f"Erreur extraction PDF : {e}")
        return ""

def chunk_text(text: str, chunk_size: int = 400, overlap: int = 50) -> List[str]:
    if not text.strip():
        return []
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunk_words = words[i:i + chunk_size]
        chunk = " ".join(chunk_words)
        if chunk.strip():
            chunks.append(chunk.strip())
        i += chunk_size - overlap
        if i >= len(words) and len(chunk_words) < overlap:
            break
    return chunks

def add_to_memory_realtime(user_msg: str, ai_response: str, subject_keywords: List[str]):
    timestamp = datetime.now().isoformat()
    memory_entry = f"""[{timestamp}]
Sujet: {', '.join(subject_keywords)}
Utilisateur: {user_msg}
Kibali: {ai_response}"""
   
    metadata = {
        "timestamp": timestamp,
        "subject_keywords": subject_keywords,
        "user_length": len(user_msg),
        "ai_length": len(ai_response),
        "hash": hashlib.md5(memory_entry.encode()).hexdigest()
    }
   
    if metadata["hash"] not in [m.get("hash") for m in memory_metadata]:
        memory_texts.append(memory_entry)
        memory_metadata.append(metadata)
        mem_emb = embed_model.encode([memory_entry], normalize_embeddings=True).astype('float32')
        memory_index.add(mem_emb)
        logger.info(f"Mémoire ajoutée en temps réel: {subject_keywords} (total: {len(memory_texts)})")
        return True
    return False

def retrieve_adaptive_memory(query: str, k: int = 5) -> tuple:
    if memory_index.ntotal == 0:
        return [], []
   
    query_emb = embed_model.encode([query], normalize_embeddings=True).astype('float32')
    k_search = min(k * 2, memory_index.ntotal)
    D, I = memory_index.search(query_emb, k=k_search)
   
    results = []
    for dist, idx in zip(D[0], I[0]):
        if 0 <= idx < len(memory_texts):
            metadata = memory_metadata[idx] if idx < len(memory_metadata) else {}
            recency_score = 1.0 / (1 + (datetime.now() - datetime.fromisoformat(metadata.get("timestamp", datetime.now().isoformat()))).seconds / 3600)
            similarity_score = 1.0 / (1 + dist)
            keyword_bonus = 0
            if conversation_ctx.subject_keywords:
                text_lower = memory_texts[idx].lower()
                keyword_bonus = sum(1 for kw in conversation_ctx.subject_keywords if kw in text_lower) * 0.1
            total_score = similarity_score * 0.6 + recency_score * 0.3 + keyword_bonus
           
            results.append({
                "text": memory_texts[idx],
                "score": total_score,
                "metadata": metadata
            })
   
    results = sorted(results, key=lambda x: x["score"], reverse=True)[:k]
    texts = [r["text"] for r in results]
    scores = [r["score"] for r in results]
    return texts, scores

# --- ROUTES ---
@app.get("/status")
async def status():
    return {
        "status": "ready",
        "doc_chunks": len(doc_chunks),
        "memory_entries": len(memory_texts),
        "pdf_library": PDF_READER,
        "model_device": str(model.device),
        "torch_cuda_available": torch.cuda.is_available(),
        "current_subject": conversation_ctx.current_subject[:100] if conversation_ctx.current_subject else None,
        "subject_message_count": conversation_ctx.message_count
    }

@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    user_message = request.messages[-1].content.strip()
    if not user_message:
        raise HTTPException(status_code=400, detail="Message vide")
   
    geo = {
        "latitude": request.latitude,
        "longitude": request.longitude,
        "city": request.city or "Libreville"
    }
   
    user_emb = embed_model.encode([user_message], normalize_embeddings=True).astype('float32')
    conversation_ctx.update_subject(user_message, user_emb)
   
    # RAG Documents PDF
    rag_context = ""
    rag_sources = []
    if doc_index.ntotal > 0 and len(doc_chunks) > 0:
        D, I = doc_index.search(user_emb, k=5)
        relevant_chunks = []
        for idx in I[0]:
            if 0 <= idx < len(doc_chunks):
                relevant_chunks.append(doc_chunks[idx][:1000])
                if idx < len(doc_metadata):
                    rag_sources.append(doc_metadata[idx].get("source", "PDF"))
        if relevant_chunks:
            rag_context = "\n\n".join([f"Document : {chunk}" for chunk in relevant_chunks])
   
    # Mémoire adaptative
    memory_context = ""
    memory_texts_filtered, memory_scores = retrieve_adaptive_memory(user_message, k=5)
    if memory_texts_filtered:
        memory_context = "\n\n".join([f"Mémoire (score: {score:.2f}): {text}"
                                      for text, score in zip(memory_texts_filtered, memory_scores)])
   
    # Réflexion stratégique
    if request.thinking_mode:
        execute_reflection_plan(
            user_message,
            geo_info=geo,
            messages=request.messages,
            current_subject=conversation_ctx.current_subject,
            subject_keywords=conversation_ctx.subject_keywords
        )
   
    # Recherche Web
    search_query = user_message
    if conversation_ctx.subject_keywords:
        search_query = f"{user_message} {' '.join(conversation_ctx.subject_keywords[:3])} Gabon"
   
    search_results = web_search(search_query)
    web_context = "\n".join([f"- {r['content'][:500]}" for r in search_results.get("results", [])[:6]])
    web_images = search_results.get("images", [])[:4]
   
    # Prompt final
    system_prompt = f"""Tu es Kibali, un assistant IA chaleureux, précis et expert du Gabon, basé à {geo['city']}.
Réponds toujours en français, de façon naturelle, concise et factuelle.
CONTEXTE CONVERSATIONNEL ACTUEL:
- Sujet en cours: {', '.join(conversation_ctx.subject_keywords) if conversation_ctx.subject_keywords else 'Nouveau sujet'}
- Nombre de messages sur ce sujet: {conversation_ctx.message_count}
PRIORITÉ DES SOURCES:
1. Documents uploadés (PDF Vault) - Source la plus fiable
2. Mémoire conversationnelle récente et pertinente
3. Informations Web actualisées
Si une information vient d'un document uploadé, mentionne-le brièvement.
Adapte-toi aux changements brusques de sujet en restant cohérent."""
   
    full_prompt = f"""### INSTRUCTIONS STRICTES :
{system_prompt}
### CONTEXTE DOCUMENTS (PDF Vault) :
{rag_context if rag_context else "Aucun document pertinent trouvé."}
### HISTORIQUE PERTINENT (Mémoire adaptative) :
{memory_context if memory_context else "Pas d'historique pertinent."}
### INFORMATIONS WEB RÉCENTES :
{web_context if web_context else "Pas d'informations web disponibles."}
### QUESTION :
{user_message}
### RÉPONSE (en français uniquement) :
"""
   
    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=8192).to(model.device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=120.0)
   
    def generate_stream():
        try:
            model.generate(
                **inputs,
                streamer=streamer,
                max_new_tokens=1024,
                temperature=0.6,
                do_sample=True,
                top_p=0.85,
                top_k=50,
                repetition_penalty=1.2,
                length_penalty=0.8
            )
        except Exception as e:
            logger.error(f"Erreur génération : {e}")
   
    thread = Thread(target=generate_stream)
    thread.start()
   
    response_text = ""
    for new_text in streamer:
        if new_text is not None:
            response_text += new_text
    response_text = response_text.strip()
   
    if response_text:
        add_to_memory_realtime(
            user_message,
            response_text,
            conversation_ctx.subject_keywords
        )
   
    context_info = {
        "subject_keywords": conversation_ctx.subject_keywords,
        "message_count": conversation_ctx.message_count,
        "memory_used": len(memory_texts_filtered),
        "rag_sources": list(set(rag_sources)),
        "web_results": len(search_results.get("results", []))
    }
   
    return ChatResponse(response=response_text, images=web_images, context_info=context_info)

@app.post("/upload")
async def upload(files: List[UploadFile] = File(...)):
    total_added = 0
    processed_files = 0
    for file in files:
        if not file.filename.lower().endswith(".pdf"):
            continue
        try:
            content = await file.read()
            text = extract_text_from_pdf(content)
            if not text:
                logger.warning(f"Aucun texte extrait de {file.filename}")
                continue
            chunks = chunk_text(text)
            if not chunks:
                continue
            timestamp = datetime.now().isoformat()
            for chunk in chunks:
                doc_metadata.append({
                    "source": file.filename,
                    "timestamp": timestamp,
                    "length": len(chunk)
                })
            embeddings = embed_model.encode(chunks, normalize_embeddings=True).astype('float32')
            doc_index.add(embeddings)
            doc_chunks.extend(chunks)
            total_added += len(chunks)
            processed_files += 1
            logger.info(f"Upload réussi : {file.filename} → {len(chunks)} chunks ajoutés")
        except Exception as e:
            logger.error(f"Erreur lors du traitement de {file.filename} : {e}")
    return {
        "status": "success",
        "files_processed": processed_files,
        "chunks_added": total_added,
        "total_doc_chunks": len(doc_chunks)
    }

@app.post("/upload-pdfs")
async def upload_pdfs(files: List[UploadFile] = File(...)):
    return await upload(files)

@app.post("/clear-memory")
async def clear_memory():
    global memory_index, memory_texts, memory_metadata
    memory_index = faiss.IndexFlatL2(dimension)
    memory_texts = []
    memory_metadata = []
    conversation_ctx.__init__()
    return {"status": "memory_cleared", "message": "Mémoire conversationnelle effacée"}

# --- DEMARRAGE ---
@app.on_event("startup")
async def startup_event():
    logger.info("🚀 Kibali AI API démarrée avec succès sur Hugging Face Spaces !")
    logger.info(f"Accès : https://your-username-your-space.hf.space | Docs : /docs")
    logger.info(f"Mémoire adaptative et réflexion contextuelle activées ✓")