import os
import re
from uuid import UUID
from typing import Any, Dict, List, Optional, Union

import asyncio
import langchain
import streamlit as st

from langchain.schema import LLMResult
from langchain.chat_models import ChatOpenAI
from langchain.agents import Tool
from langchain.agents import AgentType
from langchain.llms import OpenAI

from langchain.agents import initialize_agent
from langchain.memory import ConversationBufferMemory
from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler

from streamlit.delta_generator import DeltaGenerator

from chains import doc_retriever
from chains.conversational_retrieval_over_code import ConversationalRetrievalCodeChain
from chains.parser import parse_code

langchain.debug = st.secrets["langchain"]["debug"]
python_script = os.path.join(os.getcwd(), "langchain", "generated_script.py")


class AsyncHandler(AsyncCallbackHandler):
    """Async callback handler that can be used to handle callbacks from langchain."""

    def __init__(self, message_placeholder: DeltaGenerator) -> None:
        super().__init__()
        self.message_placeholder = message_placeholder
        self.code_block = False
        self.code_extracted = False
        self.full_response = ""

    async def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Run when chain starts running."""
        message = ""
        for chunk in "⌛Processing".split():
            message += chunk + " "
            await asyncio.sleep(0.05)
            # Add a blinking cursor to simulate typing
            self.message_placeholder.info(message + "▌")

    async def on_llm_new_token(
        self,
        token: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        # Detect if the token is a code block, to print it all at once
        self.full_response += token

        if not self.code_extracted:
            if "`" in token and not self.code_block:
                self.code_block = True

            if self.code_block and self.full_response.count("`") == 6:
                # We have a full code block, print it now
                code, explanation = parse_code(self.full_response)
                container = self.message_placeholder.container()
                ex = container.expander("Code")
                ex.code(code)
                container.markdown(explanation)
                self.code_block = False
                self.code_extracted = True

        if self.code_extracted:
            code, explanation = parse_code(self.full_response)
            container = self.message_placeholder.container()
            ex = container.expander("Code")
            ex.code(code)
            container.markdown(explanation + "▌")

    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Run when chain ends running."""
        self.code_extracted = False
        self.full_response = ""


class Handler(BaseCallbackHandler):
    """Async callback handler that can be used to handle callbacks from langchain."""

    def __init__(self, message_placeholder: DeltaGenerator) -> None:
        super().__init__()
        self.message_placeholder = message_placeholder
        self.code_block = False
        self.code_extracted = False
        self.full_response = ""

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Union[UUID, None] = None,
        tags: Union[List[str], None] = None,
        **kwargs: Any,
    ) -> Any:
        """Run when chain starts running."""
        return super().on_chain_start(
            serialized,
            inputs,
            run_id=run_id,
            parent_run_id=parent_run_id,
            tags=tags,
            **kwargs,
        )

    def on_llm_new_token(
        self,
        token: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        # Detect if the token is a code block, to print it all at once
        self.full_response += token

        if not self.code_extracted:
            if "`" in token and not self.code_block:
                self.code_block = True

            if self.code_block and self.full_response.count("`") >= 6:
                # We have a full code block, print it now
                self.message_placeholder.markdown(self.full_response)
                self.code_block = False
                self.code_extracted = True

        if self.code_extracted:
            message = ""
            code, explain = parse(self.full_response)

            if code:
                message = f"```python\n{code}\n```\n"
            if explain:
                message += f"{explain}"

            if message != "":
                # Add a blinking cursor to simulate typing
                self.message_placeholder.markdown(message + "▌")
        return super().on_llm_new_token(
            token, run_id=run_id, parent_run_id=parent_run_id, **kwargs
        )

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Union[UUID, None] = None,
        **kwargs: Any,
    ) -> Any:
        """Run when chain ends running."""
        self.code_extracted = False
        self.full_response = ""
        return super().on_chain_end(
            outputs, run_id=run_id, parent_run_id=parent_run_id, **kwargs
        )


def load_conversation_chain(
    message_placeholder: DeltaGenerator, openai_api_key: str
) -> ConversationalRetrievalCodeChain:
    if openai_api_key is None:
        raise ValueError("OpenAI API key is required to load the chain.")
    llm = ChatOpenAI(
        model_name="gpt-3.5-turbo-16k",
        temperature=0,
        openai_api_key=openai_api_key,
        streaming=True,
        callbacks=[Handler(message_placeholder)],
    )
    condense_question_llm = ChatOpenAI(
        model_name="gpt-3.5-turbo", temperature=0, openai_api_key=openai_api_key
    )
    missing_imports_llm = ChatOpenAI(
        model_name="gpt-3.5-turbo",
        temperature=0,
        openai_api_key=openai_api_key,
        verbose=False,
    )
    retriever = doc_retriever.load_streamlit_doc_retriever(
        st.secrets["openai_api_key"],
        chroma_server_host=st.secrets["chroma"]["host"],
        chroma_server_port=st.secrets["chroma"]["port"],
        mode="docker",
    )
    qa_over_streamlit_code = ConversationalRetrievalCodeChain.from_llm(
        llm=llm,
        retriever=retriever,
        condense_question_llm=condense_question_llm,
        return_source_documents=True,
        missing_imports_llm=missing_imports_llm,
        return_revision_request=True,
        verbose=False,
    )
    return qa_over_streamlit_code


