import streamlit as st
import os
from chromadb import Client as ChromaClient
from chroma_storage import ChromaStorage
from retrieval import Retriever
from langchain_groq import ChatGroq

from generation import RAGGenerator
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from constant import categorie
from dotenv import load_dotenv

import chromadb
from Reranker import Reranker 

chromadb.api.client.SharedSystemClient.clear_system_cache()
load_dotenv()
import os





CHROMA_DIR      = "src/chroma_db"
COLLECTION_NAME = "my_collection"

st.markdown("""
<style>
body {
    font-family: 'Segoe UI', Roboto, sans-serif;
}
.header {
    background-color: #4A90E2;
    color: white;
    padding: 10px;
    display: flex;
    align-items: center;
    box-shadow: 0 2px 4px rgba(0,0,0,0.2);
}
.header img {
    height: 40px;
    margin-right: 10px;
}
.header h1 {
    margin: 0;
    font-size: 24px;
}
.chat-container {
    padding: 10px;
}
.bot-message, .user-message {
    display: flex;
    align-items: flex-start;
    margin: 8px 0;
}
.bot-message .bubble {
    margin-left: 10px;
    background-color: #f0f0f0;
    color: #000;
}
.user-message {
    justify-content: flex-end;
}
.user-message .bubble {
    margin-right: 10px;
    background-color: #4A90E2;
    color: #fff;
}
.bubble {
    border-radius: 8px;
    padding: 10px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    max-width: 70%;
}
.avatar {
    width: 40px;
    height: 40px;
    border-radius: 50%;
}
@media (max-width: 600px) {
    .header h1 {
        font-size: 18px;
    }
    .bubble {
        max-width: 90%;
    }
}
</style>
""", unsafe_allow_html=True)

# --- En-tête bleu avec logo et titre ---
st.markdown(f'''
<div class="header">
  <img src="https://external-content.duckduckgo.com/iu/?u=https%3A%2F%2Fmedia.licdn.com%2Fdms%2Fimage%2FD4D0BAQFNZTx4WjU6Kw%2Fcompany-logo_200_200%2F0%2F1706688195351%3Fe%3D2147483647%26v%3Dbeta%26t%3DCBoRbqESfYGI0cxj1-5IvCpkhBTRcXEbBqS6q7rA9l0&f=1&nofb=1&ipt=a3414940fc3b31a18d97ba116cc4b5a9bc4f88e8d72bf286c029be996d8d5b0e">
  <h1>car insurance assistant</h1>
</div>
''', unsafe_allow_html=True)


#if "categories" not in st.session_state:
 #   st.session_state.categories = categorie
if "messages" not in st.session_state:
    st.session_state.messages = []
if "step" not in st.session_state:
    st.session_state.step = 0
if "selected_category" not in st.session_state:
    st.session_state.selected_category = None


# --- Affichage du chat uniquement après sélection ---
#sel = st.session_state.selected_category



# Ajouter le message initial du bot à l'étape 1
if st.session_state.step == 0:
    st.session_state.messages.append({
        "role": "assistant", 
        "content": "I am your assistant, Ominimo"
    })
  #  save_conversation_to_mongo(st.session_state.messages)
    st.session_state.step = 1
    st.rerun()



# ─────────────────────────────────────────────────────────────────────────────
# 2) INITIALISATION UNIQUE (cache via session_state)
# ─────────────────────────────────────────────────────────────────────────────
if "initialized" not in st.session_state:
    st.session_state.initialized = True
    retriever = Retriever(
        chroma_dir=CHROMA_DIR,
        collection_name=COLLECTION_NAME
       
    )
   
   
    llm = ChatGroq(
    model="meta-llama/llama-4-scout-17b-16e-instruct",  
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=2,

)
    generator = RAGGenerator()
    reranker = Reranker(model=llm)
    st.session_state.retriever   = retriever
    st.session_state.generator   = generator
    st.session_state.reranker    = reranker
   # st.session_state.categories  = categorie
   

# ─────────────────────────────────────────────────────────────────────────────
# 3) UI & État de chat
# ─────────────────────────────────────────────────────────────────────────────
query = st.chat_input("chat_query")
#if "last_rerank" in st.session_state:
   # st.code(f"Rerank result: {st.session_state.last_rerank}")

