import argparse
import asyncio
import datetime
import importlib
import sys
from typing import Any, Dict, List, Optional

import structlog

import rasa.cli.utils
import rasa.shared.utils.cli
from rasa.cli import SubParsersAction
from rasa.cli.arguments.default_arguments import (
    add_endpoint_param,
    add_model_param,
    add_remote_storage_param,
)
from rasa.core.agent import Agent
from rasa.core.exceptions import AgentNotReady
from rasa.core.processor import MessageProcessor
from rasa.core.utils import AvailableEndpoints
from rasa.dialogue_understanding.commands import Command
from rasa.dialogue_understanding.generator import (
    LLMBasedCommandGenerator,
)
from rasa.dialogue_understanding.generator.command_parser import DEFAULT_COMMANDS
from rasa.dialogue_understanding_test.command_metric_calculation import (
    calculate_command_metrics,
)
from rasa.dialogue_understanding_test.constants import (
    DEFAULT_INPUT_TESTS_PATH,
    KEY_STUB_CUSTOM_ACTIONS,
)
from rasa.dialogue_understanding_test.du_test_result import (
    DialogueUnderstandingTestResult,
    DialogueUnderstandingTestSuiteResult,
)
from rasa.dialogue_understanding_test.du_test_runner import (
    DialogueUnderstandingTestRunner,
)
from rasa.dialogue_understanding_test.io import (
    read_test_suite,
    write_test_results_to_file,
)
from rasa.dialogue_understanding_test.validation import (
    validate_cli_arguments,
    validate_test_cases,
)
from rasa.e2e_test.e2e_test_case import TestSuite
from rasa.exceptions import RasaException
from rasa.shared.constants import (
    DEFAULT_ENDPOINTS_PATH,
    LLM_CONFIG_KEY,
    ROUTE_TO_CALM_SLOT,
)
from rasa.shared.core.domain import Domain
from rasa.shared.core.flows import FlowsList
from rasa.shared.utils.llm import (
    combine_custom_and_default_config,
    resolve_model_client_config,
)
from rasa.utils.beta import ensure_beta_feature_is_enabled
from rasa.utils.endpoints import EndpointConfig

RASA_PRO_BETA_DIALOGUE_UNDERSTANDING_TEST_ENV_VAR_NAME = (
    "RASA_PRO_BETA_DIALOGUE_UNDERSTANDING_TEST"
)

structlogger = structlog.get_logger()


def add_subparser(
    subparsers: SubParsersAction, parents: List[argparse.ArgumentParser]
) -> None:
    """Add the dialogue understanding test subparser to `rasa test`.

    Args:
        subparsers: subparser we are going to attach to
        parents: Parent parsers, needed to ensure tree structure in argparse
    """
    for subparser in subparsers.choices.values():
        if subparser.prog == "rasa test":
            du_test_subparser = create_du_test_subparser(parents)

            for action in subparser._subparsers._actions:
                if action.choices is not None:
                    action.choices["du"] = du_test_subparser
                    return

    # If we get here, we couldn't hook the subparser to `rasa test`
    raise RasaException(
        "Hooking the dialogue understanding (du) test subparser to "
        "`rasa test` command could not be completed. "
        "Cannot run dialogue understanding testing."
    )


def create_du_test_subparser(
    parents: List[argparse.ArgumentParser],
) -> argparse.ArgumentParser:
    """Create dialogue understanding test subparser."""
    du_test_subparser = argparse.ArgumentParser(
        prog="rasa test du",
        parents=parents,
        conflict_handler="resolve",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Runs dialogue understanding testing.",
    )

    du_test_subparser.set_defaults(func=execute_dialogue_understanding_tests)

    add_du_test_arguments(du_test_subparser)
    add_bot_arguments(du_test_subparser)

    return du_test_subparser


def add_bot_arguments(parser: argparse.ArgumentParser) -> None:
    bot_arguments = parser.add_argument_group("Bot Settings")
    add_model_param(bot_arguments, add_positional_arg=False)
    add_endpoint_param(
        bot_arguments,
        help_text="Configuration file for the model server and the connectors as a "
        "yml file.",
    )
    add_remote_storage_param(bot_arguments)


