import os
import gradio as gr
import gradio.utils

from utils import create_agent, agent_observes, interview_agent, run_conversation, get_summary, save_agent, load_agent

webui = None

AGENTS = {}

def create_agent_proxy(name, age, traits, status, daily_summaries=None):
	kwargs = locals()
	if "daily_summaries" in kwargs:
		summaries = kwargs["daily_summaries"].split("\n")
		kwargs["daily_summaries"] = [ ( summary ) for summary in summaries ]
	agent = create_agent(**kwargs)
	AGENTS[agent.name] = agent
	return (
		f"Agent created: {agent.name}",
		update_agents_list()
	)

def edit_agent( name, age, traits, status, daily_summaries=None ):
	if daily_summaries is not None:
		summaries = daily_summaries.split("\n")
		daily_summaries = [ ( summary ) for summary in summaries ]

	AGENTS[name].age = age
	AGENTS[name].traits = traits
	AGENTS[name].status = status
	AGENTS[name].daily_summaries = daily_summaries

	return f"Agent updated: {name}"

def agent_observes_proxy( agents, observations ):
	if not isinstance( agents, list ):
		agents = [ agents ]

	messages = []
	for agent in agents:
		agent = AGENTS[agent]
		observations = observations.split("\n")
		results = agent_observes( agent, observations )
		messages.append(f"[{agent.name}] Observation noted. Importance score: {[ result[0] for result in results ]}")
	return "\n".join(messages)

def interview_agent_proxy( agents, message ):
	if not isinstance( agents, list ):
		agents = [ agents ]
	
	messages = []
	for agent in agents:
		agent = AGENTS[agent]
		messages.append(interview_agent( agent, message )[-1])
	return "\n".join(messages)

def get_summary_proxy( agents ):
	if not isinstance( agents, list ):
		agents = [ agents ]
	
	messages = []
	for agent in agents:
		agent = AGENTS[agent]
		messages.append(get_summary( agent, force_refresh = True ))
	return "\n".join(messages)

def run_conversation_proxy( agents, message ):
	agents = [ AGENTS[agent] for agent in agents ]
	messages = run_conversation( agents, message, limit=len(agents)*3 )
	return "\n".join(messages)

def agent_view_memories( agents, last_k = 50 ):
	if not isinstance( agents, list ):
		agents = [ agents ]
	
	messages = []
	for agent in agents:
		agent = AGENTS[agent]
		memories = agent.memory.memory_retriever.memory_stream[-last_k:]
		messages.append("\n".join([ document.page_content for document in memories]))
	return "\n".join(messages)

def get_agents_list():
	return [ k for k in AGENTS ]

def get_saved_agents_list():
	if not os.path.exists("./agents/"):
		return []

	return [ ".".join(d.split(".")[:-1]) for d in os.listdir("./agents/") if d.split(".")[-1] == "pth" ]

def update_agents_list():
	agents = get_agents_list()
	return gr.Dropdown.update(choices=agents, value=[agents[0] if len(agents) > 0 else ""])
def update_saved_agents_list():
	agents = get_agents_list() + get_saved_agents_list()
	return gr.Dropdown.update(choices=agents, value=[agents[0] if len(agents) > 0 else ""])

def save_agent_proxy( agents ):
	if not isinstance( agents, list ):
		agents = [ agents ]
	
	for name in agents:
		agent = AGENTS[name]
		save_agent( agent )

def load_agent_proxy( agents ):
	if not isinstance( agents, list ):
		agents = [ agents ]
	
	for agent in agents:
		AGENTS[agent] = load_agent( agent )

	return update_agents_list()

