import sys
from typing import Any
import readline


from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
import colorama

from callbacks import handlers
from config import config
from i18n import text
from utils import utils
from agent.agent import create_agent
from walrus.toolkit import WalrusToolKit
from k8s.toolkit import KubernetesToolKit

last_error = None


def setup_agent() -> Any:
    config.init()
    colorama.init()

    llm = ChatOpenAI(
        model_name="gpt-4",
        temperature=0,
        callbacks=[handlers.PrintReasoningCallbackHandler()],
    )

    text.init_system_messages(llm)

    memory = ConversationBufferMemory(memory_key="chat_history")

    enabled_toolkits = [
        toolkit.lower() for toolkit in config.APPILOT_CONFIG.toolkits
    ]

    tools = []
    if "kubernetes" in enabled_toolkits:
        kubernetes_toolkit = KubernetesToolKit(llm=llm)
        tools.extend(kubernetes_toolkit.get_tools())
    elif "walrus" in enabled_toolkits:
        walrus_toolkit = WalrusToolKit(llm=llm)
        tools.extend(walrus_toolkit.get_tools())
    else:
        print(text.get("enable_no_toolkit"))
        sys.exit(1)

    return create_agent(
        llm,
        shared_memory=memory,
        tools=tools,
        verbose=config.APPILOT_CONFIG.verbose,
    )


def run():
    appilot_agent = setup_agent()

    print(text.get("welcome"))
    user_query = None
    while True:
        user_query = input(">")
        if utils.is_inform_sent():
            continue
        elif user_query == "exit":
            break
        elif user_query == "appilot_log":
            print_last_error()
            continue
        elif user_query.startswith("#"):
            continue
        elif not user_query.strip():
            continue

        try:
            result = appilot_agent.run(user_query)
        except handlers.HumanRejectedException as he:
            utils.print_rejected_message()
            continue
        except Exception as e:
            handle_exception(e)
            continue

        utils.print_ai_response(result)


def handle_exception(e):
    global last_error
    print(text.get("response_prefix"), end="")
    print(text.get("error_occur_message"))
    last_error = e


def print_last_error():
    if last_error is None:
        print(text.get("response_prefix"), end="")
        print(text.get("no_error_message"))
    else:
        print(last_error)

langchain.memory.ConversationBufferMemory
[((442, 455), 'config.config.init', 'config.init', ([], {}), '()\n', (453, 455), False, 'from config import config\n'), ((460, 475), 'colorama.init', 'colorama.init', ([], {}), '()\n', (473, 475), False, 'import colorama\n'), ((623, 653), 'i18n.text.init_system_messages', 'text.init_system_messages', (['llm'], {}), '(llm)\n', (648, 653), False, 'from i18n import text\n'), ((668, 719), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""'}), "(memory_key='chat_history')\n", (692, 719), False, 'from langchain.memory import ConversationBufferMemory\n'), ((1211, 1307), 'agent.agent.create_agent', 'create_agent', (['llm'], {'shared_memory': 'memory', 'tools': 'tools', 'verbose': 'config.APPILOT_CONFIG.verbose'}), '(llm, shared_memory=memory, tools=tools, verbose=config.\n    APPILOT_CONFIG.verbose)\n', (1223, 1307), False, 'from agent.agent import create_agent\n'), ((908, 934), 'k8s.toolkit.KubernetesToolKit', 'KubernetesToolKit', ([], {'llm': 'llm'}), '(llm=llm)\n', (925, 934), False, 'from k8s.toolkit import KubernetesToolKit\n'), ((1400, 1419), 'i18n.text.get', 'text.get', (['"""welcome"""'], {}), "('welcome')\n", (1408, 1419), False, 'from i18n import text\n'), ((1502, 1524), 'utils.utils.is_inform_sent', 'utils.is_inform_sent', ([], {}), '()\n', (1522, 1524), False, 'from utils import utils\n'), ((2090, 2121), 'utils.utils.print_ai_response', 'utils.print_ai_response', (['result'], {}), '(result)\n', (2113, 2121), False, 'from utils import utils\n'), ((2181, 2208), 'i18n.text.get', 'text.get', (['"""response_prefix"""'], {}), "('response_prefix')\n", (2189, 2208), False, 'from i18n import text\n'), ((2228, 2259), 'i18n.text.get', 'text.get', (['"""error_occur_message"""'], {}), "('error_occur_message')\n", (2236, 2259), False, 'from i18n import text\n'), ((1052, 1074), 'walrus.toolkit.WalrusToolKit', 'WalrusToolKit', ([], {'llm': 'llm'}), '(llm=llm)\n', (1065, 1074), False, 'from walrus.toolkit import WalrusToolKit\n'), ((1187, 1198), 'sys.exit', 'sys.exit', (['(1)'], {}), '(1)\n', (1195, 1198), False, 'import sys\n'), ((2347, 2374), 'i18n.text.get', 'text.get', (['"""response_prefix"""'], {}), "('response_prefix')\n", (2355, 2374), False, 'from i18n import text\n'), ((2398, 2426), 'i18n.text.get', 'text.get', (['"""no_error_message"""'], {}), "('no_error_message')\n", (2406, 2426), False, 'from i18n import text\n'), ((569, 609), 'callbacks.handlers.PrintReasoningCallbackHandler', 'handlers.PrintReasoningCallbackHandler', ([], {}), '()\n', (607, 609), False, 'from callbacks import handlers\n'), ((1148, 1177), 'i18n.text.get', 'text.get', (['"""enable_no_toolkit"""'], {}), "('enable_no_toolkit')\n", (1156, 1177), False, 'from i18n import text\n'), ((1945, 1975), 'utils.utils.print_rejected_message', 'utils.print_rejected_message', ([], {}), '()\n', (1973, 1975), False, 'from utils import utils\n')]