def add_du_test_arguments(parser: argparse.ArgumentParser) -> None:
    """Arguments for running dialogue understanding tests."""
    du_arguments = parser.add_argument_group("Testing Settings")
    du_arguments.add_argument(
        "path-to-test-cases",
        nargs="?",
        type=str,
        default=DEFAULT_INPUT_TESTS_PATH,
        help="Input file or folder containing dialogue understanding test cases.",
    )
    du_arguments.add_argument(
        "--output-file",
        type=str,
        default="dialogue_understanding_test_{date:%Y%m%d-%H%M%S}.yml".format(
            date=datetime.datetime.now()
        ),
        help="Path to the output file to write the results to.",
    )
    du_arguments.add_argument(
        "--no-output",
        action="store_true",
        help="If set, no output file will be written to disk.",
    )
    du_arguments.add_argument(
        "--output-prompt",
        action="store_true",
        help="If set, the dialogue understanding test output will contain "
        "prompts for each failure.",
    )
    du_arguments.add_argument(
        "--additional-commands",
        type=str,
        nargs="*",
        help=(
            "List of additional custom command classes to add, separated by spaces. "
            "For example: --additional-commands my_module.MyCustomCommand"
        ),
    )
    du_arguments.add_argument(
        "--remove-default-commands",
        type=str,
        nargs="*",
        help=(
            f"List of default commands to remove, separated by spaces. "
            f"Default commands include: "
            f"{', '.join([command.__name__ for command in DEFAULT_COMMANDS])}. "
            f"For example: --remove-default-commands ClarifyCommand HumanHandoffCommand"
        ),
    )


def ensure_calm_only_bot(agent: Agent) -> None:
    if agent.domain is None or agent.processor is None:
        return

    if ROUTE_TO_CALM_SLOT in [slot.name for slot in agent.domain.slots]:
        rasa.shared.utils.cli.print_error(
            "You are using coexistence. Dialogue Understanding Tests do only work for "
            "CALM only assistants."
        )
        sys.exit(0)

    if not agent.processor.is_calm_assistant:
        rasa.shared.utils.cli.print_error(
            "Dialogue Understanding Tests do only work for CALM assistants. "
            "Your assistant does not use CALM."
        )
        sys.exit(0)


def execute_dialogue_understanding_tests(args: argparse.Namespace) -> None:
    """Run the dialogue understanding tests.

    Args:
        args: Commandline arguments.
    """
    ensure_beta_feature_is_enabled(
        "Dialogue Understanding (DU) Testing",
        env_flag=RASA_PRO_BETA_DIALOGUE_UNDERSTANDING_TEST_ENV_VAR_NAME,
    )

    # basic validation of the passed CLI arguments
    validate_cli_arguments(args)

    # initialization of endpoints
    endpoints = set_up_available_endpoints(args)

    # set up the test runner, e.g. start the agent
    try:
        test_runner = DialogueUnderstandingTestRunner(
            endpoints=endpoints,
            model_path=args.model,
            model_server=endpoints.model,
            remote_storage=args.remote_storage,
        )
    except AgentNotReady as error:
        structlogger.error(
            "rasa.dialogue_understanding_test.agent_not_ready", message=error.message
        )
        sys.exit(1)

    # Exit if the bot is not calm only
    ensure_calm_only_bot(test_runner.agent)

    # Ensure processor is not None so that we can extract the flows and the llm config
    if test_runner.agent.processor is None:
        rasa.shared.utils.cli.print_error(
            "No processor: Not able to retrieve flows and config from trained model."
        )
        sys.exit(0)

    # flows are needed in order to parse the commands when reading the test cases
    flows = asyncio.run(test_runner.agent.processor.get_flows())
    # llm config is needed for instrumentation
    llm_config = _get_llm_command_generator_config(test_runner.agent.processor)

    # read test cases from the given path
    test_suite: TestSuite = get_valid_test_suite(args, flows, test_runner.agent.domain)

    # setup stub custom actions if they are used
    set_up_stub_custom_actions(test_suite, endpoints)

    # run the actual test cases
    test_results = asyncio.run(
        test_runner.run_test_cases(
            test_suite.test_cases, test_suite.fixtures, test_suite.metadata
        )
    )

    # evaluate test results
    passing_test_results, failing_test_results = split_test_results(test_results)
    command_metrics = calculate_command_metrics(test_results)

    test_suite_result = DialogueUnderstandingTestSuiteResult.from_results(
        failing_test_results, passing_test_results, command_metrics, llm_config
    )

    # Do not move this import to the top of the file as it will break the
    # instrumentation of this function: the CLI module is initialized before the
    # instrumentation is set up, and we won't be able to "replace" the function
    # with the instrumented wrapper
    from rasa.dialogue_understanding_test.io import print_test_results

    # write results to console and file
    print_test_results(test_suite_result, output_prompt=args.output_prompt)
    if not args.no_output:
        write_test_results_to_file(
            test_suite_result,
            args.output_file,
            args.output_prompt,
        )


