#!/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 . import getLogger
from .ask import ask, single_ask, get_model_provider
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_dual, get_temp_file,
                    info_print, play_tts,
                    render_conversation_buffer, 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

    logger = getLogger.getLogger(__name__, debug=baseConf.get("debug", False))

    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:
        logger.debug("Loop start")
        # 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):
                # Get provider for accurate pricing
                model = baseConf.get('model', '')
                _, provider = get_model_provider(model)
                itkn, icost, _, _ = estimated_tokens_and_cost_dual(
                    user_input,
                    messages,
                    functionlist_to_toollist(existingFunctions.getEnabled()),
                    baseConf.get('model', ''),
                    baseConf.get("max_tokens", OWEGA_DEFAULT_MAX_TOKENS),
                    provider
                )
                print(f"\033[37mestimated input tokens: {itkn}")
                print(f"\033[37mestimated input cost: {icost:.5f}")
            pre_time = time.time()
            last_ans = messages.last_answer()
            try:
                logger.debug("Running ask()")
                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)
                )
            except KeyboardInterrupt:
                print("\r\x1b[K", end="")
                print("Request cancelled.")
            except ValueError as e:
                print()
                print(f"Error while running ask(): {e}")
            except Exception as e:
                print()
                print("An error has occurred while running ask():")
                print(e)
            else:
                # Message printing should be deferred to ask()
                # Also, message rendering should be standardized
                new_ans = messages.last_answer()
                if baseConf.get("debug", False):
                    post_time = time.time()
                    full_time = post_time-pre_time
                    print(f"\033[37mrequest took {full_time:.3f}s\033[0m")

                # Print the generated response
                # markdown_print(messages.printbuffer_get())
                render_conversation_buffer(messages)
                if (new_ans != last_ans):
                    # render_assistant(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)
        model = baseConf.get('model', '')
        _, provider = get_model_provider(model)
        itkn, icost, otkn, ocost = estimated_tokens_and_cost_dual(
            '',
            msgs,
            [],
            model,
            0,
            provider
        )
        fcost = icost + ocost
        print(f"estimated input tokens: {itkn}")
        print(f"estimated input cost: {icost:.5f}$ ({model})")
        print(f"estimated output tokens: {otkn}")
        print(f"estimated output cost: {ocost:.5f}$ ({model})")
        print(f"estimated cumulative cost for max output tokens: {fcost}$")
        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 (args.debug):  # bypass after loading conf
        baseConf["debug"] = True

    logger = getLogger.getLogger(__name__, debug=baseConf.get("debug", False))
    logger.debug("Config loaded, debug mode active.")

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

    if baseConf.get("web_access", False):
        existingFunctions.enableGroup("utility.user")
    else:
        existingFunctions.disableGroup("utility.user")

    if baseConf.get("lts_enabled", False):
        existingFunctions.enableGroup("lts")
    else:
        existingFunctions.disableGroup("lts")

    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()
