from templates.common.suffix import suffix
from templates.common.format_instructions import format_instructions
from templates.common.docs_system_instructions import docs_system_instructions
from langchain.schema import (
    # AIMessage,
    HumanMessage,
    SystemMessage
)
from langchain.tools.json.tool import JsonSpec
from langchain.agents.agent_toolkits.json.toolkit import JsonToolkit
from langchain.chat_models import ChatOpenAI, AzureChatOpenAI
from langchain.llms.openai import OpenAI
from langchain.agents import create_json_agent, ZeroShotAgent, AgentExecutor
from langchain.chains import LLMChain
from config.config import config
import openai  # required
from dotenv import load_dotenv
load_dotenv()


class OpenAPIExplorerTool:
    @staticmethod
    def create_tools(docs):
        json_spec = JsonSpec(dict_=docs)
        json_toolkit = JsonToolkit(spec=json_spec)
        tools = json_toolkit.get_tools()
        return tools


class PipedreamOpenAPIAgent:
    def __init__(self, docs, templates, auth_example, parsed_common_files):
        system_instructions = format_template(
            f"{templates.system_instructions(auth_example, parsed_common_files)}\n{docs_system_instructions}")

        tools = OpenAPIExplorerTool.create_tools(docs)
        tool_names = [tool.name for tool in tools]

        prompt_template = ZeroShotAgent.create_prompt(
            tools=tools,
            prefix=system_instructions,
            suffix=suffix,
            format_instructions=format_instructions,
            input_variables=['input', 'agent_scratchpad']
        )

        llm_chain = LLMChain(llm=get_llm(), prompt=prompt_template)
        agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names)
        verbose = True if config['logging']['level'] == 'DEBUG' else False

        self.agent_executor = AgentExecutor.from_agent_and_tools(
            agent=agent, tools=tools, verbose=verbose)

    def run(self, input):
        try:
            result = self.agent_executor.run(input)
        except Exception as e:
            result = str(e)
            if "I don't know" in result:
                return "I don't know"
            if '```' not in result:
                raise e

        return format_result(result)


def format_template(text):
    return text.replace("{", "{{").replace("}", "}}")  # escape curly braces


def format_result(result):
    if '```' in result:
        if '```javascript' in result:
            result = result.split('```javascript')[1].split('```')[0].strip()
        else:
            result = result.split('```')[1].split('```')[0].strip()
    return result


def create_user_prompt(prompt, urls_content):
    if len(urls_content) == 0:
        return prompt + "\n\n"

    user_prompt = f"{prompt}\n\n## API docs\n\n"
    for item in urls_content:
        user_prompt += f"\n\n### {item['url']}\n\n{item['content']}"
    return user_prompt + "\n\n"


def get_llm():
    if config['openai_api_type'] == "azure":
        azure_config = config["azure"]
        return AzureChatOpenAI(deployment_name=azure_config['deployment_name'],
                               model_name=azure_config["model"], temperature=config["temperature"], request_timeout=300)
    else:
        openai_config = config["openai"]
        print(f"Using OpenAI API: {openai_config['model']}")
        return ChatOpenAI(
            model_name=openai_config["model"], temperature=config["temperature"])


def ask_agent(prompt, docs, templates, auth_example, parsed_common_files, urls_content):
    agent = PipedreamOpenAPIAgent(
        docs, templates, auth_example, parsed_common_files)
    user_prompt = create_user_prompt(prompt, urls_content)
    result = agent.run(user_prompt)
    return result


def no_docs(prompt, templates, auth_example, parsed_common_files, urls_content, normal_order=True):
    user_prompt = create_user_prompt(prompt, urls_content)
    pd_instructions = format_template(
        templates.system_instructions(auth_example, parsed_common_files))

    result = get_llm()(messages=[
        SystemMessage(content="You are the most intelligent software engineer in the world. You carefully provide accurate, factual, thoughtful, nuanced code, and are brilliant at reasoning. Follow all of the instructions below — they are all incredibly important. This code will be shipped directly to production, so it's important that it's accurate and complete."),
        HumanMessage(content=user_prompt +
                     pd_instructions if normal_order else pd_instructions+user_prompt),
    ])

    return format_result(result.content)

