#!/usr/bin/env python3

# SPDX-FileCopyrightText: Portions Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import argparse
from logging import getLogger
import random
import sys

import cmd2
from colorama import Fore, Style

import garak.harnesses.probewise
from garak import _config
from garak.evaluators import ThresholdEvaluator
from garak.data import path as data_path
from garak._plugins import enumerate_plugins

logger = getLogger(__name__)


def _get_list_parser():
    # Create parser for list command
    list_parser = cmd2.Cmd2ArgumentParser(
        description="List available probes, decorators, or generators",
        epilog="This command only accepts a single string and lists the relevant attributes.",
    )
    list_parser.add_argument(
        "type",
        type=str,
        choices=("probes", "detectors", "generators"),
        help="Specify probes, detectors, or generators.",
    )
    return list_parser


global list_parser
list_parser = _get_list_parser()

probe_parser = cmd2.Cmd2ArgumentParser(
    description="Run the probe.",
    epilog="Uses set probe if no probe name is provided. "
    "If a probe name is provided, executes that probe.",
)
probe_parser.add_argument(
    "probe",
    nargs="?",
    type=str,
    help="Name of the probe to execute if not already set.",
)


def print_plugins(prefix, color):
    plugin_names = enumerate_plugins(category=prefix)
    plugin_names = [(p.replace(f"{prefix}.", ""), a) for p, a in plugin_names]
    module_names = set([(m.split(".")[0], True) for m, a in plugin_names])
    plugin_names += module_names
    for plugin_name, active in sorted(plugin_names):
        print(f"{Style.BRIGHT}{color}{prefix}: {Style.RESET_ALL}", end="")
        print(plugin_name, end="")
        if "." not in plugin_name:
            print(" 🌟", end="")
        if not active:
            print(" 💤", end="")
        print()


@cmd2.with_default_category("Garak Commands")
class GarakCommands(cmd2.CommandSet):
    def __init__(self):
        """Initialize the Garak Commands object."""
        super().__init__()

    @cmd2.with_argparser(list_parser)
    def do_list(self, args):
        if not args.type:
            print("Choose probes, detectors, or generators.")

        if args.type == "probes":
            logger.debug("Listing probes")
            print_plugins("probes", Fore.LIGHTYELLOW_EX)

        elif args.type == "detectors":
            logger.debug("Listing detectors")
            print_plugins("detectors", Fore.LIGHTBLUE_EX)

        elif args.type == "generators":
            logger.debug("Listing generators")
            print_plugins("generators", Fore.LIGHTMAGENTA_EX)

        else:
            logger.debug("Invalid choice to `list` command.")
            print("Choose probes, detectors or generators.")

        list_parser = _get_list_parser()

    @cmd2.with_argparser(probe_parser)
    def do_probe(self, args):
        if not self._cmd.target_type or not self._cmd.target_model:
            print(
                "Use the `set` command to set the target_type and target_model first."
            )
            return
        # If probe is already set, overwrite it.
        if args.probe and self._cmd.probe:
            logger.warning(f"Probe already set. Resetting probe to {args.probe}")
            print(f"Executing {args.probe}")
            self._cmd.probe = args.probe
        elif not args.probe and not self._cmd.probe:
            logger.warning("No probe set and no probe specified.")
            return None
        try:
            if self._cmd.generator:
                generator_module_name = self._cmd.generator.split(".")[0]
                generator_name = self._cmd.generator
            else:
                generator_module_name = self._cmd.target_type
                generator_name = self._cmd.target_type

            gen_conf = {generator_module_name: {"name": self._cmd.target_model}}
            _config._combine_into(gen_conf, _config.plugins.generators)

            from garak import _plugins

            generator = _plugins.load_plugin(
                f"generators.{generator_name}", config_root=_config
            )

        except ImportError as e:
            logger.error(e)
            print("Could not load generator from Garak generators.")
        except AttributeError as e:
            logger.error(e)
            print("Please check your generator model name.")

        evaluator = ThresholdEvaluator(self._cmd.eval_threshold)
        harness = garak.harnesses.probewise.ProbewiseHarness()
        harness.run(generator, [self._cmd.probe], evaluator)
        logger.info("Run complete, ending")
        print("Run complete!")


