#!/usr/bin/env python3
"""Owega's main function. Handle the CLI/TUI."""
# Import the necessary modules
import argparse
import getpass
import sys
import time

import openai
import prompt_toolkit as pt

from .ask import ask, single_ask
from .changelog import OwegaChangelog
from .config import baseConf, get_conf
from .conversation import Conversation, Conversation_from
from .constants import (
    OWEGA_DEFAULT_PROMPT,
    OWEGA_DEFAULT_MAX_TOKENS,
    OWEGA_DEFAULT_TEMPERATURE,
    OWEGA_DEFAULT_TOP_P,
    OWEGA_DEFAULT_FREQUENCY_PENALTY,
    OWEGA_DEFAULT_PRESENCE_PENALTY
)
from .license import OwegaLicense
from .OwegaFun import connectLTS, existingFunctions, functionlist_to_toollist
from .OwegaSession import OwegaSession as ps
from .OweHandlers import handle_help, handlers
from .utils import (do_quit, estimated_tokens_and_cost, get_temp_file,
                    info_print, markdown_print, play_tts, set_term_title,
                    success_msg)
from .colors import clrtxt


def get_oc_conf() -> dict:
    """Get a copy of owega's config."""
    return baseConf.copy()


def user_interaction_loop(
    temp_file: str = "",
    input_file: str = "",
    temp_is_temp: bool = False
) -> None:
    """Loop for the main interaction function."""
    if not (sys.stdin.isatty() and sys.stdout.isatty()):
        return

    if not temp_file:
        temp_is_temp = True
        temp_file = get_temp_file()

    default_prompt = OWEGA_DEFAULT_PROMPT
    # creates Conversation object and populate it
    messages: Conversation = Conversation(
        baseConf.get('default_prompt', default_prompt)
    )
    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 openai.api_key is not None:
                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)

        new_title = "Owega"
        new_title += f" v{OwegaChangelog.version}"
        new_title += f" - {baseConf.get('model', '?')}"
        new_title += f" - {baseConf.get('temperature', '?')}"
        new_title += f"/{baseConf.get('top_p', '?')}"
        new_title += f"/{baseConf.get('frequency_penalty', '?')}"
        new_title += f"/{baseConf.get('presence_penalty', '?')}"
        set_term_title(new_title)

        # get user input, and strip it (no excess spaces / tabs / newlines
        if ps['main'] is not None:
            user_input = ps['main'].prompt(pt.ANSI(input_prompt)).strip()
        else:
            user_input = input(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 user_input.strip():
                if baseConf.get("time_awareness", False):
                    date_str = time.strftime("[%Y-%m-%d %H:%M:%S]")
                    user_input = f"{date_str}\n" + user_input
            if baseConf.get("estimation", False):
                etkn, cost = estimated_tokens_and_cost(
                    user_input,
                    messages,
                    functionlist_to_toollist(existingFunctions.getEnabled()),
                    baseConf.get('model', ''),
                    baseConf.get("max_tokens", OWEGA_DEFAULT_MAX_TOKENS)
                )
                print(f"\033[37mestimated tokens: {etkn}\033[0m")
                print(f"\033[37mestimated cost: {cost:.5f}\033[0m")
            pre_time = time.time()
            last_ans = messages.last_answer()
            messages = ask(
                prompt=user_input,
                messages=messages,
                model=baseConf.get("model", ''),
                temperature=baseConf.get(
                    "temperature", OWEGA_DEFAULT_TEMPERATURE),
                max_tokens=baseConf.get("max_tokens", OWEGA_DEFAULT_MAX_TOKENS),
                top_p=baseConf.get("top_p", OWEGA_DEFAULT_TOP_P),
                frequency_penalty=baseConf.get(
                    "frequency_penalty", OWEGA_DEFAULT_FREQUENCY_PENALTY),
                presence_penalty=baseConf.get(
                    "presence_penalty", OWEGA_DEFAULT_PRESENCE_PENALTY)
            )
            new_ans = messages.last_answer()
            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
            if (new_ans != last_ans):
                print()
                print(' ' + clrtxt("magenta", " Owega ") + ": ")
                print()
                markdown_print(messages.last_answer())

                if baseConf.get('tts_enabled', False):
                    play_tts(messages.last_answer())


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    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()


def main() -> None:
    """Run the main function and handle the CLI/TUI."""
    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, cost = estimated_tokens_and_cost(
            '',
            msgs,
            [],
            baseConf.get('model', ''),
            0
        )
        print(f"estimated tokens: {etkn}")
        print(f"estimated cost: {cost:.5f}$ ({baseConf.get('model', '')})")
        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.write_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()
