import streamlit as st
import dotenv
import langchain
import json

from cassandra.cluster import Session
from cassandra.query import PreparedStatement

from langchain.agents.agent_toolkits import create_retriever_tool, create_conversational_retrieval_agent
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import OpenAIEmbeddings
from langchain.callbacks import StreamlitCallbackHandler
from langchain.schema import BaseRetriever, Document, SystemMessage

from cassandra.cluster import Cluster, Session
from cassandra.auth import PlainTextAuthProvider

# Enable langchain debug mode
langchain.debug = True

dotenv.load_dotenv(dotenv.find_dotenv())


class AstraProductRetriever(BaseRetriever):
    session: Session
    embedding: OpenAIEmbeddings
    lang: str = "English"
    search_statement_en: PreparedStatement = None
    search_statement_th: PreparedStatement = None

    class Config:
        arbitrary_types_allowed = True

    def get_relevant_documents(self, query):
        docs = []
        embeddingvector = self.embedding.embed_query(query)
        if self.lang == "Thai":
            if self.search_statement_th is None:
                self.search_statement_th = self.session.prepare("""
                    SELECT
                        product_id,
                        brand,
                        saleprice,
                        product_categories,
                        product_name,
                        short_description,
                        long_description
                    FROM hybridretail.products_cg_hybrid
                    ORDER BY openai_description_embedding_th ANN OF ?
                    LIMIT ?
                    """)
            query = self.search_statement_th
        else:
            if self.search_statement_en is None:
                self.search_statement_en = self.session.prepare("""
                    SELECT
                        product_id,
                        brand,
                        saleprice,
                        product_categories,
                        product_name_en,
                        short_description_en,
                        long_description_en
                    FROM hybridretail.products_cg_hybrid
                    ORDER BY openai_description_embedding_en ANN OF ?
                    LIMIT ?
                    """)
            query = self.search_statement_en
        results = self.session.execute(query, [embeddingvector, 5])
        top_products = results._current_rows
        for r in top_products:
            docs.append(Document(
                id=r.product_id,
                page_content=r.product_name if self.lang == "Thai" else r.product_name_en,
                metadata={"product id": r.product_id,
                          "brand": r.brand,
                          "product category": r.product_categories,
                          "product name": r.product_name if self.lang == "Thai" else r.product_name_en,
                          "description": r.short_description if self.lang == "Thai" else r.short_description_en,
                          "price": r.saleprice
                          }
            ))

        return docs


def get_session(scb: str, secrets: str) -> Session:
    """
    Connect to Astra DB using secure connect bundle and credentials.

    Parameters
    ----------
    scb : str
        Path to secure connect bundle.
    secrets : str
        Path to credentials.
    """

    cloud_config = {
        'secure_connect_bundle': scb
    }

    with open(secrets) as f:
        secrets = json.load(f)

    CLIENT_ID = secrets["clientId"]
    CLIENT_SECRET = secrets["secret"]

    auth_provider = PlainTextAuthProvider(CLIENT_ID, CLIENT_SECRET)
    cluster = Cluster(cloud=cloud_config, auth_provider=auth_provider)
    return cluster.connect()


@st.cache_resource
def create_chatbot(lang: str):
    print(f"Creating chatbot for {lang}...")
    session = get_session(scb='./config/secure-connect-multilingual.zip',
                          secrets='./config/multilingual-token.json')
    llm = ChatOpenAI(temperature=0, streaming=True)
    embedding = OpenAIEmbeddings()
    retriever = AstraProductRetriever(
        session=session, embedding=embedding, lang=lang)
    retriever_tool = create_retriever_tool(
        retriever, "products_retrevier", "Useful when searching for products from a product description. Prices are in THB.")
    system_message = "You are a customer service of a home improvement store and you are asked to pick products for a customer."
    if lang == "Thai":
        system_message = f"{system_message} All the responses should be in Thai language."
    message = SystemMessage(content=system_message)
    agent_executor = create_conversational_retrieval_agent(
        llm=llm, tools=[retriever_tool], system_message=message, verbose=True)
    return agent_executor


if 'history' not in st.session_state:
    st.session_state['history'] = {
        "English": [],
        "Thai": []
    }

st.set_page_config(layout="wide")

with st.sidebar:
    lang = st.radio(
        "Chat language",
        ["English", "Thai"],
        captions=[".", "Experimental"])

chatbot = create_chatbot(lang)


# Display chat messages from history on app rerun
for (query, answer) in st.session_state['history'][lang]:
    with st.chat_message("User"):
        st.markdown(query)
    with st.chat_message("Bot"):
        st.markdown(answer)

prompt = st.chat_input(placeholder="Ask chatbot")
if prompt:
    # Display user message in chat message container
    with st.chat_message("User"):
        st.markdown(prompt)
    # Display assistant response in chat message container
    with st.chat_message("Bot"):
        st_callback = StreamlitCallbackHandler(st.container())
        result = result = chatbot.invoke({
            "input": prompt,
            "chat_history": st.session_state['history'][lang]
        }, config={"callbacks": [st_callback]})
        st.session_state['history'][lang].append((prompt, result["output"]))
        st.markdown(result["output"])
