import os
from weaviate.connect import ConnectionParams
from weaviate import WeaviateClient
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_weaviate import WeaviateVectorStore
from langchain_mistralai import ChatMistralAI
from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langserve import add_routes
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
import warnings
import time
from threading import Lock
import logging
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.output_parsers import StrOutputParser
from langchain_core.callbacks.base import BaseCallbackHandler
from langchain.callbacks.manager import CallbackManager
import json
import atexit
import asyncio
from fastapi import WebSocketDisconnect


# Logging config
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Rate limiting
RATE_LIMIT = 3
RATE_WINDOW = 1
last_request_time = 0
request_lock = Lock()

def wait_for_rate_limit():
    global last_request_time
    with request_lock:
        current_time = time.time()
        time_since_last_request = current_time - last_request_time
        if time_since_last_request < RATE_WINDOW / RATE_LIMIT:
            sleep_time = (RATE_WINDOW / RATE_LIMIT) - time_since_last_request
            time.sleep(sleep_time)
        last_request_time = time.time()

def setup_mistral():
    logger.info("Configuration du modèle Mistral...")
    llm = ChatMistralAI(
        # Nom du modèle à utiliser (mistral-large-latest est le plus performant)
        model="mistral-large-latest",
        
        # Clé API pour l'authentification auprès de Mistral AI
        mistral_api_key=os.getenv("MISTRAL_API_KEY"),
        
        # Contrôle la créativité des réponses (0 = très déterministe, 1 = très créatif)
        temperature=0.2,
        
        # Contrôle la diversité des réponses (1 = toutes les options possibles, 0.1 = très sélectif)
        top_p=0.9,
        
        # Pénalité pour la répétition des mots (0 = pas de pénalité, 1 = forte pénalité)
        # frequency_penalty=0.5,
        
        # Pénalité pour encourager la diversité des sujets (0 = pas de pénalité, 1 = forte pénalité)
        # presence_penalty=0.5
    )
    logger.info("Modèle Mistral configuré avec succès")
    return llm

def setup_weaviate():
    logger.info("Connexion à Weaviate...")
    client = WeaviateClient(
        connection_params=ConnectionParams.from_url(
            os.getenv("WEAVIATE_URL"),
            50051
        )
    )
    client.connect()
    logger.info("Connexion à Weaviate établie avec succès")
    
    # Enregistrer la fermeture de la connexion à la sortie
    atexit.register(client.close)
    
    return client

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

def create_custom_prompt():
    template = """Tu es un assistant viticole expert en pulvérisateurs et produits phytosanitaires. \
Voici les documents techniques pertinents pour répondre à la question :

{context}

Question : {question}

Instructions :
1. Réponds de manière factuelle uniquement en te basant sur les documents fournis si cela est pertinent
2. Sois précis sur les caractéristiques techniques
3. Si la question n'est pas liée aux documents techniques, réponds de manière générale et concise


Réponse :"""
    return PromptTemplate(template=template, input_variables=["context", "question"])

def setup_rag_chain():
    logger.info("Initialisation de la chaîne RAG...")
    client = setup_weaviate()
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True}
    )
    vectorstore = WeaviateVectorStore(
        client=client,
        embedding=embeddings,
        index_name="Document",
        text_key="page_content"
    )
    llm = setup_mistral()
    memory = ConversationBufferMemory(
        memory_key="chat_history",
        return_messages=True,
        output_key="answer"
    )
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 30})
    prompt = create_custom_prompt()

    def extract_question(input_data):
        if isinstance(input_data, dict):
            question = input_data.get("question", "")
        else:
            question = str(input_data)
        
        if not question or not isinstance(question, str):
            raise ValueError("La question doit être une chaîne de caractères non vide")
        
        logger.info(f"🔍 Question extraite: {question}")
        return question

    chain = (
        RunnableLambda(extract_question)
        | {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
    )
    
    logger.info("✅ Chaîne RAG initialisée avec succès")
    return chain

# Init FastAPI
app = FastAPI(title="Vitizen Chat API", description="API de chat pour l'assistant viticole", version="1.0.1")

@app.get("/", response_class=HTMLResponse)
async def home():
    return """
    <html><body><h1>Bienvenue sur l'API Vitizen Chat</h1></body></html>
    """

# Create RAG chain
rag_chain = setup_rag_chain()
add_routes(app, rag_chain, path="/chat")

@app.websocket("/chat/ws")
async def stream_chat(websocket: WebSocket):
    await websocket.accept()
    logger.info("📡 WebSocket /chat/ws connecté")

    try:
        while True:
            try:
                data = await websocket.receive_json()
                question = data.get("message", "").strip()
                if not question:
                    logger.warning("⚠️ Message vide reçu")
                    await websocket.send_text("[START]")
                    await websocket.send_text("[Erreur] Message vide")
                    await websocket.send_text("[STOP]")
                    continue

                logger.info(f"📝 Question reçue: {question}")
                input_data = {"question": question}
                logger.info(f"🔍 Entrée préparée pour la chaîne RAG: {input_data}")

                # Envoyer le marqueur START
                await websocket.send_text("[START]")
                logger.info("🔄 Début du streaming de la réponse")

                # Utiliser astream pour le streaming sans callback
                async for chunk in rag_chain.astream(input_data):
                    if chunk:
                        logger.info(f"📦 Chunk reçu: {chunk}")
                        try:
                            if hasattr(chunk, 'content'):
                                content = chunk.content
                                if content:
                                    logger.info(f"📤 Envoi du contenu: {content}")
                                    await websocket.send_text(content)
                            elif isinstance(chunk, dict) and "answer" in chunk:
                                logger.info(f"📤 Envoi de la réponse: {chunk['answer']}")
                                await websocket.send_text(chunk["answer"])
                            else:
                                logger.info(f"📤 Envoi du chunk: {str(chunk)}")
                                await websocket.send_text(str(chunk))
                        except Exception as e:
                            logger.error(f"❌ Erreur lors de l'envoi du chunk: {str(e)}")
                            await websocket.send_text("[STOP]")
                            break

                # Envoyer le marqueur STOP
                await websocket.send_text("[STOP]")
                logger.info("✅ Streaming terminé avec succès")

            except WebSocketDisconnect:
                logger.info("🔌 Client déconnecté proprement.")
                break

            except Exception as e:
                logger.error(f"❌ Erreur pendant traitement du message: {e}")
                try:
                    await websocket.send_text("[START]")
                    await websocket.send_text("[Erreur] Erreur API")
                    await websocket.send_text("[STOP]")
                except:
                    pass

    except Exception as e:
        logger.error(f"❌ Erreur fatale: {str(e)}")
    finally:
        try:
            if not websocket.client_state.disconnected:
                await websocket.send_text("[STOP]")
                await websocket.close()
        except:
            pass
        logger.info("🔌 Connexion WebSocket fermée")

@app.get("/health")
async def health_check():
    return {"status": "ok"}