import argparse
import os
import sys
import time

from importlib.metadata import version, PackageNotFoundError

from probe.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
# New: provider selection and API key validation helpers
from probe.terminal_interface import provider_selection
from probe.core.utils import api_key_validation


def start_terminal_interface(probe):
    """
    Meant to be used from the command line. Parses arguments and starts Probe's terminal interface.
    """

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

        probe = AsyncProbe()

    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": probe, "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": probe, "attr_name": "system_message"},
        },
        {
            "name": "auto_run",
            "nickname": "y",
            "help_text": "automatically run generated code",
            "type": bool,
            "attribute": {"object": probe, "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": probe, "attr_name": "verbose"},
        },
        {
            "name": "model",
            "nickname": "m",
            "help_text": "language model to use",
            "type": str,
            "attribute": {"object": probe.llm, "attr_name": "model"},
        },
        {
            "name": "temperature",
            "nickname": "t",
            "help_text": "optional temperature setting for the language model",
            "type": float,
            "attribute": {"object": probe.llm, "attr_name": "temperature"},
        },
        {
            "name": "llm_supports_vision",
            "nickname": "lsv",
            "help_text": "inform Probe that your model supports vision, and can receive vision inputs",
            "type": bool,
            "action": argparse.BooleanOptionalAction,
            "attribute": {"object": probe.llm, "attr_name": "supports_vision"},
        },
        {
            "name": "llm_supports_functions",
            "nickname": "lsf",
            "help_text": "inform Probe that your model supports OpenAI-style functions, and can make function calls",
            "type": bool,
            "action": argparse.BooleanOptionalAction,
            "attribute": {"object": probe.llm, "attr_name": "supports_functions"},
        },
        {
            "name": "context_window",
            "nickname": "cw",
            "help_text": "optional context window size for the language model",
            "type": int,
            "attribute": {"object": probe.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": probe.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": probe.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": probe.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": probe.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": probe.llm, "attr_name": "api_version"},
        },
        {
            "name": "max_output",
            "nickname": "xo",
            "help_text": "optional maximum number of characters for code outputs",
            "type": int,
            "attribute": {"object": probe, "attr_name": "max_output"},
        },
        {
            "name": "loop",
            "help_text": "runs Probe in a loop, requiring it to admit to completing/failing task",
            "type": bool,
            "attribute": {"object": probe, "attr_name": "loop"},
        },
        {
            "name": "disable_telemetry",
            "nickname": "dt",
            "help_text": "disables sending of basic anonymous usage stats",
            "type": bool,
            "default": False,
            "attribute": {"object": probe, "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": probe, "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": probe, "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": probe, "attr_name": "safe_mode"},
        },
        {
            "name": "debug",
            "nickname": "debug",
            "help_text": "debug mode for open probe developers",
            "type": bool,
            "attribute": {"object": probe, "attr_name": "debug"},
        },
        {
            "name": "fast",
            "nickname": "f",
            "help_text": "runs `probe --model gpt-4o-mini` and asks OI to be extremely concise (shortcut for `probe --profile fast`)",
            "type": bool,
        },
        {
            "name": "multi_line",
            "nickname": "ml",
            "help_text": "enable multi-line inputs starting and ending with ```",
            "type": bool,
            "attribute": {"object": probe, "attr_name": "multi_line"},
        },
        {
            "name": "local",
            "nickname": "l",
            "help_text": "setup a local model (shortcut for `probe --profile local`)",
            "type": bool,
        },
        {
            "name": "codestral",
            "help_text": "shortcut for `probe --profile codestral`",
            "type": bool,
        },
        {
            "name": "assistant",
            "help_text": "shortcut for `probe --profile assistant.py`",
            "type": bool,
        },
        {
            "name": "llama3",
            "help_text": "shortcut for `probe --profile llama3`",
            "type": bool,
        },
        {
            "name": "groq",
            "help_text": "shortcut for `probe --profile groq`",
            "type": bool,
        },
        {
            "name": "vision",
            "nickname": "vi",
            "help_text": "experimentally use vision for supported languages (shortcut for `probe --profile vision`)",
            "type": bool,
        },
        {
            "name": "os",
            "nickname": "os",
            "help_text": "experimentally let Probe control your mouse and keyboard (shortcut for `probe --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 probe as a server",
            "type": bool,
        },
        {
            "name": "version",
            "help_text": "get Probe's version number",
            "type": bool,
        },
        {
            "name": "contribute_conversation",
            "help_text": "let Probe use the current conversation to train an Open-Source LLM",
            "type": bool,
            "attribute": {
                "object": probe,
                "attr_name": "contribute_conversation",
            },
        },
        {
            "name": "plain",
            "nickname": "pl",
            "help_text": "set output to plain text",
            "type": bool,
            "attribute": {
                "object": probe,
                "attr_name": "plain_text_display",
            },
        },
        {
            "name": "stdin",
            "nickname": "s",
            "help_text": "Run Probe in stdin mode",
            "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:])
        probe.messages.append(
            {"role": "user", "type": "message", "content": "I " + message}
        )
        sys.argv = sys.argv[:1]

        probe.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]
        probe.custom_instructions += (
            "\nThe files in CWD, which THE USER MAY BE REFERRING TO, are: "
            + ", ".join(files_in_directory)
        )

        # probe.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 = '''
Probe, 2024

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

    parser = CustomHelpParser(
        description="Probe", 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://github.com/an0dev/probe"
        )
        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:
        try:
            probe_version = version("probe")
        except Exception:
            # Fallback if version detection fails
            probe_version = "unknown"
        update_name = "Developer Preview"  # Change this with each major update
        print(f"Probe {probe_version} {update_name}")
        return

    if args.no_highlight_active_line:
        probe.highlight_active_line = False

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

    ### Set attributes on probe, 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!
            probe.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"

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

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

    set_attributes(args, arguments)
    probe.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 probe.llm.model == "gpt-4" or probe.llm.model == "openai/gpt-4":
        if probe.llm.context_window is None:
            probe.llm.context_window = 6500
        if probe.llm.max_tokens is None:
            probe.llm.max_tokens = 4096
        if probe.llm.supports_functions is None:
            probe.llm.supports_functions = (
                False if "vision" in probe.llm.model else True
            )

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

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

    ### Check for update

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

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

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

    if probe.llm.model in [
        "claude-3.5",
        "claude-3-5",
        "claude-3.5-sonnet",
        "claude-3-5-sonnet",
    ]:
        probe.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.
        try:
            # If user did not pass --model on the CLI, show provider selection menu (non-fatal)
            if not args.model:
                # debug
                if hasattr(probe, 'debug') and probe.debug:
                    print("[debug] invoking provider_selection")
                try:
                    sel = provider_selection.select_provider(
                        display_message_func=probe.display_message
                    )
                    if sel:
                        probe.llm.model = provider_selection.get_default_model_for_provider(
                            sel
                        )
                except Exception as exc:
                    # Best-effort: provider selection must not crash startup
                    if hasattr(probe, 'debug') and probe.debug:
                        print(f"Provider selection failed: {exc}", file=sys.stderr)

            # Validate LLM settings (including prompting for API keys) but do not crash on failure
            try:
                validate_llm_settings(probe)
            except KeyboardInterrupt:
                raise
            except EOFError:
                raise
            except Exception as e:
                if hasattr(probe, 'debug') and probe.debug:
                    print(f"LLM validation warning: {str(e)}", file=sys.stderr)
                # Continue even if validation had issues
        except KeyboardInterrupt:
            raise
        except EOFError:
            raise
        except Exception:
            # Defensive: do not allow provider/validation issues to stop startup
            if hasattr(probe, 'debug') and probe.debug:
                print("Unexpected error in LLM setup, continuing", file=sys.stderr)

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

    probe.in_terminal_interface = True

    contribute_conversation_launch_logic(probe)

    # Standard in mode
    if args.stdin:
        try:
            stdin_input = input()
            probe.plain_text_display = True
            probe.chat(stdin_input)
        except (KeyboardInterrupt, EOFError):
            # User cancelled input
            pass
        except Exception as e:
            # Log error if in debug mode
            if hasattr(probe, 'debug') and probe.debug:
                print(f"Error during stdin chat: {str(e)}", file=sys.stderr)
            raise
    else:
        try:
            probe.chat()
        except (KeyboardInterrupt, EOFError):
            # User cancelled chat
            pass
        except Exception:
            # Re-raise other exceptions
            raise


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 probe import probe

    try:
        start_terminal_interface(probe)
    except KeyboardInterrupt:
        try:
            probe.computer.terminate()
        except KeyboardInterrupt:
            pass
        except Exception:
            pass
    except EOFError:
        # Handle pipe closure or Ctrl+D
        try:
            probe.computer.terminate()
        except Exception:
            pass
    except Exception as e:
        # Log the error but try to clean up
        try:
            probe.computer.terminate()
        except Exception:
            pass
        # Re-raise the exception after cleanup
        raise
    finally:
        try:
            probe.computer.terminate()
        except Exception:
            pass
