#!/usr/bin/env python3
# Import the necessary modules
import openai
import os
import json5 as json
import getpass
import sys
import time
import re
import argparse
import prompt_toolkit as pt
import tempfile
from .changelog import OwegaChangelog
from .license import OwegaLicense
from .config import baseConf, get_conf, list_models
from .OwegaFun import existingFunctions, connectLTS, functionlist_to_toollist
from .OweHandlers import handlers, handle_help, handler_helps
from .utils import (
	get_home_dir,
	get_temp_file,
	info_print,
	do_quit,
	success_msg,
	clrtxt,
	print_help,
	estimated_tokens,
	play_opus,
)
from .conversation import Conversation, Conversation_from
from .ask import ask
from .OwegaSession import OwegaSession as ps


# get a copy of owega's config
def get_oc_conf():
	return baseConf.copy()


# main interaction loop
def user_interaction_loop(temp_file="", input_file="", temp_is_temp=False):
	if not temp_file:
		temp_is_temp = True
		temp_file = get_temp_file()

	# creates Conversation object and populate it
	messages = Conversation()
	connectLTS(
		messages.add_memory,
		messages.remove_memory,
		messages.edit_memory
	)
	if input_file:
		messages.load(input_file)

	# sets the input prompt
	input_prompt = '\n  ' + clrtxt("yellow", " USER ") + ": "

	# bootup info
	info_print("===== Owega =====")
	info_print(f"Owega v{OwegaChangelog.version}")
	info_print('Type "/help" for help')
	info_print(f"Default model is {baseConf.get('model', '')}")
	info_print(f"temp file is {temp_file}")

	# API key detection
	if baseConf.get("api_key", "").startswith("sk-"):
		openai.api_key = baseConf.get("api_key", "")
	else:
		# if key not in config: ask for key only if not already set (ie envvar)
		try:
			if not openai.api_key.startswith("sk-"):
				openai.api_key = getpass.getpass(prompt="OpenAI API Key: ")
		except AttributeError:
			openai.api_key = getpass.getpass(prompt="OpenAI API Key: ")
		baseConf["api_key"] = openai.api_key

	# Organization detection
	if baseConf.get("organization", "").startswith("org-"):
		openai.organization = baseConf.get("organization", "")

	# main interaction loop:
	while True:
		# save temp file
		messages.save(temp_file)

		# get user input, and strip it (no excess spaces / tabs / newlines
		user_input = ps['main'].prompt(pt.ANSI(input_prompt)).strip()

		command_found = False
		if user_input.startswith('/'):
			uinp_spl = user_input.split(' ')
			given = ' '.join(uinp_spl[1:])
			command = uinp_spl[0][1:]
			if command in handlers.keys():
				command_found = True
				current_handler = handlers.get(command, handle_help)
				messages = current_handler(
					temp_file,
					messages,
					given,
					temp_is_temp
				)
		if not command_found:
			if baseConf.get("estimation", False):
				etkn = estimated_tokens(
					user_input,
					messages,
					functionlist_to_toollist(existingFunctions.getEnabled())
				)
				cost_per_token = (
					0.03
					if 'gpt-4' in baseConf.get("model", "")
					else 0.003
				) / 1000
				cost = cost_per_token * etkn
				print(f"\033[37mestimated tokens: {etkn}\033[0m")
				print(f"\033[37mestimated cost: {cost:.5f}\033[0m")
			if baseConf.get("debug", False):
				pre_time = time.time()
			messages = ask(
				prompt=user_input,
				messages=messages,
				model=baseConf.get("model", ''),
				temperature=baseConf.get("temperature", 0.8),
				max_tokens=baseConf.get("max_tokens", 3000),
				top_p=baseConf.get("top_p", 1.0),
				frequency_penalty=baseConf.get("frequency_penalty", 0.0),
				presence_penalty=baseConf.get("presence_penalty", 0.0)
			)
			if baseConf.get("debug", False):
				post_time = time.time()
				print(f"\033[37mrequest took {post_time-pre_time:.3f}s\033[0m")

			# Print the generated response
			print()
			print(' ' + clrtxt("magenta", " Owega ") + ": ")
			print()
			print(messages.last_answer())

			if baseConf.get('tts_enabled', False):
				tmpfile = tempfile.NamedTemporaryFile(
					prefix="owegatts.",
					suffix=".opus",
					delete=False
				)
				tmpfile.close()
				tts_answer = openai.audio.speech.create(
					model='tts-1',
					voice='nova',
					input=messages.last_answer()
				)
				if baseConf.get("debug", False):
					posttts_time = time.time()
					print("\033[37mrequest took "
						+ f"{posttts_time-pre_time:.3f}s\033[0m")
				tts_answer.stream_to_file(tmpfile.name)
				play_opus(tmpfile.name)
				os.remove(tmpfile.name)