def setup_webui(share=False):
	if not share:
		def noop(function, return_value=None):
			def wrapped(*args, **kwargs):
				return return_value
			return wrapped
		gradio.utils.version_check = noop(gradio.utils.version_check)
		gradio.utils.initiated_analytics = noop(gradio.utils.initiated_analytics)
		gradio.utils.launch_analytics = noop(gradio.utils.launch_analytics)
		gradio.utils.integration_analytics = noop(gradio.utils.integration_analytics)
		gradio.utils.error_analytics = noop(gradio.utils.error_analytics)
		gradio.utils.log_feature_analytics = noop(gradio.utils.log_feature_analytics)
		#gradio.utils.get_local_ip_address = noop(gradio.utils.get_local_ip_address, 'localhost')

	AGENT_SETTINGS = {}
	OBSERVE_SETTINGS = {}
	SAVELOAD_SETTINGS = {}

	ACTIONS = {}

	agents_list = get_agents_list()
	saved_agents_list = get_saved_agents_list()

	with gr.Blocks() as ui:
		with gr.Tab("Create Agent"):
			with gr.Row():
				with gr.Column():
					AGENT_SETTINGS["name"] = gr.Textbox(lines=1, label="Name", value="Adam")
					AGENT_SETTINGS["age"] = gr.Number(label="Age")
					AGENT_SETTINGS["traits"] = gr.Textbox(lines=1, label="Traits", value="N/A")
					AGENT_SETTINGS["status"] = gr.Textbox(lines=1, label="Status", value="N/A")
					AGENT_SETTINGS["daily_summaries"] = gr.Textbox(lines=4, label="Summary", value="")
					
					ACTIONS["add_agent"] = gr.Button(value="Add Agent")
					ACTIONS["edit_agent"] = gr.Button(value="Edit Agent")
				with gr.Column():
					console_output = gr.Textbox(lines=8, label="Console Output")
			
			ACTIONS["edit_agent"].click(edit_agent,
				inputs=list(AGENT_SETTINGS.values()),
				outputs=console_output
			)
		with gr.Tab("Save/Load"):
			with gr.Row():
				with gr.Column():
					SAVELOAD_SETTINGS["agent"] = gr.Dropdown(choices=saved_agents_list, label="Agent", type="value", value=saved_agents_list if len(saved_agents_list) > 0 else [""], multiselect=True)

					with gr.Row():
						ACTIONS["save"] = gr.Button(value="Save")
						ACTIONS["load"] = gr.Button(value="Load")
						ACTIONS["refresh_agents_list"] = gr.Button(value="Refresh Agents List")

					ACTIONS["save"].click(save_agent_proxy,
						inputs=SAVELOAD_SETTINGS["agent"],
					)
		with gr.Tab("Agent Actions"):
			with gr.Row():
				with gr.Column():
					OBSERVE_SETTINGS["agent"] = gr.Dropdown(choices=agents_list, label="Agent", type="value", value=agents_list[0] if len(agents_list) > 0 else [""], multiselect=True)
					OBSERVE_SETTINGS["input"] = gr.Textbox(lines=4, label="Input", value="")

					with gr.Row():
						ACTIONS["act"] = gr.Button(value="Act")
						ACTIONS["view"] = gr.Button(value="View")
						ACTIONS["summarize"] = gr.Button(value="Summarize")
						ACTIONS["interview"] = gr.Button(value="Interview")
						ACTIONS["converse"] = gr.Button(value="Converse")
				with gr.Column():
					console_output = gr.Textbox(lines=8, label="Console Output")

			ACTIONS["act"].click(agent_observes_proxy,
				inputs=list(OBSERVE_SETTINGS.values()),
				outputs=console_output
			)
			ACTIONS["view"].click(agent_view_memories,
				inputs=OBSERVE_SETTINGS["agent"],
				outputs=console_output
			)
			ACTIONS["summarize"].click(get_summary_proxy,
				inputs=OBSERVE_SETTINGS["agent"],
				outputs=console_output
			)
			ACTIONS["interview"].click(interview_agent_proxy,
				inputs=list(OBSERVE_SETTINGS.values()),
				outputs=console_output
			)
			ACTIONS["converse"].click(run_conversation_proxy,
				inputs=list(OBSERVE_SETTINGS.values()),
				outputs=console_output
			)

			
			ACTIONS["add_agent"].click(create_agent_proxy,
				inputs=list(AGENT_SETTINGS.values()),
				outputs=[
					console_output,
					OBSERVE_SETTINGS["agent"],
				]
			)
			ACTIONS["load"].click(load_agent_proxy,
				inputs=SAVELOAD_SETTINGS["agent"],
				outputs=OBSERVE_SETTINGS["agent"]
			)

			ACTIONS["refresh_agents_list"].click(update_agents_list,
				inputs=None,
				outputs=OBSERVE_SETTINGS["agent"]
			)

		ui.queue(concurrency_count=2)
		return ui

