import inspect
import threading
from datetime import timedelta
from multiprocessing import Queue
from typing import Any, Callable

from .api_client.schemas import (
    GameConfig,
    ServiceScopedFlagIds,
    TickScopedFlagIds,
    UnscopedFlagIds,
)
from .decorator.enums import FlagIdScope, TargetingStrategy
from .decorator.schemas import Batching
from .logging import colorize, logger, suggest_closest_match
from .storage.impl import BlobStorage, FlagIdsHashStorage


class Exploit:
    def __init__(
        self,
        service: str,
        is_draft: bool,
        alias: str,
        targets_skip: set[str],
        targets_include: set[str],
        targets_explicit: set[str],
        targets_strategy: TargetingStrategy | None,
        flag_id_scope: FlagIdScope,
        delay: timedelta,
        batching: Batching,
        timeout: timedelta,
        func: Callable,
    ):
        """
        Runnable configured Exploit. Contains processed user-defined configuration and a reference to the exploit
        function.
        """
        self.service = service
        self.is_draft = is_draft
        self.alias = alias
        self.targets_skip = targets_skip
        self.targets_include = targets_include
        self.targets_explicit = targets_explicit
        self.targets_strategy = targets_strategy
        self.flag_id_scope = flag_id_scope
        self.delay = delay
        self.batching = batching
        self.timeout = timeout
        self.func = func

        # Results returned from running the attacks
        self.results: Queue[tuple[str, Any]] = Queue()

        # Helper attributes
        func_params = inspect.signature(func).parameters
        self.takes_flag_ids = "flag_ids" in func_params
        self.takes_storage = "store" in func_params

        # Internal attributes
        self._batched_target_hosts: list[list[str]]
        self._service_flag_ids: ServiceScopedFlagIds | None = None
        self._blob_storage: BlobStorage | None = None
        self._flag_ids_hash_storage: FlagIdsHashStorage | None = None

    def setup(  # noqa: C901
        self,
        *,
        game: GameConfig,
        flag_ids: UnscopedFlagIds | None = None,
        blob_storage: BlobStorage | None = None,
        flag_ids_hash_storage: FlagIdsHashStorage | None = None,
        **overrides,
    ) -> bool:
        """
        Sets up the exploit by validating the passed arguments and computing the target host lists. Returns `True` if
        the setup was successful, otherwise `False`.
        """
        # Override any attributes (commonly used for reseting batching or delay in watch mode).
        for key, value in overrides.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                raise AttributeError(f"Attempted to override a non-existent attribute '{key}'.")

        # Compute list of the target hosts
        if self.takes_flag_ids:
            if flag_ids is None:
                logger.error(
                    "🚫 Failed to setup {alias} due to missing flag ids.",
                    alias=colorize(self.alias, reset="red"),
                )
                return False
            elif self.service not in flag_ids.get_service_names():
                logger.error(
                    "🚫 Failed to setup {alias} due to invalid service '<b>{service}</>'.",
                    alias=colorize(self.alias, reset="red"),
                    service=self.service,
                )

                suggestion = suggest_closest_match(self.service, flag_ids.get_service_names())
                if suggestion:
                    logger.info(
                        "❓ Did you mean <b>{suggestion}</> instead of '{original}'?",
                        suggestion=suggestion,
                        original=self.service,
                    )
                return False
            else:
                self._service_flag_ids = flag_ids / self.service
                auto_targets = self._service_flag_ids.get_target_hosts()
        else:
            auto_targets = game.opp_team_hosts

        if self.targets_strategy == TargetingStrategy.AUTO:
            target_hosts = (
                auto_targets - (game.own_team_hosts | game.nop_team_hosts | self.targets_skip) | self.targets_include
            )
        elif self.targets_strategy == TargetingStrategy.OWN_TEAM:
            target_hosts = game.own_team_hosts - self.targets_skip | self.targets_include
        elif self.targets_strategy == TargetingStrategy.NOP_TEAM:
            target_hosts = game.nop_team_hosts - self.targets_skip | self.targets_include
        else:  # Manual targets
            target_hosts = (
                self.targets_explicit - (game.own_team_hosts | game.nop_team_hosts | self.targets_skip)
                | self.targets_include
            )
            if self._service_flag_ids is not None:
                target_hosts &= self._service_flag_ids.get_target_hosts()

        # Organize target hosts in batches
        if self.batching.size:
            self._batched_target_hosts = self._batch_target_hosts_by_size(target_hosts)
        elif self.batching.count:
            self._batched_target_hosts = self._batch_target_hosts_by_count(target_hosts)

        # Set up flag IDs hash storage
        self._flag_ids_hash_storage = flag_ids_hash_storage

        # Set up blob storage
        if self.takes_storage:
            if blob_storage is None:
                logger.error(
                    "🚫 Failed to setup {alias} due to missing Redis storage.",
                    alias=colorize(self.alias, reset="red"),
                )
                return False
            else:
                self._blob_storage = blob_storage

        return True

    def run_attacks(self, batch_idx: int = 0) -> None:  # noqa: C901
        """
        Executes attacks against the hosts in the specified batch, with each attack running in its own separate thread.
        Running attacks also includes putting the results into the results queue.

        :param batch_idx: Index of the batch of target hosts, defaults to 0. Use 0 if batching is disabled (there is
        only one batch).
        :type batch_idx: int, optional
        """
        if not hasattr(self, "_batched_target_hosts"):
            logger.error(
                "🚫 Skipped {alias} due to setup failure (_batched_target_hosts missing).",
                alias=colorize(self.alias, reset="red"),
            )
            return

        target_hosts = self._batched_target_hosts[batch_idx]
        if not target_hosts:
            logger.warning(
                "⏩ Skipping {alias} due to no targets specified.",
                alias=colorize(self.alias),
            )
            return

        attack_threads = []
        skipped_attacks = []

        for host in target_hosts:
            if self.takes_flag_ids and self._service_flag_ids is not None:
                if self.flag_id_scope == FlagIdScope.SINGLE_TICK:
                    for single_tick_flag_id in (self._service_flag_ids / host).walk():
                        if (
                            self._flag_ids_hash_storage
                            and not self.is_draft
                            and single_tick_flag_id.compute_hash(self.alias) in self._flag_ids_hash_storage
                        ):
                            skipped_attacks.append(host)
                            continue
                        thread = threading.Thread(
                            target=self._run_attack,
                            args=(host, single_tick_flag_id),
                        )
                        attack_threads.append(thread)
                        thread.start()
                elif self.flag_id_scope == FlagIdScope.LAST_N_TICKS:
                    single_tick_flag_id_values = []
                    for single_tick_flag_id in (self._service_flag_ids / host).walk():
                        if (
                            self._flag_ids_hash_storage
                            and not self.is_draft
                            and single_tick_flag_id.compute_hash(self.alias) in self._flag_ids_hash_storage
                        ):
                            skipped_attacks.append(host)
                        else:
                            single_tick_flag_id_values.append(single_tick_flag_id.value)

                    thread = threading.Thread(
                        target=self._run_attack,
                        args=(host, single_tick_flag_id_values),
                    )
                    attack_threads.append(thread)
                    thread.start()
            else:
                thread = threading.Thread(
                    target=self._run_attack,
                    args=(host, None),
                )
                attack_threads.append(thread)
                thread.start()

        if skipped_attacks:
            logger.info(
                "⏩ Skipped total <b>{count}</> completed attacks via <b>{alias}</> on hosts {hosts}.",  # noqa: E501
                count=len(skipped_attacks),
                alias=colorize(self.alias),
                hosts=", ".join([colorize(host) for host in set(skipped_attacks)]),
            )

        try:
            for thread in attack_threads:
                thread.join()
        except (KeyboardInterrupt, SystemExit):
            logger.info(
                "🛑 Stopping {alias} (batch {batch_idx}/{batch_count}) due to user interruption.",
                alias=colorize(self.alias),
                batch_idx=batch_idx + 1,
                batch_count=len(self._batched_target_hosts),
            )

    def get_batch_count(self) -> int:
        """
        Returns the number of batches that will be run.

        :return: Number of batches.
        :rtype: int
        """
        return len(self._batched_target_hosts)

    def get_batch_interval(self) -> timedelta:
        """
        Returns the interval between batches.

        :return: Interval between batches.
        :rtype: timedelta
        """
        return self.batching.interval  # type: ignore[return-value]

    def _batch_target_hosts_by_size(self, hosts_set: set[str]) -> list[list[str]]:
        """
        Groups target hosts into batches of a specified size.

        :param hosts_set: Set of hosts to be distributed into batches.
        :type hosts_set: set[str]
        :return: List of batches, where each batch is a list of hosts.
        :rtype: list[list[str]]
        """
        if hosts_set is None or self.batching is None or self.batching.size is None:
            return []

        hosts_list = list(hosts_set)

        return [hosts_list[i : i + self.batching.size] for i in range(0, len(hosts_list), self.batching.size)]

    def _batch_target_hosts_by_count(self, hosts_set: set[str]) -> list[list[str]]:
        """
        Distributes target hosts into a specified number of batches as evenly as possible.

        :param hosts_set: Set of hosts to be distributed into batches.
        :type hosts_set: set[str]
        :return: List of batches, where each batch is a list of hosts.
        :rtype: list[list[str]]
        """
        if hosts_set is None or self.batching is None or self.batching.count is None:
            return []

        hosts_list = list(hosts_set)

        size = len(hosts_list) // self.batching.count
        remainder = len(hosts_list) % self.batching.count
        batches = [hosts_list[i * size : (i + 1) * size] for i in range(self.batching.count)]
        for i in range(remainder):
            batches[i].append(hosts_list[self.batching.count * size + i])

        return [batch for batch in batches if batch]

    def _is_first_batch(self, batch_idx: int):
        return batch_idx == 0

    def _is_last_batch(self, batch_idx: int):
        return batch_idx == len(self._batched_target_hosts) - 1

    def _run_attack(self, host: str, flag_ids: TickScopedFlagIds | list | None = None):
        """
        Runs a single attack against the specified target host.

        :param host: Target host to attack.
        :type host: str
        :param flag_ids: TickScopedFlagIds object if , defaults to None
        :type flag_ids: TickScopedFlagIds | list | None, optional
        """
        kwargs: dict[str, Any] = {"target": host}
        if self.takes_storage:
            kwargs["store"] = self._blob_storage
        if self.takes_flag_ids:
            kwargs["flag_ids"] = flag_ids.value if isinstance(flag_ids, TickScopedFlagIds) else flag_ids

        try:
            result = self.func(**kwargs)
        except Exception as e:
            logger.error(
                "An error has occured while attacking <b>{host}</> via <b>{alias}</>: {error}",
                host=colorize(host, reset="red"),
                alias=colorize(self.alias, reset="red"),
                error=e,
            )
        else:
            self.results.put((host, result))

            if isinstance(flag_ids, TickScopedFlagIds) and not self.is_draft and self._flag_ids_hash_storage:
                self._flag_ids_hash_storage.add(flag_ids.compute_hash(self.alias))