# parse command-line arguments
def parse_args():
	parser = argparse.ArgumentParser(description="Owega main application")
	parser.add_argument("-d", "--debug", action='store_true',
		help="Enable debug output")
	parser.add_argument("-c", "--changelog", action='store_true',
		help="Display changelog and exit")
	parser.add_argument("-l", "--license", action='store_true',
		help="Display license and exit")
	parser.add_argument("-v", "--version", action='store_true',
		help="Display version and exit")
	parser.add_argument("-f", "--config-file", type=str,
		help="Specify path to config file")

	parser.add_argument("-i", "--history", type=str,
		help="Specify the history file to import")

	parser.add_argument("-a", "--ask", type=str,
		help="Asks a question directly from the command line")

	parser.add_argument("-o", "--output", type=str,
		help="Saves the history to the specified file")

	parser.add_argument("-t", "--tts", action='store_true',
		help="Enables TTS generation when asking")
	parser.add_argument("-s", "--ttsfile", type=str,
		help="Outputs a generated TTS file single-ask mode")

	parser.add_argument("-T", "--training", action='store_true',
		help="outputs training data from -i file")
	parser.add_argument("-e", "--estimate", action='store_true',
		help="shows estimate token usage / cost from a request from -i file")

	return parser.parse_args()


# finds if any element in the list a is present in the list b
def is_la_in_lb(a, b):
	for e in a:
		if e in b:
			return True
	return False


# ask a single question (with a new context)
def single_ask(
	user_prompt,
	temp_file: str = "",
	input_file: str = "",
	temp_is_temp: bool = False,
	should_print: bool = False
):
	if not temp_file:
		temp_is_temp = True
		temp_file = get_temp_file()
	messages = Conversation()
	connectLTS(
		messages.add_memory,
		messages.remove_memory,
		messages.edit_memory
	)
	if input_file:
		messages.load(input_file)
	messages = ask(
		prompt=user_prompt,
		messages=messages,
		model=baseConf.get("model", ''),
		temperature=baseConf.get("temperature", 0.8),
		max_tokens=baseConf.get("max_tokens", 3000),
		top_p=baseConf.get('top_p', 1.0),
		frequency_penalty=baseConf.get('frequency_penalty', 0.0),
		presence_penalty=baseConf.get('presence_penalty', 0.0)
	)
	if should_print:
		print(messages.last_answer())
	if baseConf.get('tts_enabled', False):
		tmpfile = tempfile.NamedTemporaryFile(
			prefix="owegatts.",
			suffix=".opus",
			delete=False
		)
		tmpfile.close()
		tts_answer = openai.audio.speech.create(
			model='tts-1',
			voice='nova',
			input=messages.last_answer()
		)
		tts_answer.stream_to_file(tmpfile.name)
		play_opus(tmpfile.name)
		os.remove(tmpfile.name)
	return messages.last_answer()
	if not temp_is_temp:
		messages.save(temp_file)


# main function, runs if __name__ is "__main__" ;)
def main():
	args = parse_args()

	if (args.debug):  # bypass before loading conf
		baseConf["debug"] = True

	if args.changelog:
		print(OwegaChangelog.log)
	if args.license:
		print(OwegaLicense)
	if args.version:
		print(f"Owega v{OwegaChangelog.version}")
	if (args.changelog or args.license or args.version):
		do_quit(value=1)
	if (args.training and not args.history):
		do_quit("Can't generate training data without a history", value=1)
	if args.training:
		msgs = Conversation_from(args.history)
		print(msgs.generate_training())
		sys.exit(0)
	if (args.estimate and not args.history):
		do_quit(
			"Can't estimate token consumption/cost without a history", value=1)
	if args.estimate:
		msgs = Conversation_from(args.history)
		etkn = estimated_tokens(
			'',
			msgs,
			''
		)
		cost_per_token = (
			0.003
		) / 1000
		cost = cost_per_token * etkn
		print(f"estimated tokens: {etkn}")
		print(f"estimated cost: {cost:.5f}$ (gpt-3) / {cost*10:.5f}$ (gpt-4)")
		sys.exit(0)

	input_history = ""
	if (args.history):
		input_history = args.history

	temp_file = get_temp_file()
	temp_is_temp = True
	if (args.output):
		temp_is_temp = False
		temp_file = args.output

	get_conf(args.config_file)
	if baseConf.get("commands", False):
		existingFunctions.enableGroup("utility.system")
	else:
		existingFunctions.disableGroup("utility.system")

	if (args.debug):  # bypass after loading conf
		baseConf["debug"] = True

	if (args.tts):
		baseConf["tts_enabled"] = True

	if (args.ask):
		answer = single_ask(
			args.ask,
			temp_file,
			input_history,
			temp_is_temp,
			True
		)
		if (args.ttsfile):
			tts_answer = openai.audio.speech.create(
				model="tts-1",
				voice="nova",
				input=answer
			)
			if (("opus" not in args.ttsfile)
				and ("mp3" not in args.ttsfile)
				and ("aac" not in args.ttsfile)
				and ("flac" not in args.ttsfile)):
				args.ttsfile = args.ttsfile + '.opus'
			tts_answer.stream_to_file(args.ttsfile)
	else:
		try:
			user_interaction_loop(
				temp_file=temp_file,
				input_file=input_history,
				temp_is_temp=temp_is_temp
			)
		except EOFError:
			do_quit(
				success_msg(),
				temp_file=temp_file,
				is_temp=temp_is_temp,
				should_del=temp_is_temp
			)
		except KeyboardInterrupt:
			do_quit(
				success_msg(),
				temp_file=temp_file,
				is_temp=temp_is_temp,
				should_del=False
			)


if __name__ == "__main__":
	main()
