from typing import TYPE_CHECKING, Literal, Optional, TypeVar, cast

import fmf.utils

import tmt.log
import tmt.steps
import tmt.utils
from tmt.container import container, simple_field
from tmt.guest import Guest
from tmt.plugins import PluginRegistry
from tmt.result import PhaseResult, ResultGuestData, ResultOutcome
from tmt.steps import (
    Action,
    PluginOutcome,
    PluginTask,
    PullTask,
    PushTask,
    sync_with_guests,
)
from tmt.utils import uniq

if TYPE_CHECKING:
    import tmt.base.core
    import tmt.guest
    from tmt.base.plan import Plan


@container
class PrepareStepData(tmt.steps.WhereableStepData, tmt.steps.StepData):
    pass


PrepareStepDataT = TypeVar('PrepareStepDataT', bound=PrepareStepData)


class _RawPrepareStepData(tmt.steps._RawStepData, tmt.steps.RawWhereableStepData, total=False):
    pass


class PreparePlugin(tmt.steps.Plugin[PrepareStepDataT, PluginOutcome]):
    """
    Common parent of prepare plugins
    """

    # ignore[assignment]: as a base class, PrepareStepData is not included in
    # PrepareStepDataT.
    _data_class = PrepareStepData  # type: ignore[assignment]

    # Methods ("how: ..." implementations) registered for the same step.
    _supported_methods: PluginRegistry[tmt.steps.Method] = PluginRegistry('step.prepare')

    def go(
        self,
        *,
        guest: 'tmt.guest.Guest',
        environment: Optional[tmt.utils.Environment] = None,
        logger: tmt.log.Logger,
    ) -> PluginOutcome:
        """
        Prepare the guest (common actions)
        """

        self.go_prolog(logger)

        # Show guest name first in multihost scenarios
        if self.step.plan.provision.is_multihost:
            logger.info('guest', guest.name, 'green')

        # Show requested role if defined
        if self.data.where:
            logger.info('where', fmf.utils.listed(self.data.where), 'green')

        return PluginOutcome()


# Required & recommended packages
#
# Structures and code for collecting requirements for different guests
# or their groups to avoid installing all test requirements on all
# guests. For example, a test running on the "server" guest might
# require package `foo` while the test running on the "client" might
# require package `bar`, and `foo` and `bar` cannot be installed at the
# same time.


@container
class DependencyCollection:
    """
    Bundle guests and packages to install on them
    """

    # Guest*s*, not a guest. The list will start with just one guest at
    # first, but when grouping guests by same requirements, we'd start
    # adding guests to the list when spotting same set of dependencies.
    guests: list[Guest]
    dependencies: list['tmt.base.core.DependencySimple'] = simple_field(default_factory=list)

    @property
    def as_key(self) -> 'DependencyCollectionKey':
        return frozenset(self.dependencies)


DependencyCollectionKey = frozenset['tmt.base.core.DependencySimple']


