import gradio as gr
from pathlib import Path
from reactagent.environment import Environment
from reactagent.agents.agent_research import ResearchAgent
from reactagent.runner import create_parser
from reactagent import llm
from reactagent.users.user import User

class SessionInfo:
    def __init__(self):
        self.coro_cache = {}
        self.parser = create_parser()

    def make_session(self, prompt, session_hash):
        id = session_hash

        llm_name='claude-3-5-sonnet-20240620'
        fastllm_name='claude-3-haiku-20240307'
        rawargs = [
            '--research-problem', prompt,
            '--log-dir', str(Path('logs', id)),
            '--work-dir', str(Path('workspaces', id)),
            '--llm-name', llm_name,
            '--edit-script-llm-name', llm_name,
            '--fast-llm-name', fastllm_name,
        ]

        args = self.parser.parse_args(rawargs)
        llm.FAST_MODEL = args.fast_llm_name
        env = Environment(args)
        agent = ResearchAgent(args, env)
        coro = agent.run(env)

        self.coro_cache[id] = coro
        return id

    def get_response(self, human_input, session_hash):
        coro_input = human_input
        if session_hash not in self.coro_cache:
            self.make_session(human_input, session_hash)
            coro_input = None

        try:
            output = self.coro_cache[session_hash].send(coro_input)
        except StopIteration:
            output = None
            del self.coro_cache[session_hash]

        return output

session_info = SessionInfo()

def info_to_message(info):
    msg = "Agent Response:\n"
    for k, v in info.items():
        if isinstance(v, dict):
            tempv = v
            v = ""
            for k2, v2 in tempv.items():
                v += f"{k2}:\n  {v2}\n"
        v = User.indent_text(v, 2)
        msg += '-' * 64
        msg += '\n'
        msg += f"{k}:\n{v}\n"

    msg += "Please provide feedback based on the history, response entries, and observation, and questions: "
    return msg

def predict(message, history, request: gr.Request):
    response = session_info.get_response(message, request.session_hash)
    if response is None:
        return "Agent is finished. Enter a new instruction."
    return info_to_message(response)

if __name__ == "__main__":
    chatbot = gr.Chatbot(render_markdown=False)
    demo = gr.ChatInterface(predict, chatbot=chatbot)
    demo.launch()