import argparse
import os
import sys
import time

from importlib.metadata import version, PackageNotFoundError

from unishell.terminal_interface.contributing_conversations import (
    contribute_conversation_launch_logic,
    contribute_conversations,
)

from .conversation_navigator import conversation_navigator
from .profiles.profiles import open_storage_dir, profile, reset_profile
from .utils.check_for_update import check_for_update
from .validate_llm_settings import validate_llm_settings


def start_terminal_interface(unishell):
    """
    Meant to be used from the command line. Parses arguments, starts OI's terminal interface.
    """

    # Instead use an async unishell, which has a server. Set settings on that
    if "--server" in sys.argv:
        from unishell import AsyncInterpreter

        unishell = AsyncInterpreter()

    arguments = [
        {
            "name": "profile",
            "nickname": "p",
            "help_text": "name of profile. run `--profiles` to open profile directory",
            "type": str,
            "default": "default.yaml",
        },
        {
            "name": "custom_instructions",
            "nickname": "ci",
            "help_text": "custom instructions for the language model. will be appended to the system_message",
            "type": str,
            "attribute": {"object": unishell, "attr_name": "custom_instructions"},
        },
        {
            "name": "system_message",
            "nickname": "sm",
            "help_text": "(we don't recommend changing this) base prompt for the language model",
            "type": str,
            "attribute": {"object": unishell, "attr_name": "system_message"},
        },
        {
            "name": "auto_run",
            "nickname": "y",
            "help_text": "automatically run generated code",
            "type": bool,
            "attribute": {"object": unishell, "attr_name": "auto_run"},
        },
        {
            "name": "no_highlight_active_line",
            "nickname": "nhl",
            "help_text": "turn off active line highlighting in code blocks",
            "type": bool,
            "action": "store_true",
            "default": False,  # Default to False, meaning highlighting is on by default
        },
        {
            "name": "verbose",
            "nickname": "v",
            "help_text": "print detailed logs",
            "type": bool,
            "attribute": {"object": unishell, "attr_name": "verbose"},
        },
        {
            "name": "model",
            "nickname": "m",
            "help_text": "language model to use",
            "type": str,
            "attribute": {"object": unishell.llm, "attr_name": "model"},
        },
        {
            "name": "temperature",
            "nickname": "t",
            "help_text": "optional temperature setting for the language model",
            "type": float,
            "attribute": {"object": unishell.llm, "attr_name": "temperature"},
        },
        {
            "name": "llm_supports_vision",
            "nickname": "lsv",
            "help_text": "inform OI that your model supports vision, and can receive vision inputs",
            "type": bool,
            "action": argparse.BooleanOptionalAction,
            "attribute": {"object": unishell.llm, "attr_name": "supports_vision"},
        },
        {
            "name": "llm_supports_functions",
            "nickname": "lsf",
            "help_text": "inform OI that your model supports OpenAI-style functions, and can make function calls",
            "type": bool,
            "action": argparse.BooleanOptionalAction,
            "attribute": {"object": unishell.llm, "attr_name": "supports_functions"},
        },
        {
            "name": "context_window",
            "nickname": "cw",
            "help_text": "optional context window size for the language model",
            "type": int,
            "attribute": {"object": unishell.llm, "attr_name": "context_window"},
        },
        {
            "name": "max_tokens",
            "nickname": "x",
            "help_text": "optional maximum number of tokens for the language model",
            "type": int,
            "attribute": {"object": unishell.llm, "attr_name": "max_tokens"},
        },
        {
            "name": "max_budget",
            "nickname": "b",
            "help_text": "optionally set the max budget (in USD) for your llm calls",
            "type": float,
            "attribute": {"object": unishell.llm, "attr_name": "max_budget"},
        },
        {
            "name": "api_base",
            "nickname": "ab",
            "help_text": "optionally set the API base URL for your llm calls (this will override environment variables)",
            "type": str,
            "attribute": {"object": unishell.llm, "attr_name": "api_base"},
        },
        {
            "name": "api_key",
            "nickname": "ak",
            "help_text": "optionally set the API key for your llm calls (this will override environment variables)",
            "type": str,
            "attribute": {"object": unishell.llm, "attr_name": "api_key"},
        },
        {
            "name": "api_version",
            "nickname": "av",
            "help_text": "optionally set the API version for your llm calls (this will override environment variables)",
            "type": str,
            "attribute": {"object": unishell.llm, "attr_name": "api_version"},
        },
        {
            "name": "max_output",
            "nickname": "xo",
            "help_text": "optional maximum number of characters for code outputs",
            "type": int,
            "attribute": {"object": unishell, "attr_name": "max_output"},
        },
        {
            "name": "loop",
            "help_text": "runs OI in a loop, requiring it to admit to completing/failing task",
            "type": bool,
            "attribute": {"object": unishell, "attr_name": "loop"},
        },
        {
            "name": "disable_telemetry",
            "nickname": "dt",
            "help_text": "disables sending of basic anonymous usage stats",
            "type": bool,
            "default": False,
            "attribute": {"object": unishell, "attr_name": "disable_telemetry"},
        },
        {
            "name": "offline",
            "nickname": "o",
            "help_text": "turns off all online features (except the language model, if it's hosted)",
            "type": bool,
            "attribute": {"object": unishell, "attr_name": "offline"},
        },
        {
            "name": "speak_messages",
            "nickname": "sp",
            "help_text": "(Mac only, experimental) use the applescript `say` command to read messages aloud",
            "type": bool,
            "attribute": {"object": unishell, "attr_name": "speak_messages"},
        },
        {
            "name": "safe_mode",
            "nickname": "safe",
            "help_text": "optionally enable safety mechanisms like code scanning; valid options are off, ask, and auto",
            "type": str,
            "choices": ["off", "ask", "auto"],
            "default": "off",
            "attribute": {"object": unishell, "attr_name": "safe_mode"},
        },
        {
            "name": "debug",
            "nickname": "debug",
            "help_text": "debug mode for open unishell developers",
            "type": bool,
            "attribute": {"object": unishell, "attr_name": "debug"},
        },
        {
            "name": "fast",
            "nickname": "f",
            "help_text": "runs `unishell --model gpt-4o-mini` and asks OI to be extremely concise (shortcut for `unishell --profile fast`)",
            "type": bool,
        },
        {
            "name": "multi_line",
            "nickname": "ml",
            "help_text": "enable multi-line inputs starting and ending with ```",
            "type": bool,
            "attribute": {"object": unishell, "attr_name": "multi_line"},
        },
        {
            "name": "local",
            "nickname": "l",
            "help_text": "setup a local model (shortcut for `unishell --profile local`)",
            "type": bool,
        },
        {
            "name": "codestral",
            "help_text": "shortcut for `unishell --profile codestral`",
            "type": bool,
        },
        {
            "name": "assistant",
            "help_text": "shortcut for `unishell --profile assistant.py`",
            "type": bool,
        },
        {
            "name": "llama3",
            "help_text": "shortcut for `unishell --profile llama3`",
            "type": bool,
        },
        {
            "name": "groq",
            "help_text": "shortcut for `unishell --profile groq`",
            "type": bool,
        },
        {
            "name": "vision",
            "nickname": "vi",
            "help_text": "experimentally use vision for supported languages (shortcut for `unishell --profile vision`)",
            "type": bool,
        },
        {
            "name": "os",
            "nickname": "os",
            "help_text": "experimentally let Open UniShell control your mouse and keyboard (shortcut for `unishell --profile os`)",
            "type": bool,
        },
        # Special commands
        {
            "name": "reset_profile",
            "help_text": "reset a profile file. run `--reset_profile` without an argument to reset all default profiles",
            "type": str,
            "default": "NOT_PROVIDED",
            "nargs": "?",  # This means you can pass in nothing if you want
        },
        {"name": "profiles", "help_text": "opens profiles directory", "type": bool},
        {
            "name": "local_models",
            "help_text": "opens local models directory",
            "type": bool,
        },
        {
            "name": "conversations",
            "help_text": "list conversations to resume",
            "type": bool,
        },
        {
            "name": "server",
            "help_text": "start open unishell as a server",
            "type": bool,
        },
        {
            "name": "version",
            "help_text": "get Open UniShell's version number",
            "type": bool,
        },
        {
            "name": "contribute_conversation",
            "help_text": "let Open UniShell use the current conversation to train an Open-Source LLM",
            "type": bool,
            "attribute": {
                "object": unishell,
                "attr_name": "contribute_conversation",
            },
        },
        {
            "name": "plain",
            "nickname": "pl",
            "help_text": "set output to plain text",
            "type": bool,
            "attribute": {
                "object": unishell,
                "attr_name": "plain_text_display",
            },
        },
        {
            "name": "stdin",
            "nickname": "s",
            "help_text": "Run OI in stdin mode",
            "type": bool,
        },
        {
            "name": "use_actions",
            "nickname": "ua",
            "help_text": "enable dynamic action execution (tries actions before code generation)",
            "type": bool,
            "action": argparse.BooleanOptionalAction,
            "attribute": {"object": unishell, "attr_name": "use_actions"},
        },
        {
            "name": "agent",
            "nickname": "ag",
            "help_text": "start desktop agent mode for gateway integration",
            "type": bool,
        },
    ]

    if "--stdin" in sys.argv and "--plain" not in sys.argv:
        sys.argv += ["--plain"]

    # i shortcut
    if len(sys.argv) > 1 and not sys.argv[1].startswith("-"):
        message = " ".join(sys.argv[1:])
        unishell.messages.append(
            {"role": "user", "type": "message", "content": "I " + message}
        )
        sys.argv = sys.argv[:1]

        unishell.custom_instructions = "UPDATED INSTRUCTIONS: You are in ULTRA FAST, ULTRA CERTAIN mode. Do not ask the user any questions or run code to gathet information. Go as quickly as you can. Run code quickly. Do not plan out loud, simply start doing the best thing. The user expects speed. Trust that the user knows best. Just interpret their ambiguous command as quickly and certainly as possible and try to fulfill it IN ONE COMMAND, assuming they have the right information. If they tell you do to something, just do it quickly in one command, DO NOT try to get more information (for example by running `cat` to get a file's infomration— this is probably unecessary!). DIRECTLY DO THINGS AS FAST AS POSSIBLE."

        files_in_directory = os.listdir()[:100]
        unishell.custom_instructions += (
            "\nThe files in CWD, which THE USER MAY BE REFERRING TO, are: "
            + ", ".join(files_in_directory)
        )

        # unishell.debug = True

    # Check for deprecated flags before parsing arguments
    deprecated_flags = {
        "--debug_mode": "--verbose",
    }

    for old_flag, new_flag in deprecated_flags.items():
        if old_flag in sys.argv:
            print(f"\n`{old_flag}` has been renamed to `{new_flag}`.\n")
            time.sleep(1.5)
            sys.argv.remove(old_flag)
            sys.argv.append(new_flag)

    class CustomHelpParser(argparse.ArgumentParser):
        def print_help(self, *args, **kwargs):
            super().print_help(*args, **kwargs)
            special_help_message = '''
Open UniShell, 2024

Use """ to write multi-line messages.
            '''
            print(special_help_message)

    parser = CustomHelpParser(
        description="Open UniShell", usage="%(prog)s [options]"
    )

    # Add arguments
    for arg in arguments:
        default = arg.get("default")
        action = arg.get("action", "store_true")
        nickname = arg.get("nickname")

        name_or_flags = [f'--{arg["name"]}']
        if nickname:
            name_or_flags.append(f"-{nickname}")

        # Construct argument name flags
        flags = (
            [f"-{nickname}", f'--{arg["name"]}'] if nickname else [f'--{arg["name"]}']
        )

        if arg["type"] == bool:
            parser.add_argument(
                *flags,
                dest=arg["name"],
                help=arg["help_text"],
                action=action,
                default=default,
            )
        else:
            choices = arg.get("choices")
            parser.add_argument(
                *flags,
                dest=arg["name"],
                help=arg["help_text"],
                type=arg["type"],
                choices=choices,
                default=default,
                nargs=arg.get("nargs"),
            )

    args, unknown_args = parser.parse_known_args()

    # handle unknown arguments
    if unknown_args:
        print(f"\nUnrecognized argument(s): {unknown_args}")
        parser.print_usage()
        print(
            "For detailed documentation of supported arguments, please visit: https://docs.unishell.com/settings/all-settings"
        )
        sys.exit(1)

    if args.profiles:
        open_storage_dir("profiles")
        return

    if args.local_models:
        open_storage_dir("models")
        return

    if args.reset_profile is not None and args.reset_profile != "NOT_PROVIDED":
        reset_profile(
            args.reset_profile
        )  # This will be None if they just ran `--reset_profile`
        return

    if args.version:
        oi_version = version("unishell")
        update_name = "Developer Preview"  # Change this with each major update
        print(f"Open UniShell {oi_version} {update_name}")
        return

    if args.no_highlight_active_line:
        unishell.highlight_active_line = False

    # if safe_mode and auto_run are enabled, safe_mode disables auto_run
    if unishell.auto_run and (
        unishell.safe_mode == "ask" or unishell.safe_mode == "auto"
    ):
        setattr(unishell, "auto_run", False)

    ### Set attributes on unishell, so that a profile script can read the arguments passed in via the CLI

    set_attributes(args, arguments)

    ### Apply profile

    # Profile shortcuts, which should probably not exist:

    if args.fast:
        args.profile = "fast.yaml"

    if args.vision:
        args.profile = "vision.yaml"

    if args.os:
        args.profile = "os.py"

    if args.local:
        args.profile = "local.py"
        if args.vision:
            # This is local vision, set up moondream!
            unishell.computer.vision.load()
        if args.os:
            args.profile = "local-os.py"

    if args.codestral:
        args.profile = "codestral.py"
        if args.vision:
            args.profile = "codestral-vision.py"
        if args.os:
            args.profile = "codestral-os.py"

    if args.assistant:
        args.profile = "assistant.py"

    if args.llama3:
        args.profile = "llama3.py"
        if args.vision:
            args.profile = "llama3-vision.py"
        if args.os:
            args.profile = "llama3-os.py"

    if args.groq:
        args.profile = "groq.py"

    unishell = profile(
        unishell,
        args.profile or get_argument_dictionary(arguments, "profile")["default"],
    )

    ### Set attributes on unishell, because the arguments passed in via the CLI should override profile

    set_attributes(args, arguments)
    unishell.disable_telemetry = (
        os.getenv("DISABLE_TELEMETRY", "false").lower() == "true"
        or args.disable_telemetry
    )

    ### Set some helpful settings we know are likely to be true

    if unishell.llm.model == "gpt-4" or unishell.llm.model == "openai/gpt-4":
        if unishell.llm.context_window is None:
            unishell.llm.context_window = 6500
        if unishell.llm.max_tokens is None:
            unishell.llm.max_tokens = 4096
        if unishell.llm.supports_functions is None:
            unishell.llm.supports_functions = (
                False if "vision" in unishell.llm.model else True
            )

    elif unishell.llm.model.startswith("gpt-4") or unishell.llm.model.startswith(
        "openai/gpt-4"
    ):
        if unishell.llm.context_window is None:
            unishell.llm.context_window = 123000
        if unishell.llm.max_tokens is None:
            unishell.llm.max_tokens = 4096
        if unishell.llm.supports_functions is None:
            unishell.llm.supports_functions = (
                False if "vision" in unishell.llm.model else True
            )

    if unishell.llm.model.startswith(
        "gpt-3.5-turbo"
    ) or unishell.llm.model.startswith("openai/gpt-3.5-turbo"):
        if unishell.llm.context_window is None:
            unishell.llm.context_window = 16000
        if unishell.llm.max_tokens is None:
            unishell.llm.max_tokens = 4096
        if unishell.llm.supports_functions is None:
            unishell.llm.supports_functions = True

    ### Check for update

    try:
        if not unishell.offline and not args.stdin:
            # This message should actually be pushed into the utility
            if check_for_update():
                unishell.display_message(
                    "> **A new version of Open UniShell is available.**\n>Please run: `pip install --upgrade open-unishell`\n\n---"
                )
    except:
        # Doesn't matter
        pass

    if unishell.llm.api_base:
        if (
            not unishell.llm.model.lower().startswith("openai/")
            and not unishell.llm.model.lower().startswith("azure/")
            and not unishell.llm.model.lower().startswith("ollama")
            and not unishell.llm.model.lower().startswith("jan")
            and not unishell.llm.model.lower().startswith("local")
        ):
            unishell.llm.model = "openai/" + unishell.llm.model
        elif unishell.llm.model.lower().startswith("jan/"):
            # Strip jan/ from the model name
            unishell.llm.model = unishell.llm.model[4:]

    # If --conversations is used, run conversation_navigator
    if args.conversations:
        conversation_navigator(unishell)
        return

    if unishell.llm.model in [
        "claude-3.5",
        "claude-3-5",
        "claude-3.5-sonnet",
        "claude-3-5-sonnet",
    ]:
        unishell.llm.model = "claude-3-5-sonnet-20240620"

    if not args.server:
        # This SHOULD RUN WHEN THE SERVER STARTS. But it can't rn because
        # if you don't have an API key, a prompt shows up, breaking the whole thing.
        validate_llm_settings(
            unishell
        )  # This should actually just run unishell.llm.load() once that's == to validate_llm_settings

    if args.server:
        unishell.server.run()
        return

    unishell.in_terminal_interface = True

    contribute_conversation_launch_logic(unishell)

    # Standard in mode
    if args.stdin:
        stdin_input = input()
        unishell.plain_text_display = True
        unishell.chat(stdin_input)
    elif args.agent:
        # Agent mode
        import asyncio
        from unishell.core.agent import DesktopAgent, AgentConfig
        
        print("\n🤖 Desktop Agent Configuration")
        print("=" * 50)
        gateway_url = input("Gateway URL: ").strip()
        desktop_id = input("Desktop ID: ").strip()
        device_token = input("Device Token: ").strip()
        
        config = AgentConfig(
            gateway_url=gateway_url,
            desktop_id=desktop_id,
            device_token=device_token
        )
        
        agent = DesktopAgent(config, unishell)
        asyncio.run(agent.run())
    else:
        unishell.chat()