class Prepare(tmt.steps.StepWithQueue[PrepareStepData, PluginOutcome]):
    """
    Prepare the environment for testing.

    Use the 'order' attribute to select in which order preparation
    should happen if there are multiple configs. Default order is '50'.
    Default order of required packages installation is '70', for the
    recommended packages it is '75'.
    """

    _plugin_base_class = PreparePlugin

    results: list[PhaseResult]

    @property
    def _preserved_workdir_members(self) -> set[str]:
        """
        A set of members of the step workdir that should not be removed.
        """

        members = {
            *super()._preserved_workdir_members,
        }

        if self.plan.my_run:
            members = {*members, f'results{self.plan.my_run.state_format.suffix}'}

        return members

    def __init__(
        self,
        *,
        plan: 'Plan',
        raw_data: list[tmt.steps._RawStepData],
        logger: tmt.log.Logger,
    ) -> None:
        """
        Initialize prepare step data
        """

        super().__init__(plan=plan, raw_data=raw_data, logger=logger)

        self.results = []
        self.preparations_applied = 0

    def load(self) -> None:
        super().load()

        self.results = self._load_results(PhaseResult, allow_missing=True)

    def save(self) -> None:
        super().save()

        self._save_results(self.results)

    def wake(self) -> None:
        """
        Wake up the step (process workdir and command line)
        """

        super().wake()

        # Choose the right plugin and wake it up
        for data in self.data:
            # FIXME: cast() - see https://github.com/teemtee/tmt/issues/1599
            plugin = cast(PreparePlugin[PrepareStepData], PreparePlugin.delegate(self, data=data))
            plugin.wake()
            # Add plugin only if there are data
            if not plugin.data.is_bare:
                self._phases.append(plugin)

        # Nothing more to do if already done and not asked to run again
        if self.status() == 'done' and not self.should_run_again:
            self.debug('Prepare wake up complete (already done before).', level=2)
        # Save status and step data (now we know what to do)
        else:
            self.status('todo')
            self.save()

    def summary(self) -> None:
        """
        Give a concise summary of the preparation
        """

        preparations = fmf.utils.listed(self.preparations_applied, 'preparation')
        self.info('summary', f'{preparations} applied', 'green', shift=1)

    def go(self, force: bool = False) -> None:
        """
        Prepare the guests
        """

        super().go(force=force)

        # Nothing more to do if already done
        if self.status() == 'done':
            self.info('status', 'done', 'green', shift=1)
            self.summary()
            self.actions()
            return

        import tmt.base.core

        # All phases from all steps.
        phases = [
            phase
            for step in (
                self.plan.discover,
                self.plan.provision,
                self.plan.prepare,
                self.plan.execute,
                self.plan.finish,
                self.plan.report,
            )
            for phase in step.phases(classes=step._plugin_base_class)
        ]

        # All provisioned guests.
        guests = self.plan.provision.ready_guests

        # 1. collect all requirements, per guest. For each phase, test,
        # check and so on, find out on which guest it needs to run, and
        # add its requirements to a guest-specific collection.

        # Collecting all essential requirements, per guest.
        collected_essential_requires = {
            guest: DependencyCollection(guests=[guest]) for guest in guests
        }

        # Collecting all required packages, per guest.
        collected_requires = {guest: DependencyCollection(guests=[guest]) for guest in guests}

        # Collecting all recommended packages, per guest.
        collected_recommends = {guest: DependencyCollection(guests=[guest]) for guest in guests}

        # For each guest, check everything that can run, and if enabled
        # for the given guest, add requirements into the correct
        # collection.
        for guest in guests:
            # First, check phases - plugins have their own requirements,
            # the essential requirements.
            for phase in phases:
                if not phase.enabled_by_when:
                    continue
                if not phase.enabled_on_guest(guest):
                    continue

                collected_essential_requires[
                    guest
                ].dependencies += tmt.base.core.assert_simple_dependencies(
                    # ignore[attr-defined]: mypy thinks that phase is Phase type, while its
                    # actually PluginClass
                    phase.essential_requires(),  # type: ignore[attr-defined]
                    'After beakerlib processing, tests may have only simple requirements',
                    self._logger,
                )

            # The `discover` step is different: no phases, just query tests
            # collected by the step itself. Maybe we could iterate over
            # `discover` phases, but I think re-runs and workdir reuse would
            # use what the step loads from its storage, `tests.yaml`. Which
            # means, there probably would be no phases to inspect from time to
            # time, therefore going after the step itself.
            for test_origin in self.plan.discover.tests(enabled=True):
                test = test_origin.test

                if not test.enabled_on_guest(guest):
                    continue

                # Collect required packages (require, framework and checks)

                collected_requires[guest].dependencies += tmt.base.core.assert_simple_dependencies(
                    test.require,
                    'After beakerlib processing, tests may have only simple requirements',
                    self._logger,
                )

                collected_requires[guest].dependencies += test.test_framework.get_requirements(
                    test, self._logger
                )

                for check in test.check:
                    collected_requires[guest].dependencies += check.plugin.essential_requires(
                        guest, test, self._logger
                    )

                # Collect recommended packages

                collected_recommends[
                    guest
                ].dependencies += tmt.base.core.assert_simple_dependencies(
                    test.recommend,
                    'After beakerlib processing, tests may have only simple requirements',
                    self._logger,
                )

        # 2. Now we have guests and all their requirements. There can be
        # duplicities, multiple tests requesting the same package, but also
        # some guests may share the set of packages to be installed on them.
        # Let's say N guests share a `role`, all their tests would add the same
        # requirements to these guests.
        #
        # So the final 2 steps:
        #
        # 1. make the list of requirements unique,
        # 2. group guests with same requirements.
        def _prune_collections(
            collections: dict[Guest, DependencyCollection],
        ) -> list[DependencyCollection]:
            pruned: dict[DependencyCollectionKey, DependencyCollection] = {}

            for guest, collection in collections.items():
                collection.dependencies = uniq(collection.dependencies)

                if collection.as_key in pruned:
                    pruned[collection.as_key].guests.append(guest)

                else:
                    pruned[collection.as_key] = collection

            return list(pruned.values())

        pruned_essential_requires = _prune_collections(collected_essential_requires)
        pruned_requires = _prune_collections(collected_requires)
        pruned_recommends = _prune_collections(collected_recommends)

        # 3. for each collection, which now groups a set of packages and
        # all guests they need to be installed on, add new phase that
        # would take care of installation.
        def _emit_phase(
            pruned_collections: list[DependencyCollection],
            name: str,
            summary: str,
            order: int,
            missing: Literal['skip', 'fail'] = 'fail',
        ) -> None:
            from tmt.steps.prepare.install import PrepareInstallData

            for collection in pruned_collections:
                if not collection.dependencies:
                    continue

                data = PrepareInstallData(
                    name=name,
                    how='install',
                    summary=summary,
                    order=order,
                    where=[guest.name for guest in collection.guests],
                    package=collection.dependencies,
                    missing=missing,
                )

                self._phases.append(PreparePlugin.delegate(self, data=data))

        _emit_phase(
            pruned_essential_requires,
            'essential-requires',
            'Install essential required packages',
            tmt.steps.PHASE_ORDER_PREPARE_INSTALL_ESSENTIAL_REQUIRES,
        )

        _emit_phase(
            pruned_requires,
            'requires',
            'Install required packages',
            tmt.steps.PHASE_ORDER_PREPARE_INSTALL_REQUIRES,
        )

        _emit_phase(
            pruned_recommends,
            'recommends',
            'Install recommended packages',
            tmt.steps.PHASE_ORDER_PREPARE_INSTALL_RECOMMENDS,
            missing='skip',
        )

        if self._steppified_ready_guests:
            sync_with_guests(
                self, 'push', PushTask(self._steppified_ready_guests, self._logger), self._logger
            )

            # To separate "push" from "prepare" queue visually
            self.info('')

        self._queue.reset()

        for prepare_phase in self.phases(classes=(Action, PreparePlugin)):
            if isinstance(prepare_phase, Action):
                self._queue.enqueue_action(phase=prepare_phase)

            elif prepare_phase.enabled_by_when:
                self._queue.enqueue_plugin(
                    phase=prepare_phase,  # type: ignore[arg-type]
                    guests=[
                        guest
                        for guest in self._steppified_ready_guests
                        if prepare_phase.enabled_on_guest(guest)
                    ],
                )

        self.results: list[PhaseResult] = []
        exceptions: list[Exception] = []

        def _record_exception(
            outcome: PluginTask[PrepareStepData, PluginOutcome], exc: Exception
        ) -> None:
            outcome.logger.fail(str(exc))

            exceptions.append(exc)

        def _is_failed() -> bool:
            return bool(exceptions) or any(
                result.result in (ResultOutcome.ERROR, ResultOutcome.FAIL)
                for result in self.results
            )

        for outcome in self._queue.run():
            if not isinstance(outcome.phase, PreparePlugin):
                continue

            # Possible outcomes: plugin crashed, raised an exception,
            # and that exception has been delivered to the top of the
            # phase's thread and propagated to us in the task outcome.
            #
            # Log the failure, save the exception, and add an error
            # result to represent the crash. Plugin did not return any
            # usable results, otherwise it would not have ended with
            # an exception...
            if outcome.exc:
                assert outcome.guest is not None  # narrow type

                _record_exception(outcome, outcome.exc)

                self.results.append(
                    PhaseResult(
                        name=outcome.phase.name,
                        result=ResultOutcome.ERROR,
                        note=['Plugin raised an unhandled exception.'],
                        guest=ResultGuestData.from_guest(guest=outcome.guest),
                    )
                )

                # Do not let the queue continue.
                self._queue.stop()

                break

            # Or, plugin finished successfully - not necessarily after
            # achieving its goals successfully. Save results, and if
            # plugin returned also some exceptions, do the same as above:
            # log them and save them, but do not emit any special result.
            # Plugin was alive till the very end, and returned results.
            if outcome.result:
                self.results += outcome.result.results

                if outcome.result.exceptions:
                    for exc in outcome.result.exceptions:
                        _record_exception(outcome, exc)

            if _is_failed():
                self._queue.stop()

                break

            self.preparations_applied += 1

        self._save_results(self.results)

        if _is_failed():
            # TODO: needs a better message...
            raise tmt.utils.PrepareError(
                'prepare step failed',
                causes=exceptions,
            )

        # Notify guests that the prepare step has completed.
        for guest in self._steppified_ready_guests:
            guest.on_step_complete(self)

        self.info('')

        # Pull artifacts created in the plan data directory
        # if there was at least one plugin executed
        if self.phases() and self._steppified_ready_guests:
            sync_with_guests(
                self,
                'pull',
                PullTask(self._steppified_ready_guests, self.plan.data_directory, self._logger),
                self._logger,
            )

            # To separate "prepare" from "pull" queue visually
            self.info('')

        # Give a summary, update status and save
        self.summary()
        self.status('done')
        self.save()


# Establish the "plugin class -> step class" link.
PreparePlugin._step_class = Prepare