chat.py

from fastapi import APIRouter, Request, HTTPException, Depends
from fastapi.responses import JSONResponse, StreamingResponse
from datetime import datetime
from bson.objectid import ObjectId
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import re
import json
import asyncio
from auth import get_current_user
from database import get_db
from config import HF_TOKEN, MAX_TOKENS, EMBEDDING_MODEL

router = APIRouter(prefix="/api", tags=["Chat"])
db=get_db()
conversation_history = {}
hf_client = InferenceClient(token=HF_TOKEN)


def save_bot_response(conversation_id, current_user, text, current_tokens=0, message_tokens=0):
    if not conversation_id or not current_user:
        print(" Impossible de sauvegarder la réponse")
        return None
    
    try:
        message_id = db.messages.insert_one({
            "conversation_id": conversation_id,
            "user_id": str(current_user["_id"]),
            "sender": "bot", 
            "text": text,
            "timestamp": datetime.utcnow()
        }).inserted_id
        
        response_tokens = int(len(text.split()) * 1.3) if text else 0
        total_tokens = current_tokens + message_tokens + response_tokens
        
        db.conversations.update_one(
            {"_id": ObjectId(conversation_id)},
            {"$set": {
                "last_message": text[:100] + ("..." if len(text) > 100 else ""),
                "updated_at": datetime.utcnow(),
                "token_count": total_tokens
            }}
        )
        
        print(f"Réponse du bot sauvegardée : {message_id}")
        return message_id
    except Exception as e:
        print(f" Erreur: {str(e)}")
        return None


try:
    from langchain_community.embeddings import HuggingFaceEmbeddings
    embedding_model = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
    print("Modèle d'embedding médical chargé avec succès")
except Exception as e:
    print(f"Erreur chargement embedding: {str(e)}")
    embedding_model = None

# Fonctions de RAG
def retrieve_relevant_context(query, embedding_model, mongo_collection, k=5):
    query_embedding = embedding_model.embed_query(query)
    
    docs = list(mongo_collection.find({}, {"text": 1, "embedding": 1}))
    
    print(f"[DEBUG] Recherche de contexte pour: '{query}'")
    print(f"[DEBUG] {len(docs)} documents trouvés dans la base de données")
    
    if not docs:
        print("[DEBUG] Aucun document dans la collection. RAG désactivé.")
        return ""
    
    similarities = []
    for i, doc in enumerate(docs):
        if "embedding" not in doc or not doc["embedding"]:
            print(f"[DEBUG] Document {i} sans embedding")
            continue
            
        sim = cosine_similarity([query_embedding], [doc["embedding"]])[0][0]
        similarities.append((sim, i, doc["text"]))
    
    similarities.sort(reverse=True)
    
    print("\n=== CONTEXTE SÉLECTIONNÉ ===")
    top_k_docs = []
    for i, (score, idx, text) in enumerate(similarities[:k]):
        doc_preview = text[:100] + "..." if len(text) > 100 else text
        print(f"Document #{i+1} (score: {score:.4f}): {doc_preview}")
        top_k_docs.append(text)
    print("==========================\n")
    
    return "\n\n".join(top_k_docs)

