import asyncio
from src.agents.qdrantretrieveuserproxyagent import TrackableQdrantRetrieveUserProxyAgent
from src.agents.retrieveassistantagent import TrackableRetrieveAssistantAgent
import streamlit as st
from qdrant_client import QdrantClient
import glob
import os
from sentence_transformers import SentenceTransformer

class MultiAgentRAGChat:
    def __init__(self, assistant_name, user_proxy_name, llm_config, problem):
        self.assistant = TrackableRetrieveAssistantAgent(name=assistant_name,
                                                 system_message="""you are helpful assistant. Reply "TERMINATE" in 
                                                 the end when everything is done """,
                                                 human_input_mode="NEVER",
                                                 llm_config=llm_config,
                                                 )
        self.user_proxy = TrackableQdrantRetrieveUserProxyAgent(name=user_proxy_name,
                                                    human_input_mode="NEVER",
                                                    max_consecutive_auto_reply=4,
                                                    retrieve_config={
                                                        "task": "code",
                                                        #"docs_path": self.list_files(st.session_state["docs_path"]), This is for loading custion files from directory - logic implemented
                                                        "docs_path": st.session_state["docs_path"],
                                                        "chunk_token_size": 500,
                                                        "model": llm_config["config_list"][0]["model"],
                                                        "client": QdrantClient(":memory:"),
                                                        "embedding_model": self.embeddings_model()
                                                    },
                                                    code_execution_config=False,
                                                    is_termination_msg=lambda x: x.get("content", "").strip().endswith(
                                                      "TERMINATE"))
        self.problem = problem
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        
        
    
    def embeddings_model(self):
        sentence_transformer_ef = SentenceTransformer("all-distilroberta-v1").encode    
        return sentence_transformer_ef
   

    def list_files(self,directory):
        # Ensure the directory path ends with a slash
        if not directory.endswith('/'):
            directory += '/'
        try :
            
            # Use glob to get the list of files
            files = glob.glob(os.path.join(directory, '*'))
            file_list = [path.replace('\\', '/') for path in files]
            
            if file_list.count == 0:
                raise ValueError('list of files is zero')
            
        except Exception as ex:
            raise ValueError('issue with file path')
            
        
        return file_list



    async def initiate_chat(self):
        await self.user_proxy.a_initiate_chat(self.assistant,  message=self.user_proxy.message_generator, problem=self.problem, clear_history=st.session_state["chat_with_history"])

    def run(self):
        self.loop.run_until_complete(self.initiate_chat())