def set_attributes(args, arguments):
    for argument_name, argument_value in vars(args).items():
        if argument_value is not None:
            if argument_dictionary := get_argument_dictionary(arguments, argument_name):
                if "attribute" in argument_dictionary:
                    attr_dict = argument_dictionary["attribute"]
                    setattr(attr_dict["object"], attr_dict["attr_name"], argument_value)

                    if args.verbose:
                        print(
                            f"Setting attribute {attr_dict['attr_name']} on {attr_dict['object'].__class__.__name__.lower()} to '{argument_value}'..."
                        )


def get_argument_dictionary(arguments: list[dict], key: str) -> dict:
    if (
        len(
            argument_dictionary_list := list(
                filter(lambda x: x["name"] == key, arguments)
            )
        )
        > 0
    ):
        return argument_dictionary_list[0]
    return {}


def main():
    from unishell import unishell

    try:
        start_terminal_interface(unishell)
    except KeyboardInterrupt:
        try:
            unishell.computer.terminate()

            if not unishell.offline and not unishell.disable_telemetry:
                feedback = None
                if len(unishell.messages) > 3:
                    feedback = (
                        input("\n\nWas Open UniShell helpful? (y/n): ")
                        .strip()
                        .lower()
                    )
                    if feedback == "y":
                        feedback = True
                    elif feedback == "n":
                        feedback = False
                    else:
                        feedback = None
                    if feedback != None and not unishell.contribute_conversation:
                        if unishell.llm.model == "i":
                            contribute = "y"
                        else:
                            print(
                                "\nThanks for your feedback! Would you like to send us this chat so we can improve?\n"
                            )
                            contribute = input("(y/n): ").strip().lower()

                        if contribute == "y":
                            unishell.contribute_conversation = True
                            unishell.display_message(
                                "\n*Thank you for contributing!*\n"
                            )

                if (
                    unishell.contribute_conversation or unishell.llm.model == "i"
                ) and unishell.messages != []:
                    conversation_id = (
                        unishell.conversation_id
                        if hasattr(unishell, "conversation_id")
                        else None
                    )
                    contribute_conversations(
                        [unishell.messages], feedback, conversation_id
                    )

        except KeyboardInterrupt:
            pass
    finally:
        unishell.computer.terminate()