langchain.chat_models.ChatOpenAI langchain.schema.Document langchain.agents.agent_toolkits.create_conversational_retrieval_agent langchain.agents.agent_toolkits.create_retriever_tool langchain.schema.SystemMessage langchain.embeddings.OpenAIEmbeddings
[((5021, 5054), 'streamlit.set_page_config', 'st.set_page_config', ([], {'layout': '"""wide"""'}), "(layout='wide')\n", (5039, 5054), True, 'import streamlit as st\n'), ((5462, 5502), 'streamlit.chat_input', 'st.chat_input', ([], {'placeholder': '"""Ask chatbot"""'}), "(placeholder='Ask chatbot')\n", (5475, 5502), True, 'import streamlit as st\n'), ((648, 668), 'dotenv.find_dotenv', 'dotenv.find_dotenv', ([], {}), '()\n', (666, 668), False, 'import dotenv\n'), ((3694, 3741), 'cassandra.auth.PlainTextAuthProvider', 'PlainTextAuthProvider', (['CLIENT_ID', 'CLIENT_SECRET'], {}), '(CLIENT_ID, CLIENT_SECRET)\n', (3715, 3741), False, 'from cassandra.auth import PlainTextAuthProvider\n'), ((3756, 3812), 'cassandra.cluster.Cluster', 'Cluster', ([], {'cloud': 'cloud_config', 'auth_provider': 'auth_provider'}), '(cloud=cloud_config, auth_provider=auth_provider)\n', (3763, 3812), False, 'from cassandra.cluster import Cluster, Session\n'), ((4093, 4134), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0)', 'streaming': '(True)'}), '(temperature=0, streaming=True)\n', (4103, 4134), False, 'from langchain.chat_models import ChatOpenAI\n'), ((4151, 4169), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (4167, 4169), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((4287, 4435), 'langchain.agents.agent_toolkits.create_retriever_tool', 'create_retriever_tool', (['retriever', '"""products_retrevier"""', '"""Useful when searching for products from a product description. Prices are in THB."""'], {}), "(retriever, 'products_retrevier',\n    'Useful when searching for products from a product description. Prices are in THB.'\n    )\n", (4308, 4435), False, 'from langchain.agents.agent_toolkits import create_retriever_tool, create_conversational_retrieval_agent\n'), ((4693, 4730), 'langchain.schema.SystemMessage', 'SystemMessage', ([], {'content': 'system_message'}), '(content=system_message)\n', (4706, 4730), False, 'from langchain.schema import BaseRetriever, Document, SystemMessage\n'), ((4752, 4864), 'langchain.agents.agent_toolkits.create_conversational_retrieval_agent', 'create_conversational_retrieval_agent', ([], {'llm': 'llm', 'tools': '[retriever_tool]', 'system_message': 'message', 'verbose': '(True)'}), '(llm=llm, tools=[retriever_tool],\n    system_message=message, verbose=True)\n', (4789, 4864), False, 'from langchain.agents.agent_toolkits import create_retriever_tool, create_conversational_retrieval_agent\n'), ((5084, 5162), 'streamlit.radio', 'st.radio', (['"""Chat language"""', "['English', 'Thai']"], {'captions': "['.', 'Experimental']"}), "('Chat language', ['English', 'Thai'], captions=['.', 'Experimental'])\n", (5092, 5162), True, 'import streamlit as st\n'), ((3585, 3597), 'json.load', 'json.load', (['f'], {}), '(f)\n', (3594, 3597), False, 'import json\n'), ((5339, 5362), 'streamlit.chat_message', 'st.chat_message', (['"""User"""'], {}), "('User')\n", (5354, 5362), True, 'import streamlit as st\n'), ((5372, 5390), 'streamlit.markdown', 'st.markdown', (['query'], {}), '(query)\n', (5383, 5390), True, 'import streamlit as st\n'), ((5400, 5422), 'streamlit.chat_message', 'st.chat_message', (['"""Bot"""'], {}), "('Bot')\n", (5415, 5422), True, 'import streamlit as st\n'), ((5432, 5451), 'streamlit.markdown', 'st.markdown', (['answer'], {}), '(answer)\n', (5443, 5451), True, 'import streamlit as st\n'), ((5576, 5599), 'streamlit.chat_message', 'st.chat_message', (['"""User"""'], {}), "('User')\n", (5591, 5599), True, 'import streamlit as st\n'), ((5609, 5628), 'streamlit.markdown', 'st.markdown', (['prompt'], {}), '(prompt)\n', (5620, 5628), True, 'import streamlit as st\n'), ((5697, 5719), 'streamlit.chat_message', 'st.chat_message', (['"""Bot"""'], {}), "('Bot')\n", (5712, 5719), True, 'import streamlit as st\n'), ((6051, 6080), 'streamlit.markdown', 'st.markdown', (["result['output']"], {}), "(result['output'])\n", (6062, 6080), True, 'import streamlit as st\n'), ((5768, 5782), 'streamlit.container', 'st.container', ([], {}), '()\n', (5780, 5782), True, 'import streamlit as st\n'), ((2574, 2981), 'langchain.schema.Document', 'Document', ([], {'id': 'r.product_id', 'page_content': "(r.product_name if self.lang == 'Thai' else r.product_name_en)", 'metadata': "{'product id': r.product_id, 'brand': r.brand, 'product category': r.\n    product_categories, 'product name': r.product_name if self.lang ==\n    'Thai' else r.product_name_en, 'description': r.short_description if \n    self.lang == 'Thai' else r.short_description_en, 'price': r.saleprice}"}), "(id=r.product_id, page_content=r.product_name if self.lang ==\n    'Thai' else r.product_name_en, metadata={'product id': r.product_id,\n    'brand': r.brand, 'product category': r.product_categories,\n    'product name': r.product_name if self.lang == 'Thai' else r.\n    product_name_en, 'description': r.short_description if self.lang ==\n    'Thai' else r.short_description_en, 'price': r.saleprice})\n", (2582, 2981), False, 'from langchain.schema import BaseRetriever, Document, SystemMessage\n')]