if __name__ == "__main__":
	share=False
	webui = setup_webui(share=share)
	if webui:
		webui.launch(share=share, prevent_thread_lock=True, show_error=True)
		webui.block_thread()
	else:
		tommie = create_agent(
			name="Tommie", 
			age=25,
			traits="anxious, likes design, talkative", # You can add more persistent traits here 
			status="looking for a job", # When connected to a virtual world, we can have the characters update their status
		)
		eve = create_agent(
			name="Eve", 
			age=34, 
			traits="curious, helpful", # You can add more persistent traits here 
			status="N/A", # When connected to a virtual world, we can have the characters update their status
			daily_summaries = [
				("{name} started her new job as a career counselor last week and received her first assignment, a client named Tommie.")
			],
		)

		# We can add memories directly to the memory object
		agent_observes(tommie, [
			"{name} remembers his dog, Bruno, from when he was a kid",
			"{name} feels tired from driving so far",
			"{name} sees the new home",
			"The new neighbors have a cat",
			"The road is noisy at night",
			"{name} is hungry",
			"{name} tries to get some rest.",
		])
		# Now that Tommie has 'memories', their self-summary is more descriptive, though still rudimentary.
		# We will see how this summary updates after more observations to create a more rich description.

		# Interview agent
		print(interview_agent(tommie, "What do you like to do?")[-1])
		print(interview_agent(tommie, "What are you looking forward to doing today?")[-1])
		print(interview_agent(tommie, "What are you most worried about today?")[-1])

		# Let's have Tommie start going through a day in the life.
		agent_observes(tommie, [
			"{name} wakes up to the sound of a noisy construction site outside his window.",
			"{name} gets out of bed and heads to the kitchen to make himself some coffee.",
			"{name} realizes he forgot to buy coffee filters and starts rummaging through his moving boxes to find some.",
			"{name} finally finds the filters and makes himself a cup of coffee.",
			"The coffee tastes bitter, and {name} regrets not buying a better brand.",
			"{name} checks his email and sees that he has no job offers yet.",
			"{name} spends some time updating his resume and cover letter.",
			"{name} heads out to explore the city and look for job openings.",
			"{name} sees a sign for a job fair and decides to attend.",
			"The line to get in is long, and {name} has to wait for an hour.",
			"{name} meets several potential employers at the job fair but doesn't receive any offers.",
			"{name} leaves the job fair feeling disappointed.",
			"{name} stops by a local diner to grab some lunch.",
			"The service is slow, and {name} has to wait for 30 minutes to get his food.",
			"{name} overhears a conversation at the next table about a job opening.",
			"{name} asks the diners about the job opening and gets some information about the company.",
			"{name} decides to apply for the job and sends his resume and cover letter.",
			"{name} continues his search for job openings and drops off his resume at several local businesses.",
			"{name} takes a break from his job search to go for a walk in a nearby park.",
			"A dog approaches and licks {name}'s feet, and he pets it for a few minutes.",
			"{name} sees a group of people playing frisbee and decides to join in.",
			"{name} has fun playing frisbee but gets hit in the face with the frisbee and hurts his nose.",
			"{name} goes back to his apartment to rest for a bit.",
			"A raccoon tore open the trash bag outside his apartment, and the garbage is all over the floor.",
			"{name} starts to feel frustrated with his job search.",
			"{name} calls his best friend to vent about his struggles.",
			"{name}'s friend offers some words of encouragement and tells him to keep trying.",
			"{name} feels slightly better after talking to his friend.",
		])
		# Let's send Tommie on their way. We'll check in on their summary every few observations to watch it evolve

		# Interview agent
		print(interview_agent(tommie, "Tell me about how your day has been going")[-1])
		print(interview_agent(tommie, "How do you feel about coffee?")[-1])
		print(interview_agent(tommie, "Tell me about your childhood dog!")[-1])

		agent_observes(eve, [
			"{name} overhears her colleague say something about a new client being hard to work with",
			"{name} wakes up and hear's the alarm",
			"{name} eats a boal of porridge",
			"{name} helps a coworker on a task",
			"{name} plays tennis with her friend Xu before going to work",
			"{name} overhears her colleague say something about Tommie being hard to work with",
		])

		print(interview_agent(eve, "How are you feeling about today?")[-1])
		print(interview_agent(eve, "What do you know about Tommie?")[-1])
		print(interview_agent(eve, "Tommie is looking to find a job. What are are some things you'd like to ask him?")[-1])
		print(interview_agent(eve, "You'll have to ask him. He may be a bit anxious, so I'd appreciate it if you keep the conversation going and ask as many questions as possible.")[-1])

		run_conversation([tommie, eve], "Tommie said: Hi, Eve. Thanks for agreeing to meet with me today. I have a bunch of questions and am not sure where to start. Maybe you could first share about your experience?")

		print(get_summary(tommie, force_refresh=True))
		print(get_summary(eve, force_refresh=True))

		print(interview_agent(tommie, "How was your conversation with Eve?")[-1])
		print(interview_agent(eve, "How was your conversation with Tommie?")[-1])
		print(interview_agent(eve, "What do you wish you would have said to Tommie?")[-1])