from typing import List
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from extract import extract_list
from openai import OpenAI
from pinecone import Pinecone
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.chains.question_answering import load_qa_chain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate

import os, json, random
from dotenv import load_dotenv
import gradio as gr

load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")
pinecone_index = os.getenv("PINECONE_INDEX")
pinecone_api_key = os.getenv("PINECONE_API_KEY")

openai_client = OpenAI(api_key=openai_api_key)

practice_list = ', '.join(extract_list())

description = "Extract mediator's practice field that user want to search. Available mediator practice fields are " + practice_list
metadata_list = ['fullname', 'mediator profile on mediate.com', 'mediator Biography', 'mediator state']
metadata_value = ['Name', "Profile", "Biography", "State"]
class MediatorRetriever(BaseRetriever):
    def getMetadata(self, message):
        tools = [
            {
                "type": "function",
                "name": "get_info",
                "description": "Extract the information of mediator.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "mediator state": {
                            "type": "string",
                            "description": "Extract mediator's state that user wants to search. If both mediator city and mediator state are possible, please extract as mediator state."
                        },
                        "mediator country": {
                            "type": "string",
                            "description": "Extract mediator's country that user wants to search."
                        },
                        "mediator city": {
                            "type": "string",
                            "description": "Extract mediator's city that user wants to search."
                        },
                        "mediator areas of practice": {
                            "type": "array",
                            "description": description,
                            "enum": extract_list()
                        }
                    }
                }
            }
        ]
        
        response = openai_client.chat.completions.create(
            model="gpt-4-1106-preview",
            messages=[
                {"role": "system", "content": "You are a professional arbitration field searcher. Your role is to extract information about the moderator from the user's message. If you have less than 3 mediators in a particular city, you should search for mediators in your state. If you have less than 3 mediators in a particular state that you are searching for, you should search for mediators in the states closest to the state that you are searching for.  For example, if you have less than 3 mediators in California, you should search for mediators in Oregon, Arizona, or Nevada closest to California. Important: You must provide at least three moderators."},
                {"role": "user", "content": message}
            ],
            functions=tools,
            stream=True
        )
        new_message = ""
        for chunk in response:
            delta = chunk.choices[0].delta
            if delta and delta.content:
                yield {"search_status" : False, "data" : delta.content}
            elif delta and delta.function_call:
                new_message += delta.function_call.arguments
        if new_message:
            yield {"search_status" : True, "data" : new_message}

    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        """Sync implementations for retriever."""
        matching_documents = []
        message = ""
        metaData = self.getMetadata(query)
        
        for data in metaData:
            if data['data'] is None:
                break

            search_status = data['search_status']
            if search_status == True:
                metadata = json.loads(data['data'])
                practice_data = ""
                try:
                    if 'mediator areas of practice' in metadata:
                        practice_data = metadata['mediator areas of practice']
                        del metadata['mediator areas of practice']
                    elif 'mediator practice field' in metadata:
                        practice_data = metadata['mediator practice field']
                        del metadata['mediator practice field']
                    
                except:
                    practice_data = ""

                if "mediator city" in metadata:
                    try:
                        del metadata['mediator state']
                        del metadata['mediator country']
                    except:
                        pass
                if not practice_data:
                    response = openai_client.chat.completions.create(
                        model="gpt-4-1106-preview",
                        messages=[
                            {"role": "system", "content": "Generate a message asking the user about the conflict to better match them with mediators."},
                            {"role": "user", "content": query}
                        ],
                        stream=True
                    )
                    
                    for chunk in response:
                        delta = chunk.choices[0].delta
                        if delta and delta.content:
                            message += delta.content
                            yield {"message": message}
                
                elif not any(k in metadata for k in ['mediator country', 'mediator state', 'mediator city']):
                    response = openai_client.chat.completions.create(
                        model="gpt-4-1106-preview",
                        messages=[
                            {"role": "system", "content": "Generate a message asking the user for their location to find mediators in their area."},
                            {"role": "user", "content": query}
                        ],
                        stream=True
                    )
                  
                    for chunk in response:
                        delta = chunk.choices[0].delta
                        if delta and delta.content:
                            message += delta.content
                            yield {"message": message}
                else:
                    tools = [
                        {
                            "type": "function",
                            "name": "mediator_search",
                            "description": "Extract how many mediators the user wants to search.",
                            "parameters": {
                                "type": "object",
                                "properties": {
                                    "mediator": {
                                        "type": "number",
                                        "description": "The number of mediators the user wants to search for. If a user requests a list of mediators, this means that the user wants to search for at least three mediators. If your message does not include information about the number of mediators, you must provide at least three mediators. If you only have one arbitrator, you must search three mediators in the state and out of state. In example, If there is no three mediators in California search the mediators from Oregon, Arizona, or Nevada. So you must provide at least three mediators. IMPORTANT: you must provide at least three mediators.",
                                        "default": 3
                                    }
                                },
                                "required": ["mediator"]
                            }
                        }
                    ]
                    
                    response = openai_client.chat.completions.create(
                        model="gpt-4-1106-preview",
                        messages=[
                            {"role": "system", "content": "Please extract how many mediators users want to search for."},
                            {"role": "user", "content": query}
                        ],
                        functions=tools
                    )
                    try:
                        number_str = response.choices[0].message.function_call.arguments
                        mediator_num = max(int(json.loads(number_str)['mediator']), 3)
                    except:
                        mediator_num = 3

                    template = """"""
                    prompt = "You are a professional mediator information analyzer. You have to analyze the following mediators based on the user's message. You shouldn't write the mediator's information again. You shouldn't write the mediators in context as the excellent choice or ideal candidate. You have to analyze the mediators at once. Please respond with no more than 300 characters."
                    end = """Context: {context}
                        Chat history: {chat_history}
                        Human: {human_input}
                        Your Response as Chatbot:"""
                    
                    template += prompt + end

                    prompt = PromptTemplate(
                        input_variables=["chat_history", "human_input", "context"], 
                        template=template
                    )

                    pc = Pinecone(api_key=pinecone_api_key)

                    embeddings = OpenAIEmbeddings(api_key=openai_api_key)
                    
                    index = pc.Index(pinecone_index)
                        
                    results = index.query(
                        vector=embeddings.embed_query(query),
                        top_k=5,
                        filter=metadata,
                        include_metadata=True
                    )

                    new_data = []
                    for result in results['matches']:
                        data = {}
                        for metadata_key in metadata_list:      
                            data[metadata_key] = result['metadata'].get(metadata_key, "N/A")
                        # practice_data = result['metadata'].get('mediator areas of practice', "N/A")
                        if not practice_data:
                            continue
                        
                        if practice_data:
                            new_data.append(data)
                    random.shuffle(new_data)

                    if new_data:
                        if practice_data and mediator_num == 1:
                            message += f"I have located a mediator who specializes in {practice_data}. Here are their details:\n\n"
                        elif practice_data and mediator_num > 1:
                            message += f"I have located mediators who specialize in {practice_data}. Here are their details:\n\n"
                        elif not practice_data and mediator_num == 1:
                            message += f"I have located a mediator. Here are their details:\n\n"
                        elif not practice_data and mediator_num > 1:
                            message += f"I have located mediators. Here are their details:\n\n"

                    for index, new_datum in enumerate(new_data):
                        if index < mediator_num:
                            content = ""

                            for metadata_index, metadata_key in enumerate(metadata_list):
                                value = new_datum[metadata_key]
                                content += f"<b>{metadata_value[metadata_index]}</b>: {value if value else 'N/A'} \n"
                                message += f"<b>{metadata_value[metadata_index]}</b>: {value if value else 'N/A'} \n"

                            message += "\n\n"
                            new_doc = Document(page_content=content)
                            matching_documents.append(new_doc)
                        else:
                            break
                    
                    chat_openai = ChatOpenAI(model='gpt-4-1106-preview', 
                            openai_api_key=openai_api_key)
                    
                    memory = ConversationBufferMemory(memory_key="chat_history", input_key="human_input")
                    chain = load_qa_chain(chat_openai, chain_type="stuff", prompt=prompt, memory=memory)

                    output = chain.invoke(
                        {"input_documents": matching_documents, "human_input": query}, 
                        return_only_outputs=False
                    )
                    
                    message += f"Why appropriate: {output['output_text']}"

            else:
                message += data['data']

            yield {"documents": matching_documents, "message": message}