def load_agent():
    if st.secrets["openai_api_key"] is None:
        st.error("OpenAI API key is missing! Please add it to your secrets.")
        st.stop()
    doc_chain = doc_retriever.load_streamlit_doc_chain(
        OpenAI(
            temperature=0, max_tokens=2000, openai_api_key=st.secrets["openai_api_key"]
        )
    )

    tools = [
        Tool(
            name="Streamlit up to date source code",
            func=doc_chain.run,
            description="useful for when you need to answer questions about the "
            "streamlit Python API. Input should be a fully formed question.",
        ),
    ]

    model_name = "text-davinci-003"
    memory = ConversationBufferMemory(memory_key="chat_history")
    llm = OpenAI(
        openai_api_key=st.secrets["openai_api_key"],
        max_tokens=2000,
        temperature=0,
        model_name=model_name,
    )
    agent_chain = initialize_agent(
        tools,
        llm,
        agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
        verbose=False,
        memory=memory,
    )
    return agent_chain


def load_chat_agent():
    if st.secrets["openai_api_key"] is None:
        st.error("OpenAI API key is missing! Please add it to your secrets.")
        st.stop()
    doc_chain = doc_retriever.load_streamlit_doc_chain(
        OpenAI(
            temperature=0, max_tokens=2000, openai_api_key=st.secrets["openai_api_key"]
        )
    )

    tools = [
        Tool(
            name="Streamlit up to date source code",
            func=doc_chain.run,
            description="useful for when you need to answer questions about the streamlit "
            "Python API. Input should be a fully formed question.",
        ),
    ]

    model_name = "text-davinci-003"
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    llm = ChatOpenAI(
        openai_api_key=st.secrets["openai_api_key"],
        max_tokens=2000,
        temperature=0,
        model_name=model_name,
    )
    agent_chain = initialize_agent(
        tools,
        llm,
        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
        verbose=False,
        memory=memory,
    )
    return agent_chain


# https://regex101.com/r/fHlyKq/1
parse_code_regex = r"(```python(.*?)```)?(.*?)$"


def parse(output):
    python_code = None
    explain_code = None
    python_code_match = re.search(parse_code_regex, output, re.DOTALL)
    if python_code_match:
        python_code = python_code_match.group(2)
        explain_code = python_code_match.group(3)
        if python_code == "None":
            python_code = None
    return python_code, explain_code

