import copy
from typing import Optional

import frostfs_testlib.resources.optionals as optionals
from frostfs_testlib.load.interfaces import ScenarioRunner
from frostfs_testlib.load.load_config import (
    EndpointSelectionStrategy,
    LoadParams,
    LoadScenario,
    LoadType,
)
from frostfs_testlib.load.load_report import LoadReport
from frostfs_testlib.load.load_verifiers import LoadVerifier
from frostfs_testlib.reporter import get_reporter
from frostfs_testlib.storage.cluster import ClusterNode
from frostfs_testlib.storage.dataclasses.frostfs_services import S3Gate, StorageNode
from frostfs_testlib.storage.dataclasses.wallet import WalletInfo
from frostfs_testlib.testing.test_control import run_optionally

reporter = get_reporter()


class BackgroundLoadController:
    k6_dir: str
    load_params: LoadParams
    original_load_params: LoadParams
    verification_params: LoadParams
    nodes_under_load: list[ClusterNode]
    load_counter: int
    loaders_wallet: WalletInfo
    load_summaries: dict
    endpoints: list[str]
    runner: ScenarioRunner
    started: bool

    def __init__(
        self,
        k6_dir: str,
        load_params: LoadParams,
        loaders_wallet: WalletInfo,
        nodes_under_load: list[ClusterNode],
        runner: ScenarioRunner,
    ) -> None:
        self.k6_dir = k6_dir
        self.original_load_params = load_params
        self.load_params = copy.deepcopy(self.original_load_params)
        self.nodes_under_load = nodes_under_load
        self.load_counter = 1
        self.loaders_wallet = loaders_wallet
        self.runner = runner
        self.started = False
        if load_params.endpoint_selection_strategy is None:
            raise RuntimeError("endpoint_selection_strategy should not be None")

        self.endpoints = self._get_endpoints(
            load_params.load_type, load_params.endpoint_selection_strategy
        )

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED, [])
    def _get_endpoints(
        self, load_type: LoadType, endpoint_selection_strategy: EndpointSelectionStrategy
    ):
        all_endpoints = {
            LoadType.gRPC: {
                EndpointSelectionStrategy.ALL: list(
                    set(
                        endpoint
                        for node_under_load in self.nodes_under_load
                        for endpoint in node_under_load.service(StorageNode).get_all_rpc_endpoint()
                    )
                ),
                EndpointSelectionStrategy.FIRST: list(
                    set(
                        node_under_load.service(StorageNode).get_rpc_endpoint()
                        for node_under_load in self.nodes_under_load
                    )
                ),
            },
            # for some reason xk6 appends http protocol on its own
            LoadType.S3: {
                EndpointSelectionStrategy.ALL: list(
                    set(
                        endpoint.replace("http://", "")
                        for node_under_load in self.nodes_under_load
                        for endpoint in node_under_load.service(S3Gate).get_all_endpoints()
                    )
                ),
                EndpointSelectionStrategy.FIRST: list(
                    set(
                        node_under_load.service(S3Gate).get_endpoint().replace("http://", "")
                        for node_under_load in self.nodes_under_load
                    )
                ),
            },
        }

        return all_endpoints[load_type][endpoint_selection_strategy]

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Prepare load instances")
    def prepare(self):
        self.runner.prepare(self.load_params, self.nodes_under_load, self.k6_dir)
        self.runner.init_k6_instances(self.load_params, self.endpoints, self.k6_dir)

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    def start(self):
        with reporter.step(f"Start load on nodes {self.nodes_under_load}"):
            self.runner.start()
        self.started = True

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Stop load")
    def stop(self):
        self.runner.stop()

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED, True)
    def is_running(self) -> bool:
        return self.runner.is_running

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Reset load")
    def _reset_for_consequent_load(self):
        """This method is required if we want to run multiple loads during test run.
        Raise load counter by 1 and append it to load_id
        """
        self.load_counter += 1
        self.load_params = copy.deepcopy(self.original_load_params)
        self.load_params.set_id(f"{self.load_params.load_id}_{self.load_counter}")

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Startup load")
    def startup(self):
        self.prepare()
        self.start()

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Stop and get results of load")
    def teardown(self, load_report: Optional[LoadReport] = None):
        if not self.started:
            return

        self.stop()
        self.load_summaries = self._get_results()
        self.started = False
        if load_report:
            load_report.add_summaries(self.load_summaries)

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Verify results of load")
    def verify(self):
        try:
            if self.load_params.verify:
                self.verification_params = LoadParams(
                    verify_clients=self.load_params.verify_clients,
                    scenario=LoadScenario.VERIFY,
                    registry_file=self.load_params.registry_file,
                    verify_time=self.load_params.verify_time,
                    load_type=self.load_params.load_type,
                    load_id=self.load_params.load_id,
                    working_dir=self.load_params.working_dir,
                    endpoint_selection_strategy=self.load_params.endpoint_selection_strategy,
                    k6_process_allocation_strategy=self.load_params.k6_process_allocation_strategy,
                    setup_timeout="1s",
                )
                self._run_verify_scenario()
                verification_summaries = self._get_results()
                self.verify_summaries(self.load_summaries, verification_summaries)
        finally:
            self._reset_for_consequent_load()

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Verify summaries from k6")
    def verify_summaries(self, load_summaries: dict, verification_summaries: dict):
        verifier = LoadVerifier(self.load_params)
        for node_or_endpoint in load_summaries:
            with reporter.step(f"Verify load summaries for {node_or_endpoint}"):
                verifier.verify_summaries(
                    load_summaries[node_or_endpoint], verification_summaries[node_or_endpoint]
                )

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    def wait_until_finish(self):
        self.runner.wait_until_finish()

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Run verify scenario")
    def _run_verify_scenario(self):
        if self.verification_params.verify_time is None:
            raise RuntimeError("verify_time should not be none")

        self.runner.init_k6_instances(self.verification_params, self.endpoints, self.k6_dir)
        with reporter.step("Run verify load data"):
            self.runner.start()
            self.runner.wait_until_finish()

    @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED)
    @reporter.step_deco("Get load results")
    def _get_results(self) -> dict:
        return self.runner.get_results()