retriever = MediatorRetriever()

from langchain_openai import ChatOpenAI
from langchain.chains import create_history_aware_retriever
from langchain import hub
from langchain_core.messages import HumanMessage

rephrase_prompt = hub.pull("langchain-ai/chat-langchain-rephrase")
llm = ChatOpenAI(model="gpt-4-1106-preview", api_key=openai_api_key)

history_aware_retriever = create_history_aware_retriever(
    llm, retriever, rephrase_prompt
)

chat_history = []

def search(query, history):
    try:
        response = history_aware_retriever.stream({"input": query, "chat_history": chat_history})
        for chunks in response:
            for chunk in chunks:
                if 'documents' in chunk:
                    chat_history.extend([HumanMessage(content=query), chunk['documents']])
                yield chunk['message']
    except StopAsyncIteration:
        print("Async iteration ended unexpectedly")
    except Exception as e:
        print(f"An error occurred: {e}")
        yield f"An error occurred: {e}"

chatbot = gr.Chatbot(avatar_images=["user_.png", "bot.png"], height="100%",container=False)

js_code = """
<style>

    html body {
        background: transparent !important;
    }
    
    ::-webkit-scrollbar {
      width: 10px;
    }
    
    ::-webkit-scrollbar-track {
      box-shadow: inset 0 0 5px grey; 
      border-radius: 10px;
    }
     
    ::-webkit-scrollbar-thumb {
      background: gray; 
      border-radius: 10px;
    }
    
    ::-webkit-scrollbar-thumb:hover {
      background: #b30000; 
    }
    
    #component-0 {
        border-radius: 0px;
    }
    
    #component-4 h1 {
        color: white;
        text-align: center;
        font-size: 22px;
        padding-top: 13px;
    }
    .avatar-container img {
        border: 1px solid white;
        padding: 4px;
    }
    
    textarea {
        height: 90px !important;
        padding: 25px !important;
        font-size: 16px !important;
    }
    
    #component-6 {
        margin-top: 0%;
        border-top-left-radius: 0px;
        border-top-right-radius: 0px;
        border-bottom-left-radius: 10px;
        border-bottom-right-radius: 10px;
    }
    
    #component-9 {
        display: none;
    }
    
    #component-10 {
        display: none;
    }
    
    footer {
        display: none !important;
    }
    
    .svelte-1b6s6s {
        display: none !important;
    }
    
    #message-btn {
        position: fixed;
        right: 5px;
        bottom: 5px;
        transition: transform 0.5s ease-out; 
    }
    #custom-btn {
        padding: 15px;
        width: 70px;
        height: 70px;
        cursor: pointer;
        border-radius: 50%;
        background: #6f6f6f;
        transition: transform 0.5s ease-out;
    }
    #component-0 {
        background-color: #dcdcdc;
    }
    
    #component-2 {
        padding: 64px 0px 0px 0px;
        border-radius: 9px;
        height: 92%;
        opacity: 0;
        transform: scaleY(0); 
        transform-origin: bottom; 
        transition: transform 0.5s ease-out, opacity 0.5s ease-out;
        overflow: hidden;
        position:fixed;
        bottom:5.7rem;
        gap: unset;
        box-shadow: -14px 5px 16px 0px rgba(0,0,0,0.08),
            0px 2px 2px 0px rgba(0,0,0,0.12),
            0px 4px 4px 0px rgba(0,0,0,0.16),
            0px 8px 8px 0px rgba(0,0,0,0.2);
    }
    #component-4 {
        background-color: #6f6f6f;
        height: 85px;
        display: flex;
        justify-content: center;
        align-items: center;
        border-top-left-radius: 10px;
        border-top-right-radius: 10px;
        border-bottom-left-radius: 0px;
        border-bottom-right-radius: 0px;
    }

    .section::-webkit-scrollbar-thumb {
        background-image: linear-gradient(180deg, #d0368a 0%, #708ad4 99%);
        box-shadow: inset 2px 2px 5px 0 rgba(#fff, 0.5);
        border-radius: 100px;
    }
    
    .gradio-container{
        padding:0 !important;
        background-color:transparent;
    }
    
    gradio-app {
        background: none !important;
    }
    
</style>
<script>
    function toggleChatbox() {
        window.parent.postMessage(
            {
                type: "toggleChat",
            },
            "*"
        );
        
        const chatbox = document.getElementById('component-2');
        const buttonImg = document.getElementById('custom-btn');
        const messageBtn = document.getElementById('message-btn');
        if ((chatbox.style.transform === 'scaleY(0)') || (chatbox.style.transform === '')) {
            chatbox.style.transform = 'scaleY(1)';
            chatbox.style.opacity = '1';
            buttonImg.src = "/file=close.png";
            messageBtn.style.transform = 'rotate(360deg)';
        } else {
            chatbox.style.transform = 'scaleY(0)';
            chatbox.style.opacity = '0';
            buttonImg.src = "/file=chat.png";
            messageBtn.style.transform = 'rotate(0deg)';
        }
    }
    
    function autoScroll() {
        const chatbox = document.querySelector('.bubble-wrap');
        if (chatbox) {
            chatbox.scrollTop = chatbox.scrollHeight;
        }
    }
    
    const timer = setInterval(isElements, 1000);
    
    function isElements() {
        const messageBtn = document.getElementById('message-btn');
        if (messageBtn) {
            messageBtn.addEventListener('click', toggleChatbox);
        }
        const userInputBox = document.querySelector('textarea.scroll-hide');
        if (userInputBox) {
            userInputBox.addEventListener('keypress', autoScroll, false);
        }
        if (messageBtn && userInputBox) {
            clearInterval(timer);
        }
    }
</script>
"""

def toggle_chatbox():
    global chatbox_visible
    chatbox_visible = not chatbox_visible
    return gr.update(visible=chatbox_visible)

chatbox_visible = True

with gr.Blocks(head=js_code) as demo:
    with gr.Column(visible=chatbox_visible) as chatbox_container:
        chatbox = gr.ChatInterface(fn=search, title="Mediate Chatbot", multimodal=False, retry_btn=None, clear_btn=None, undo_btn=None, chatbot=chatbot)
    
 
    gr.HTML(f'<button id="message-btn"/><img style="display: inline;" src="/file=chat.png" id="custom-btn"/><button>')

if __name__ == "__main__":
    demo.launch(debug=True, allowed_paths=["./"])