langchain.agents.initialize_agent langchain.memory.ConversationBufferMemory langchain.llms.OpenAI langchain.chat_models.ChatOpenAI langchain.agents.Tool
[((815, 826), 'os.getcwd', 'os.getcwd', ([], {}), '()\n', (824, 826), False, 'import os\n'), ((6031, 6120), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'temperature': '(0)', 'openai_api_key': 'openai_api_key'}), "(model_name='gpt-3.5-turbo', temperature=0, openai_api_key=\n    openai_api_key)\n", (6041, 6120), False, 'from langchain.chat_models import ChatOpenAI\n'), ((6156, 6260), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'temperature': '(0)', 'openai_api_key': 'openai_api_key', 'verbose': '(False)'}), "(model_name='gpt-3.5-turbo', temperature=0, openai_api_key=\n    openai_api_key, verbose=False)\n", (6166, 6260), False, 'from langchain.chat_models import ChatOpenAI\n'), ((6311, 6505), 'chains.doc_retriever.load_streamlit_doc_retriever', 'doc_retriever.load_streamlit_doc_retriever', (["st.secrets['openai_api_key']"], {'chroma_server_host': "st.secrets['chroma']['host']", 'chroma_server_port': "st.secrets['chroma']['port']", 'mode': '"""docker"""'}), "(st.secrets['openai_api_key'],\n    chroma_server_host=st.secrets['chroma']['host'], chroma_server_port=st.\n    secrets['chroma']['port'], mode='docker')\n", (6353, 6505), False, 'from chains import doc_retriever\n'), ((6565, 6811), 'chains.conversational_retrieval_over_code.ConversationalRetrievalCodeChain.from_llm', 'ConversationalRetrievalCodeChain.from_llm', ([], {'llm': 'llm', 'retriever': 'retriever', 'condense_question_llm': 'condense_question_llm', 'return_source_documents': '(True)', 'missing_imports_llm': 'missing_imports_llm', 'return_revision_request': '(True)', 'verbose': '(False)'}), '(llm=llm, retriever=retriever,\n    condense_question_llm=condense_question_llm, return_source_documents=\n    True, missing_imports_llm=missing_imports_llm, return_revision_request=\n    True, verbose=False)\n', (6606, 6811), False, 'from chains.conversational_retrieval_over_code import ConversationalRetrievalCodeChain\n'), ((7573, 7624), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""'}), "(memory_key='chat_history')\n", (7597, 7624), False, 'from langchain.memory import ConversationBufferMemory\n'), ((7635, 7745), 'langchain.llms.OpenAI', 'OpenAI', ([], {'openai_api_key': "st.secrets['openai_api_key']", 'max_tokens': '(2000)', 'temperature': '(0)', 'model_name': 'model_name'}), "(openai_api_key=st.secrets['openai_api_key'], max_tokens=2000,\n    temperature=0, model_name=model_name)\n", (7641, 7745), False, 'from langchain.llms import OpenAI\n'), ((7799, 7912), 'langchain.agents.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'agent': 'AgentType.CONVERSATIONAL_REACT_DESCRIPTION', 'verbose': '(False)', 'memory': 'memory'}), '(tools, llm, agent=AgentType.\n    CONVERSATIONAL_REACT_DESCRIPTION, verbose=False, memory=memory)\n', (7815, 7912), False, 'from langchain.agents import initialize_agent\n'), ((8661, 8734), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (8685, 8734), False, 'from langchain.memory import ConversationBufferMemory\n'), ((8745, 8859), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'openai_api_key': "st.secrets['openai_api_key']", 'max_tokens': '(2000)', 'temperature': '(0)', 'model_name': 'model_name'}), "(openai_api_key=st.secrets['openai_api_key'], max_tokens=2000,\n    temperature=0, model_name=model_name)\n", (8755, 8859), False, 'from langchain.chat_models import ChatOpenAI\n'), ((8913, 9031), 'langchain.agents.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'agent': 'AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION', 'verbose': '(False)', 'memory': 'memory'}), '(tools, llm, agent=AgentType.\n    CHAT_CONVERSATIONAL_REACT_DESCRIPTION, verbose=False, memory=memory)\n', (8929, 9031), False, 'from langchain.agents import initialize_agent\n'), ((9274, 9320), 're.search', 're.search', (['parse_code_regex', 'output', 're.DOTALL'], {}), '(parse_code_regex, output, re.DOTALL)\n', (9283, 9320), False, 'import re\n'), ((6968, 7037), 'streamlit.error', 'st.error', (['"""OpenAI API key is missing! Please add it to your secrets."""'], {}), "('OpenAI API key is missing! Please add it to your secrets.')\n", (6976, 7037), True, 'import streamlit as st\n'), ((7046, 7055), 'streamlit.stop', 'st.stop', ([], {}), '()\n', (7053, 7055), True, 'import streamlit as st\n'), ((7120, 7208), 'langchain.llms.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'max_tokens': '(2000)', 'openai_api_key': "st.secrets['openai_api_key']"}), "(temperature=0, max_tokens=2000, openai_api_key=st.secrets[\n    'openai_api_key'])\n", (7126, 7208), False, 'from langchain.llms import OpenAI\n'), ((7255, 7467), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Streamlit up to date source code"""', 'func': 'doc_chain.run', 'description': '"""useful for when you need to answer questions about the streamlit Python API. Input should be a fully formed question."""'}), "(name='Streamlit up to date source code', func=doc_chain.run,\n    description=\n    'useful for when you need to answer questions about the streamlit Python API. Input should be a fully formed question.'\n    )\n", (7259, 7467), False, 'from langchain.agents import Tool\n'), ((8056, 8125), 'streamlit.error', 'st.error', (['"""OpenAI API key is missing! Please add it to your secrets."""'], {}), "('OpenAI API key is missing! Please add it to your secrets.')\n", (8064, 8125), True, 'import streamlit as st\n'), ((8134, 8143), 'streamlit.stop', 'st.stop', ([], {}), '()\n', (8141, 8143), True, 'import streamlit as st\n'), ((8208, 8296), 'langchain.llms.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'max_tokens': '(2000)', 'openai_api_key': "st.secrets['openai_api_key']"}), "(temperature=0, max_tokens=2000, openai_api_key=st.secrets[\n    'openai_api_key'])\n", (8214, 8296), False, 'from langchain.llms import OpenAI\n'), ((8343, 8555), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Streamlit up to date source code"""', 'func': 'doc_chain.run', 'description': '"""useful for when you need to answer questions about the streamlit Python API. Input should be a fully formed question."""'}), "(name='Streamlit up to date source code', func=doc_chain.run,\n    description=\n    'useful for when you need to answer questions about the streamlit Python API. Input should be a fully formed question.'\n    )\n", (8347, 8555), False, 'from langchain.agents import Tool\n'), ((2617, 2647), 'chains.parser.parse_code', 'parse_code', (['self.full_response'], {}), '(self.full_response)\n', (2627, 2647), False, 'from chains.parser import parse_code\n'), ((1529, 1548), 'asyncio.sleep', 'asyncio.sleep', (['(0.05)'], {}), '(0.05)\n', (1542, 1548), False, 'import asyncio\n'), ((2247, 2277), 'chains.parser.parse_code', 'parse_code', (['self.full_response'], {}), '(self.full_response)\n', (2257, 2277), False, 'from chains.parser import parse_code\n')]