import streamlit as st
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings, FastEmbedEmbeddings
from langchain_core.output_parsers import StrOutputParser
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import PromptTemplate   # ✅ updated
from langchain_community.vectorstores.utils import filter_complex_metadata
from langchain.chains import LLMChain
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.messages import HumanMessage
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
import pandas as pd
import chardet
import os

# ---------- Classe Document ----------
class Document:
    def __init__(self, page_content, metadata=None):
        self.page_content = page_content
        self.metadata = metadata or {}


# ---------- Classe Chat principale ----------

class Chat:
    vector_store = None
    retriever = None
    chain = None

    def __init__(self, file_path):
        # LLM conversationnel
        llm = HuggingFaceEndpoint(
            repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1",  # ou Mistral-7B-Instruct-v0.2
            task="text-generation",
            temperature=0.1,
            max_new_tokens=800,
            
        )

        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1024, chunk_overlap=150, length_function=len
        )
        self.model = ChatHuggingFace(llm=llm)

        # Split des documents
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1024, chunk_overlap=150, length_function=len
        )

        # Prompt pour le raisonnement médical
        prompt_template = """
<s> [INST] Vous êtes un assistant intelligent et expert en sciences médicales. 
Votre tâche est d'analyser précisément les descriptions fournies et de rechercher le "code acte" correspondant 
dans le fichier (contexte) fourni. Vous devez suivre ces instructions pour chaque question :

1. Correspondance exacte : Si vous trouvez une correspondance exacte pour la description dans le fichier, 
renvoyez le "code acte" avec sa description, formatés en Markdown.

2. Correspondances proches : Si aucune correspondance exacte n'est trouvée, citez jusqu'à 4 codes actes 
les plus proches de la description, avec leurs descriptions respectives, également formatés en Markdown. 
Expliquez en détail pourquoi ces codes ont été suggérés, en précisant :
   - Type de procédure (diagnostique, chirurgicale, thérapeutique, etc.)
   - Emplacement anatomique concerné
   - Type de dispositif (drain, cathéter, implant, etc.)
   - Technique ou méthode utilisée (fluoroscopie, laparoscopie, etc.)

3. Aucune correspondance pertinente : Si aucune correspondance pertinente n'est trouvée, 
indiquez clairement que le code est "Introuvable".

4. Méthode de recherche : Basez votre recherche sur une correspondance exacte de la description 
ou une technique similaire utilisée. Indiquez les critères de similarité.

### Exemple de Réponse :
Pour la description **"groupage sanguin"** :
- **Code acte** : `B229`
- **Description** : "CROSS MATCH : GROUPE ABO ET RHESUS"

Répondez toujours en Markdown. [/INST] </s>
[INST]
**Question** : {question}
**Contexte** : {context}
**Réponse (format Markdown)** : [/INST]
"""
        self.prompt = PromptTemplate.from_template(prompt_template)

        # Charger et vectoriser le fichier CSV
        self.ingest(file_path)

    # ---------- Ingestion des données ----------
    def ingest(self, file_path: str):
        import chardet
        import pandas as pd
        from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
        from langchain.vectorstores import Chroma

        # Détection automatique de l'encodage
        with open(file_path, 'rb') as f:
            result = chardet.detect(f.read())
            charenc = result['encoding']

        # Lecture CSV
        df = pd.read_csv(file_path, encoding=charenc, on_bad_lines='warn')

        # Conversion en documents
        docs = [Document(page_content=str(row.dropna().to_dict()), metadata={})
                for index, row in df.iterrows()]

        # Embeddings
        embedding_doc = FastEmbedEmbeddings(model_name="intfloat/multilingual-e5-large")

        self.vector_store = Chroma.from_documents(
            documents=docs,
            embedding=embedding_doc,
            persist_directory="./chroma_db"
        )
        self.vector_store.persist()

        k = min(30, len(docs))
        self.retriever = self.vector_store.as_retriever(
            search_type="similarity",  # Changed search type to similarity
            search_kwargs={"k": k}
        )

    # ---------- Fonction d'interaction ----------
    def ask(self, query: str):
        # Récupérer les documents pertinents
        docs = self.retriever.get_relevant_documents(query)
        context_text = "\n".join([doc.page_content for doc in docs]) or "Aucun code pertinent trouvé dans le fichier."

        # Construire le prompt avec le contexte
        prompt_input = self.prompt.format(question=query, context=context_text)

        # Envoyer au modèle conversationnel
        from langchain.schema import HumanMessage
        response = self.model([HumanMessage(content=prompt_input)])

        return response.content

    # ---------- Reset ----------
    def clear(self):
        self.vector_store = None
        self.retriever = None
        self.chain = None

# ---------- Interface Streamlit ----------
def main():
    st.title("🧠 Assistant Médical - Recherche de Code Acte")

    file_path = "actes.csv"

    @st.cache_resource
    def load_chat(file_path):
        return Chat(file_path)

    chat = load_chat(file_path)
    st.success(f"✅ Fichier '{file_path}' chargé et vectorisé avec succès !")

    user_input = st.text_input("🩺 Posez votre question médicale :", key="user_input")

    if st.button("Analyser"):
        if user_input:
            with st.spinner("💬 Analyse en cours..."):
                response = chat.ask(user_input)

                start_marker = "**Réponse (format Markdown)** : [/INST]"
                start_pos = response.find(start_marker)

                if start_pos == -1:
                    st.markdown(f"**Assistant :**\n\n{response}")
                else:
                    response_text = response[start_pos + len(start_marker):].strip()
                    st.markdown(f"**Assistant :**\n\n{response_text}")


if __name__ == "__main__":
    main()