def _import_custom_command_class(class_path: str) -> Command:
    """Dynamically import a command class from a string path."""
    try:
        module_name, class_name = class_path.rsplit(".", 1)
        module = importlib.import_module(module_name)
        clz = getattr(module, class_name)
    except (ImportError, AttributeError, ValueError) as e:
        raise ValueError(f"Failed to import class '{class_path}': {e}")
    if not issubclass(clz, Command):
        structlogger.error(
            "rasa.dialogue_understanding_test.invalid_additional_command",
            event_info="The custom command class must be a subclass of Command.",
            class_path=class_path,
        )
        sys.exit(1)
    return clz


def _extract_additional_command_classes_from_cli_args(
    args: argparse.Namespace,
) -> List[Command]:
    """Extract additional command classes from the CLI arguments."""
    additional_commands = getattr(args, "additional_commands", [])
    if not additional_commands:
        return []
    return [
        _import_custom_command_class(command_module)
        for command_module in additional_commands
    ]


def get_valid_test_suite(
    args: argparse.Namespace, flows: FlowsList, domain: Optional[Domain]
) -> TestSuite:
    """Read the test cases from the given test case path and validate them."""
    path_to_test_cases = getattr(args, "path-to-test-cases", DEFAULT_INPUT_TESTS_PATH)
    remove_default_commands = getattr(args, "remove_default_commands", [])
    custom_command_classes = _extract_additional_command_classes_from_cli_args(args)
    test_suite = read_test_suite(
        path_to_test_cases, flows, custom_command_classes, remove_default_commands
    )
    validate_test_cases(test_suite.test_cases, domain)
    return test_suite


def set_up_available_endpoints(args: argparse.Namespace) -> AvailableEndpoints:
    """Set up the available endpoints for the test runner."""
    args.endpoints = rasa.cli.utils.get_validated_path(
        args.endpoints, "endpoints", DEFAULT_ENDPOINTS_PATH, True
    )
    endpoints = AvailableEndpoints.get_instance(args.endpoints)

    # Ignore all endpoints apart from action server, model, and nlu
    # to ensure InMemoryTrackerStore is being used instead of production
    # tracker store
    endpoints.tracker_store = None
    endpoints.lock_store = None
    endpoints.event_broker = None

    # disable nlg endpoint as we don't need it for dialogue understanding tests
    endpoints.nlg = None

    return endpoints


def set_up_stub_custom_actions(
    test_suite: TestSuite, endpoints: AvailableEndpoints
) -> None:
    """Set up the stub custom actions if they are used."""
    if test_suite.stub_custom_actions:
        if not endpoints.action:
            endpoints.action = EndpointConfig()

        endpoints.action.kwargs[KEY_STUB_CUSTOM_ACTIONS] = (
            test_suite.stub_custom_actions
        )


def split_test_results(
    results: List[DialogueUnderstandingTestResult],
) -> tuple[
    List[DialogueUnderstandingTestResult], List[DialogueUnderstandingTestResult]
]:
    """Split the test results into passed and failed test cases."""
    passed_cases = [r for r in results if r.passed]
    failed_cases = [r for r in results if not r.passed]

    return passed_cases, failed_cases


def _get_llm_command_generator_config(
    processor: MessageProcessor,
) -> Optional[Dict[str, Any]]:
    from rasa.dialogue_understanding.generator.constants import DEFAULT_LLM_CONFIG

    train_schema = processor.model_metadata.train_schema

    for node_name, node in train_schema.nodes.items():
        if node.matches_type(LLMBasedCommandGenerator, include_subtypes=True):
            # Configurations can reference model groups defined in the endpoints.yml
            resolved_config = resolve_model_client_config(
                node.config.get(LLM_CONFIG_KEY, {}), node_name
            )
            return combine_custom_and_default_config(
                resolved_config, DEFAULT_LLM_CONFIG
            )

    return None