langchain.agents.AgentExecutor.from_agent_and_tools langchain.agents.agent_toolkits.json.toolkit.JsonToolkit langchain.agents.ZeroShotAgent.create_prompt langchain.agents.ZeroShotAgent langchain.chat_models.ChatOpenAI langchain.schema.HumanMessage langchain.schema.SystemMessage langchain.chat_models.AzureChatOpenAI langchain.tools.json.tool.JsonSpec
[((701, 714), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (712, 714), False, 'from dotenv import load_dotenv\n'), ((810, 830), 'langchain.tools.json.tool.JsonSpec', 'JsonSpec', ([], {'dict_': 'docs'}), '(dict_=docs)\n', (818, 830), False, 'from langchain.tools.json.tool import JsonSpec\n'), ((854, 881), 'langchain.agents.agent_toolkits.json.toolkit.JsonToolkit', 'JsonToolkit', ([], {'spec': 'json_spec'}), '(spec=json_spec)\n', (865, 881), False, 'from langchain.agents.agent_toolkits.json.toolkit import JsonToolkit\n'), ((1343, 1524), 'langchain.agents.ZeroShotAgent.create_prompt', 'ZeroShotAgent.create_prompt', ([], {'tools': 'tools', 'prefix': 'system_instructions', 'suffix': 'suffix', 'format_instructions': 'format_instructions', 'input_variables': "['input', 'agent_scratchpad']"}), "(tools=tools, prefix=system_instructions, suffix\n    =suffix, format_instructions=format_instructions, input_variables=[\n    'input', 'agent_scratchpad'])\n", (1370, 1524), False, 'from langchain.agents import create_json_agent, ZeroShotAgent, AgentExecutor\n'), ((1670, 1730), 'langchain.agents.ZeroShotAgent', 'ZeroShotAgent', ([], {'llm_chain': 'llm_chain', 'allowed_tools': 'tool_names'}), '(llm_chain=llm_chain, allowed_tools=tool_names)\n', (1683, 1730), False, 'from langchain.agents import create_json_agent, ZeroShotAgent, AgentExecutor\n'), ((1837, 1914), 'langchain.agents.AgentExecutor.from_agent_and_tools', 'AgentExecutor.from_agent_and_tools', ([], {'agent': 'agent', 'tools': 'tools', 'verbose': 'verbose'}), '(agent=agent, tools=tools, verbose=verbose)\n', (1871, 1914), False, 'from langchain.agents import create_json_agent, ZeroShotAgent, AgentExecutor\n'), ((3038, 3201), 'langchain.chat_models.AzureChatOpenAI', 'AzureChatOpenAI', ([], {'deployment_name': "azure_config['deployment_name']", 'model_name': "azure_config['model']", 'temperature': "config['temperature']", 'request_timeout': '(300)'}), "(deployment_name=azure_config['deployment_name'], model_name\n    =azure_config['model'], temperature=config['temperature'],\n    request_timeout=300)\n", (3053, 3201), False, 'from langchain.chat_models import ChatOpenAI, AzureChatOpenAI\n'), ((3351, 3436), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': "openai_config['model']", 'temperature': "config['temperature']"}), "(model_name=openai_config['model'], temperature=config['temperature']\n    )\n", (3361, 3436), False, 'from langchain.chat_models import ChatOpenAI, AzureChatOpenAI\n'), ((4061, 4430), 'langchain.schema.SystemMessage', 'SystemMessage', ([], {'content': '"""You are the most intelligent software engineer in the world. You carefully provide accurate, factual, thoughtful, nuanced code, and are brilliant at reasoning. Follow all of the instructions below — they are all incredibly important. This code will be shipped directly to production, so it\'s important that it\'s accurate and complete."""'}), '(content=\n    "You are the most intelligent software engineer in the world. You carefully provide accurate, factual, thoughtful, nuanced code, and are brilliant at reasoning. Follow all of the instructions below — they are all incredibly important. This code will be shipped directly to production, so it\'s important that it\'s accurate and complete."\n    )\n', (4074, 4430), False, 'from langchain.schema import HumanMessage, SystemMessage\n'), ((4430, 4537), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': '(user_prompt + pd_instructions if normal_order else pd_instructions +\n    user_prompt)'}), '(content=user_prompt + pd_instructions if normal_order else \n    pd_instructions + user_prompt)\n', (4442, 4537), False, 'from langchain.schema import HumanMessage, SystemMessage\n')]