# Copyright 2026 Arturo Roberti
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
from dataclasses import dataclass, field
from textwrap import dedent

import networkx as nx

from gurk.lib.context import get_logger
from gurk.lib.core.plugins import (
    GurkArgumentParser,
    create_plugin_venv,
    get_available_plugin_tasks,
    get_venv_package_version,
    remove_venv,
    venv_exists,
)
from gurk.lib.shared.dicts import fill_typed_dict
from gurk.lib.shared.scripts import Command, SchedulerTask
from gurk.lib.shared.tasks import (
    CustomTaskDictCollection,
    ResolvedArgsDefinitionCollection,
    ResolvedTaskDict,
    ResolvedTaskDictCollection,
)
from gurk.lib.utils import (
    GURK_VERSION,
    PACKAGE_NAME,
    generate_random_path,
    overlay_dicts,
    typecheck,
)


@dataclass
class Processor:
    """Class to process tasks based on provided options and CLI arguments."""

    # fmt: off
    option:      CustomTaskDictCollection = field()
    cli_args:    list[str]                = field()
    parser_base: GurkArgumentParser       = field()
    tasks:       list[SchedulerTask]      = field(init=False, default_factory=list)
    # fmt: on

    def __post_init__(self):
        # Get logger
        logger = get_logger()

        # Get all tasks
        tasks = get_available_plugin_tasks()

        # Extract task args
        task_args: ResolvedArgsDefinitionCollection = dict()
        for task_name, task in tasks.items():
            task_args[task_name] = task.pop("args")

        # Overlay option
        tasks = overlay_dicts([tasks, self.option])

        # Fill missing properties
        tasks = fill_typed_dict(tasks, ResolvedTaskDictCollection)

        # Enable dependencies of enabled tasks
        tasks = self.enable_dependencies(tasks)

        # Filter only enabled tasks
        tasks = {
            task_name: task
            for task_name, task in tasks.items()
            if task_name in self.option
        }
        task_args = {
            task_name: task_args[task_name] for task_name in tasks.keys()
        }
        logger.debug(f"Enabled tasks: {list(tasks.keys())}")

        # Extract run args
        run_args = dict()
        for task_name, task in self.option.items():
            run_args[task_name] = self.option.pop("args", [])

        # Create argparser, removing args already passed in option
        for task_name, args in task_args.items():
            if args:
                used_args = set(args.keys()) & set(
                    self.option.get("args", {}).keys()
                )
                unused_args = {
                    arg_name: args[arg_name]
                    for arg_name in args.keys()
                    if arg_name not in used_args
                }
                self.parser_base.extend_arguments(unused_args)

        # Parse cli args
        parsed_cli_args = self.parser_base.parse_args(self.cli_args)
        logger.debug("Parsed CLI args successfully")

        # If all args are good, re-insert them back
        for task_name, task in tasks.items():
            task["args"] = []

            # Re-assign run args, if any
            if task_name in run_args:
                task["args"].extend(run_args[task_name])

            # Extend with always-allowed args, if any
            if {"-f", "--force"} & set(self.cli_args):
                task["args"].append("--force")

            # Extend with cli args that belong to this task, if any
            common_arg_names = set(task_args[task_name].keys()) & set(
                self.cli_args
            )
            common_args = []
            for arg in task_args[task_name].keys():
                if arg not in common_arg_names:
                    continue

                arg_attr = arg.lstrip("-").replace("-", "_")
                value = getattr(parsed_cli_args, arg_attr)

                common_args.append(arg)
                if isinstance(value, bool):
                    # Already appended
                    pass
                elif isinstance(value, str):
                    common_args.append(value)
                elif isinstance(value, list):
                    for v in value:
                        common_args.append(v)

            task["args"].extend(common_args)

        logger.debug("Assigned args to resp. tasks successfully")

        # Ensure valid plugin states
        self.ensure_plugins(tasks)

        # Prepend system preparation task
        if "pytest" not in sys.modules:
            tasks = self.add_preparation_task(tasks)

        # Convert to SchedulerTask list
        for task_name, task in tasks.items():
            if task_name not in self.option:
                continue

            resolved_task = SchedulerTask(
                name=task_name,
                command=Command(task["script"], task["function"]),
                config_file=(
                    task["config_file"].as_posix()
                    if task["config_file"]
                    else None
                ),
                depends_on=tuple(task["depends_on"]),
                privileged=task["privileged"],
                args=tuple(task["args"]),
            )
            self.tasks.append(resolved_task)

    @typecheck
    def enable_dependencies(
        self, tasks: ResolvedTaskDictCollection
    ) -> ResolvedTaskDictCollection:
        """
        Enable dependencies of enabled tasks.

        :param tasks: Tasks to process
        :type tasks: ResolvedTaskDictCollection
        :return: Processed tasks with dependencies enabled
        :rtype: ResolvedTaskDictCollection
        """
        # Get logger
        logger = get_logger()

        # Build dependency graph
        dependency_graph = nx.DiGraph()
        for task_name, task in tasks.items():
            dependency_graph.add_node(task_name)
            for dep in task["depends_on"]:
                dependency_graph.add_edge(dep, task_name)

        # Enable dependencies of enabled tasks
        for node in nx.topological_sort(dependency_graph):
            if node in self.option:
                for dep in nx.ancestors(dependency_graph, node):
                    if dep not in self.option:
                        self.option[dep] = {}
                        logger.debug(f"Enabling dependency '{dep}'")

        return tasks

    @typecheck
    def ensure_plugins(self, tasks: ResolvedTaskDictCollection) -> None:
        """
        Ensure all plugins used and their virtual environments exist validly.

        :param tasks: Tasks to process
        :type tasks: ResolvedTaskDictCollection
        """
        # Get logger
        logger = get_logger()

        # Get plugins used in tasks
        plugin_names = set(
            task_name.split("/")[0] for task_name in tasks.keys()
        ) - {PACKAGE_NAME}

        for plugin in plugin_names:
            # Remove plugin venv (if any) with different gurk version
            if (
                venv_exists(plugin)
                and get_venv_package_version(plugin, PACKAGE_NAME)
                != GURK_VERSION
            ):
                logger.debug(
                    f"Removing existing virtual environment for plugin '{plugin}' "
                    "to ensure it is re-created with the current gurk version."
                )
                if not remove_venv(plugin):
                    logger.fatal(
                        f"Failed to remove existing virtual environment for plugin '{plugin}'."
                    )
            # Create plugin venv
            if not venv_exists(plugin):
                if not create_plugin_venv(plugin):
                    logger.fatal(
                        f"Failed to create virtual environment for plugin '{plugin}'",
                    )
                else:
                    logger.debug(
                        f"Created virtual environment for plugin '{plugin}' successfully."
                    )

    @typecheck
    def add_preparation_task(
        self,
        tasks: ResolvedTaskDictCollection,
    ) -> ResolvedTaskDictCollection:
        """
        Prepend a system preparation task to update and upgrade apt packages.

        :param tasks: Tasks to process
        :type tasks: ResolvedTaskDictCollection
        :return: Processed tasks with preparation task added
        :rtype: ResolvedTaskDictCollection
        """
        prepare_task_name = "gurk/prepare-system"
        safe_task_name = prepare_task_name.replace("/", "-")

        # Create temporary bash script for preparation task
        tmp_path = generate_random_path(suffix=f"{safe_task_name}.bash")
        with open(tmp_path, "w") as f:
            f.write(dedent("""\
                if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
                    # (STEP) Updating apt packages...
                    sudo apt-get update -y
                fi
                """))

        # Define preparation task
        prep_task = ResolvedTaskDict(
            description="Prepare system by updating and upgrading apt packages",
            script=tmp_path,
            function=None,
            config_file=None,
            depends_on=[],
            privileged=False,
            supercedes=[],
            args=[],
        )

        # Enable prep task
        self.option[prepare_task_name] = {}

        # Prepend preparation task to all tasks' dependencies
        for _, task in tasks.items():
            task["depends_on"].insert(0, prepare_task_name)
        tasks[prepare_task_name] = prep_task

        return tasks