@router.post("/chat")
async def chat(request: Request):
    global conversation_history
    
    data = await request.json()
    user_message = data.get("message", "").strip()
    conversation_id = data.get("conversation_id")
    skip_save = data.get("skip_save", False)

    if not user_message:
        raise HTTPException(status_code=400, detail="Le champ 'message' est requis.")

    current_user = None
    try:
        current_user = await get_current_user(request)
    except HTTPException:
        pass

    if not skip_save and conversation_id and current_user: 
        db.messages.insert_one({
            "conversation_id": conversation_id,
            "user_id": str(current_user["_id"]),
            "sender": "user",
            "text": user_message,
            "timestamp": datetime.utcnow()
        })

    current_tokens = 0
    message_tokens = 0
    if current_user and conversation_id:
        conv = db.conversations.find_one({
            "_id": ObjectId(conversation_id),
            "user_id": str(current_user["_id"])
        })
        if conv:
            current_tokens = conv.get("token_count", 0)
        message_tokens = int(len(user_message.split()) * 1.3)
        if current_tokens + message_tokens > MAX_TOKENS:
            error_message = "⚠️ **Limite de taille de conversation atteinte**\n\nCette conversation est devenue trop longue. Pour continuer à discuter, veuillez créer une nouvelle conversation."
            
            if conversation_id and current_user:
                save_bot_response(conversation_id, current_user, error_message, current_tokens, message_tokens)
            
            return JSONResponse({
                "error": "token_limit_exceeded",
                "message": error_message,
                "tokens_used": current_tokens,
                "tokens_limit": MAX_TOKENS
            }, status_code=403)

    is_history_question = any(
        phrase in user_message.lower()
        for phrase in [
            "ma première question", "ma précédente question", "ma dernière question",
            "ce que j'ai demandé", "j'ai dit quoi", "quelles questions",
            "c'était quoi ma", "quelle était ma", "mes questions", "questions précédentes"
        ]
    ) or re.search(r"(?:quelle|quelles|quoi).*?(\d+)[a-z]{2}.*?question", user_message.lower()) \
       or re.search(r"derni[eè]re question", user_message.lower()) \
       or re.search(r"premi[eè]re question", user_message.lower()) \
       or re.search(r"question pr[eé]c[eé]dente", user_message.lower()) \
       or re.search(r"(toutes|liste|quelles|quoi).*questions", user_message.lower())

    if conversation_id not in conversation_history:
        conversation_history[conversation_id] = []
        if current_user and conversation_id:
            previous_messages = list(db.messages.find(
                {"conversation_id": conversation_id}
            ).sort("timestamp", 1))
            
            for msg in previous_messages:
                if msg["sender"] == "user":
                    conversation_history[conversation_id].append(f"Question : {msg['text']}")
                else:
                    conversation_history[conversation_id].append(f"Réponse : {msg['text']}")

    if is_history_question:
        actual_questions = []
        
        if conversation_id in conversation_history:
            for msg in conversation_history[conversation_id]:
                if msg.startswith("Question : "):
                    q_text = msg.replace("Question : ", "")
                    is_meta = any(phrase in q_text.lower() for phrase in [
                        "ma première question", "ma précédente question", "ma dernière question",
                        "ce que j'ai demandé", "j'ai dit quoi", "quelles questions",
                        "c'était quoi ma", "quelle était ma", "mes questions"
                    ]) or re.search(r"(?:quelle|quelles|quoi).*?(\d+)[a-z]{2}.*?question", q_text.lower()) \
                    or re.search(r"derni[eè]re question", q_text.lower()) \
                    or re.search(r"premi[eè]re question", q_text.lower()) \
                    or re.search(r"question pr[eé]c[eé]dente", q_text.lower()) \
                    or re.search(r"(toutes|liste|quelles|quoi).*questions", q_text.lower())
                    if not is_meta:
                        actual_questions.append(q_text)
        
        history_response = ""
        
        if not actual_questions:
            history_response = "Vous n'avez pas encore posé de question dans cette conversation. C'est notre premier échange."
        else:
            
            if any(phrase in user_message.lower() for phrase in ["question précédente", "dernière question"]) and len(actual_questions) > 1:
                prev_question = actual_questions[-1] if actual_questions else "Aucune question précédente trouvée."
                history_response = f"**Votre question précédente était :**\n\n\"{prev_question}\""
            
            elif any(phrase in user_message.lower() for phrase in ["première question", "1ère question", "1ere question"]):
                first_question = actual_questions[0] if actual_questions else "Aucune première question trouvée."
                history_response = f"**Votre première question était :**\n\n\"{first_question}\""
            
            elif re.search(r"(\d+)[èeme]{1,3}", user_message.lower()):
                match = re.search(r"(\d+)[èeme]{1,3}", user_message.lower())
                if match:
                    question_num = int(match.group(1))
                    if 0 < question_num <= len(actual_questions):
                        specific_question = actual_questions[question_num-1]
                        history_response = f"**Votre question n°{question_num} était :**\n\n\"{specific_question}\""
                    else:
                        history_response = f"Je ne trouve pas de question n°{question_num} dans notre conversation. Vous n'avez posé que {len(actual_questions)} question(s)."
            
            else:
                history_response = "**Voici les questions que vous avez posées dans cette conversation :**\n\n"
                for i, question in enumerate(actual_questions, 1):
                    history_response += f"{i}. {question}\n"
                    
                if len(actual_questions) > 3:
                    history_response += f"\nVous avez posé {len(actual_questions)} questions dans cette conversation."

        if conversation_id:
            conversation_history[conversation_id].append(f"Réponse : {history_response}")
            
        if conversation_id and current_user:
            save_bot_response(conversation_id, current_user, history_response, current_tokens, message_tokens)
            print(f"Réponse à la question d'historique sauvegardée pour conversation {conversation_id}")
        
        return JSONResponse({"response": history_response})

    if conversation_id:
        conversation_history[conversation_id].append(f"Question : {user_message}")

    context = None
    if not is_history_question and embedding_model:
        context = retrieve_relevant_context(user_message, embedding_model, db.connaissances, k=5)
        if context and conversation_id:
            conversation_history[conversation_id].append(f"Contexte : {context}")

    system_prompt = (
        "Tu es un chatbot spécialisé dans la santé mentale, et plus particulièrement la schizophrénie. "
        "Tu réponds de façon fiable, claire et empathique, en t'appuyant uniquement sur des sources médicales et en français. "
        "IMPORTANT: Fais particulièrement attention aux questions de suivi. Si l'utilisateur pose une question qui ne précise "
        "pas clairement le sujet mais qui fait suite à votre échange précédent, comprends que cette question fait référence "
        "au contexte de la conversation précédente. Par exemple, si l'utilisateur demande 'Comment les traite-t-on?' après "
        "avoir parlé des symptômes positifs de la schizophrénie, ta réponse doit porter spécifiquement sur le traitement "
        "des symptômes positifs, et non sur la schizophrénie en général. IMPORTANT: Vise tes réponses sous forme de Markdown."
        "IMPÉRATIF: Structure tes réponses en Markdown, utilisant **des gras** pour les points importants, "
    "des titres avec ## pour les sections principales, des listes à puces avec * pour énumérer des points, "
    "et > pour les citations importantes. Cela rend ton contenu plus facile à lire et à comprendre."
    
    )

    enriched_context = ""

    if conversation_id in conversation_history:
        actual_questions = []
        for msg in conversation_history[conversation_id]:
            if msg.startswith("Question : "):
                q_text = msg.replace("Question : ", "")
                is_meta = any(phrase in q_text.lower() for phrase in [
                    "ma première question", "ma précédente question", "ma dernière question",
                    "ce que j'ai demandé", "j'ai dit quoi", "quelles questions",
                    "c'était quoi ma", "quelle était ma", "mes questions"
                ]) or re.search(r"(?:quelle|quelles|quoi).*?(\d+)[a-z]{2}.*?question", q_text.lower()) \
                   or re.search(r"derni[eè]re question", q_text.lower()) \
                   or re.search(r"premi[eè]re question", q_text.lower()) \
                   or re.search(r"question pr[eé]c[eé]dente", q_text.lower()) \
                   or re.search(r"(toutes|liste|quelles|quoi).*questions", q_text.lower())
                if not is_meta and q_text != user_message:  
                    actual_questions.append(q_text)
        
        if actual_questions:
            recent_questions = actual_questions[-5:] 
            enriched_context += "Historique récent des questions:\n"
            for i, q in enumerate(recent_questions):
                enriched_context += f"- Question précédente {len(recent_questions)-i}: {q}\n"
            enriched_context += "\n"

    if context:
        enriched_context += "Contexte médical pertinent:\n"
        enriched_context += context
        enriched_context += "\n\n"

    if enriched_context:
        system_prompt += (
            f"\n\n{enriched_context}\n\n"
            "Utilise ces informations pour répondre de manière plus précise et contextuelle. "
            "Ne pas inventer d'informations. Si tu ne sais pas, redirige vers un professionnel de santé. "
            "Tu dois donner une réponse complète, bien structurée et ne jamais couper ta réponse brutalement. "
            "Si tu n'as pas assez de place pour finir, indique-le clairement à l'utilisateur."
        )
    else:
        system_prompt += (
            "Tu dois répondre uniquement à partir de connaissances médicales factuelles. "
            "Si tu ne sais pas répondre, indique-le clairement et suggère de consulter un professionnel de santé. "
            "Tu dois donner une réponse complète et bien structurée."
        )

    messages = [{"role": "system", "content": system_prompt}]
    
    if conversation_id and len(conversation_history.get(conversation_id, [])) > 0:
        history = conversation_history[conversation_id]
        
        user_messages = []
        bot_messages = []
        
        for i in range(len(history)):
            if i < len(history) and history[i].startswith("Question :"):
                user_text = history[i].replace("Question : ", "")
                user_messages.append(user_text)
                
            if i+1 < len(history) and history[i+1].startswith("Réponse :"):
                bot_text = history[i+1].replace("Réponse : ", "")
                bot_messages.append(bot_text)
        
        valid_pairs = min(len(user_messages), len(bot_messages))
        
        for i in range(valid_pairs):
            messages.append({"role": "user", "content": user_messages[i]})
            messages.append({"role": "assistant", "content": bot_messages[i]})
    
    messages.append({"role": "user", "content": user_message})

    async def generate_stream():
        try:
            collected_response = ""
            
            yield "data: {\"type\": \"start\"}\n\n"
            
            completion_stream = hf_client.chat.completions.create(
                model="mistralai/Mistral-7B-Instruct-v0.3",
                messages=messages,
                max_tokens=1024,
                temperature=0.7,
                stream=True
            )
            chunk_buffer = ""
            chunk_count = 0
            MAX_CHUNKS_BEFORE_SEND = 3 
            for chunk in completion_stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    content = chunk.choices[0].delta.content
                    collected_response += content
                    
                    # Envoyer chaque token individuellement sans buffering
                    yield f"data: {json.dumps({'content': content})}\n\n"
                    
                    # Petit sleep pour éviter le buffering par le serveur ASGI
                    await asyncio.sleep(0)
            
            if collected_response.endswith((".", "!", "?")) == False and len(collected_response) > 500:
                suffix = "\n\n(Note: Ma réponse a été limitée par des contraintes de taille. N'hésitez pas à me demander de poursuivre si vous souhaitez plus d'informations.)"
                collected_response += suffix
                yield f"data: {json.dumps({'content': suffix})}\n\n"
            
            if conversation_id:
                conversation_history[conversation_id].append(f"Réponse : {collected_response}")
                
                if len(conversation_history[conversation_id]) > 50: 
                    conversation_history[conversation_id] = conversation_history[conversation_id][-50:]
            
            if conversation_id and current_user:
                save_bot_response(conversation_id, current_user, collected_response, current_tokens, message_tokens)
            
            yield "data: {\"type\": \"end\"}\n\n"
            
        except Exception as e:
            error_message = str(e)
            print(f"❌ Streaming error: {error_message}")
            
            try:
                fallback = hf_client.text_generation(
                    model="mistralai/Mistral-7B-Instruct-v0.3",
                    prompt=f"<s>[INST] {system_prompt}\n\nQuestion: {user_message} [/INST]",
                    max_new_tokens=512,
                    temperature=0.7
                )
                yield f"data: {json.dumps({'content': fallback})}\n\n"
                
                if conversation_id:
                    conversation_history[conversation_id].append(f"Réponse : {fallback}")
                
                if conversation_id and current_user:
                    save_bot_response(conversation_id, current_user, fallback, current_tokens, message_tokens)
                    
            except Exception as fallback_error:
                print(f" Erreur: {str(fallback_error)}")
                error_response = "Je suis désolé, je rencontre actuellement des difficultés techniques"
                yield f"data: {json.dumps({'content': error_response})}\n\n"
                
                if conversation_id and current_user:
                    save_bot_response(conversation_id, current_user, error_response, current_tokens, message_tokens)
            
            yield "data: {\"type\": \"end\"}\n\n"
    
    return StreamingResponse(
        generate_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache, no-transform",
            "X-Accel-Buffering": "no"  # Important pour Nginx
        }
)