class GarakTerminal(cmd2.Cmd):
    """Terminal class for Interactive Garak CLI"""

    _cmd = None

    def __init__(self):
        super().__init__(allow_cli_args=False, auto_load_commands=False)
        self._load_garak()
        self.prompt = "garak> "
        self.quit_message = "Thanks for hacking!"
        # Create settable parameters
        self.add_settable(cmd2.Settable("target_type", str, "Type of the target", self))
        self.add_settable(
            cmd2.Settable("target_model", str, "Name of the target", self)
        )
        self.add_settable(cmd2.Settable("probe", str, "Probe to execute", self))
        self.add_settable(cmd2.Settable("detector", str, "Detector to execute", self))
        self.add_settable(
            cmd2.Settable("generator", str, "Generator settings path", self)
        )
        self.add_settable(
            cmd2.Settable(
                "eval_threshold", float, "Evaluation threshold for success", self
            )
        )
        # Set default parameter values
        self.target_type = ""
        self.target_model = ""
        self.probe = ""
        self.detector = ""
        self.generator = ""
        self.eval_threshold = 0.5
        # Disable shell, script, alias, and edit commands
        self.disable_command("alias", "Command not available.")
        self.disable_command("edit", "Command not available.")
        self.disable_command("macro", "Command not available.")
        self.disable_command("run_pyscript", "Command not available.")
        self.disable_command("run_script", "Command not available.")
        self.disable_command("shell", "Command not available.")
        # We don't have any shortcuts. Disable this too.
        self.disable_command("shortcuts", "Command not available.")
        # Remove excess settables
        self.remove_settable("timing")
        self.remove_settable("quiet")
        self.remove_settable("max_completion_items")
        self.remove_settable("allow_style")
        self.remove_settable("always_show_hint")
        self.remove_settable("echo")
        self.remove_settable("editor")
        self.remove_settable("feedback_to_output")

    def default(self, command: str) -> None:
        """Execute when a command isn't recognized"""
        print(f"Command does not exist.\n")
        return None

    def postcmd(self, stop, line):
        """Set the prompt to reflect interaction changes."""
        target_type = self.target_type
        target_model = self.target_model
        active_probe = self.probe
        if not target_type or not target_model:
            self.prompt = "garak> "
            return None
        if not active_probe:
            self.prompt = f"{target_type}: {target_model}> "
        else:
            self.prompt = f"{target_type}: {target_model}>{active_probe}> "

        self._load_garak()
        return stop

    def _load_garak(self):
        if self._cmd:
            self.unregister_command_set(self._cmd)
        self._cmd = GarakCommands()
        self.register_command_set(self._cmd)

    @cmd2.with_argument_list
    def do_quit(self, args):
        print(self.quit_message)
        sys.exit(0)

    def settings_ns_provider(self) -> argparse.Namespace:
        """Populate an argparse Namespace with current settings"""
        ns = argparse.Namespace()
        ns.app_settings = self.settings
        return ns


def banner():
    """Display a random banner"""
    with open(data_path / "banners", "r", encoding="utf-8") as bannerfile:
        banners = bannerfile.read().lstrip().split("\n\n")
        print(Fore.LIGHTRED_EX + f"\n{random.choice(banners)}\n" + Fore.RESET)


def interactive_mode():
    """
    Launch garak in interactive mode. Allows users to explore and use garak without CLI flags.
    """
    banner()
    terminal = GarakTerminal()
    sys.exit(terminal.cmdloop())


if __name__ == "__main__":
    sys.stdout.reconfigure(encoding="utf-8")
    interactive_mode()