# Si l'utilisateur a saisi une requête
if query:
    # Si on est en phase de clarification
    if "qa_pairs" in st.session_state and "clarif_idx" in st.session_state:
        idx = st.session_state.clarif_idx
        st.session_state.messages.append({"role": "user", "content": query})
      #  save_conversation_to_mongo(st.session_state.messages)
        st.session_state.qa_pairs[idx]["response"] = query
        st.session_state.clarif_idx += 1
    else:
        # Requête initiale normale
        st.session_state.messages.append({"role": "user", "content": query})
      #  save_conversation_to_mongo(st.session_state.messages)
        try:
            # 1. Récupération des documents initiaux
           # with st.spinner():
            docs = st.session_state.retriever.retrieve(query)
                

            passages = [(doc_tuple[0].page_content, doc_tuple[1]) for doc_tuple in docs]


           # st.code(passages)
            
            # 2. Reranking des documents pour améliorer la pertinence
            if docs:
                rerank_result = st.session_state.reranker.rerank(query, passages)
                
                # Affichage du résultat brut pour debug
               # st.code(f"Rerank raw result: {rerank_result}")
                if rerank_result is not None:
                    st.session_state.last_rerank = rerank_result.model_dump()
                    
                    # Filtrage et réorganisation des documents selon le reranking
                    if rerank_result.top_indexes:
                        reranked_docs = [docs[i] for i in rerank_result.top_indexes if i < len(docs)]
                        docs = reranked_docs
                    else:
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": rerank_result.message or "Aucun document pertinent trouvé. Veuillez reformuler votre question."
                        })
                        

                # Si le reranker a échoué
                # Si le reranker a échoué
                else:
                    pass

                
                

                

            
                      



                    
            # 3. Génération de la réponse avec les documents reranked
            if docs:  # Vérifier que nous avons des documents avant de continuer
                qa_pairs, solution, lang = st.session_state.generator.retrieve_qa(
                    query=query,
                    docs_scores=docs
                )
                st.session_state.qa_pairs = qa_pairs
                st.session_state.solution = solution
                st.session_state.lang = lang
            
                
            else:
                # Aucun document disponible
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": "Aucun document pertinent trouvé pour votre requête. Pouvez-vous la reformuler ?"
                })
                  
                
        except ValueError as e:
            st.session_state.messages.append({
                "role": "assistant",
                "content": str(e)
            })
         #   save_conversation_to_mongo(st.session_state.messages)
        except Exception as e:
            st.error(f"❌ Erreur : {e}")

# ───────────────────────────────────────────────
# Gestion des clarifications
# ───────────────────────────────────────────────
if "qa_pairs" in st.session_state:
    if "clarif_shown" not in st.session_state:
        questions_text = "\n".join([
            f"{i+1}. {pair['question']}" for i, pair in enumerate(st.session_state.qa_pairs)
        ])
        st.session_state.messages.append({
            "role": "assistant",
            "content": questions_text
        })
      #  save_conversation_to_mongo(st.session_state.messages)
        st.session_state.clarif_shown = True

    if "clarif_idx" not in st.session_state:
        st.session_state.clarif_idx = 0

    # Lorsque toutes les clarifications sont remplies, générer la réponse finale
    if st.session_state.clarif_idx >= len(st.session_state.qa_pairs):
        final = st.session_state.generator.generate_answer(
            st.session_state.messages,
            st.session_state.solution,
            st.session_state.lang
        )

        if not final.strip() or "sorry" in final.lower():
            st.session_state.messages.append({
                "role": "assistant",
                "content": "I'm sorry, I cannot answer this request. It is beyond my capabilities or not related to car insurance."
            })
        else:
            st.session_state.messages.append({"role": "assistant", "content": final})    

    for key in ["qa_pairs", "solution", "lang", "clarif_idx", "clarif_shown"]:
        st.session_state.pop(key, None)

    
        for key in ["qa_pairs", "solution", "lang", "clarif_idx", "clarif_shown"]:
            st.session_state.pop(key, None)


st.markdown('<div class="chat-container">', unsafe_allow_html=True)
for msg in st.session_state.messages:
    if not msg.get("content"):
        continue
    if msg["role"] == "assistant":
        st.markdown(f'''
        <div class="bot-message">
          <img src="https://external-content.duckduckgo.com/iu/?u=https%3A%2F%2Fmedia.licdn.com%2Fdms%2Fimage%2FD4D0BAQFNZTx4WjU6Kw%2Fcompany-logo_200_200%2F0%2F1706688195351%3Fe%3D2147483647%26v%3Dbeta%26t%3DCBoRbqESfYGI0cxj1-5IvCpkhBTRcXEbBqS6q7rA9l0&f=1&nofb=1&ipt=a3414940fc3b31a18d97ba116cc4b5a9bc4f88e8d72bf286c029be996d8d5b0e" class="avatar">
          <div class="bubble">{msg["content"]}</div>
        </div>
        ''', unsafe_allow_html=True)
    else:
        st.markdown(f'''
        <div class="user-message">
          <div class="bubble">{msg["content"]}</div>
          <img src="https://img.icons8.com/ios-filled/50/4A90E2/user-male-circle.png" class="avatar">
        </div>
        ''', unsafe_allow_html=True)
st.markdown('</div>', unsafe_allow_html=True)