import logging
import time
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager, nullcontext
from enum import Enum
from typing import Any

from sdcm.sct_events.system import SoftTimeoutEvent, HardTimeoutEvent
from sdcm.utils.adaptive_timeouts.load_info_store import (
    NodeLoadInfoService,
    AdaptiveTimeoutStore,
    ArgusAdaptiveTimeoutStore,
    NodeLoadInfoServices,
)
from sdcm.utils.features import is_tablets_feature_enabled

LOGGER = logging.getLogger(__name__)


def _get_operation_timeout_factor(params, operation: "Operations") -> float:
    """Return timeout multiplier for the current operation.

    Looks up operation.value[0] in the configured multipliers dict.
    Missing config or missing operation key returns 1.0.
    """
    multipliers = params.get("adaptive_timeout_multipliers")
    if multipliers is None:
        return 1.0

    return multipliers.get_multiplier(operation.value[0])


TABLETS_SOFT_TIMEOUT = 1 * 60 * 60
TABLETS_HARD_TIMEOUT = 3 * 60 * 60
_STREAMING_OVERHEAD = 600  # add 10 minutes overhead for operations other than streaming


def _get_decommission_timeout(
    node_info_service: NodeLoadInfoService, tablets_enabled: bool = False
) -> tuple[tuple[int, int | None], dict[str, Any]]:
    """Calculate timeout for decommission operation based on node load info. Still experimental, used to gather historical data."""
    try:
        node_info = node_info_service.as_dict()
        LOGGER.debug(f"Estimating decommission timeout with {node_info=}")
        if tablets_enabled:
            node_info["tablets_enabled"] = True
            estimated = int(node_info_service.node_data_size_mb / node_info_service.expected_throughput)
            soft_timeout = estimated * 2 + _STREAMING_OVERHEAD
            hard_timeout = estimated * 4 + _STREAMING_OVERHEAD
            return (soft_timeout, hard_timeout), node_info
        else:
            # For non-tablet cases, calculate based on data size
            # rough estimation from previous runs almost 9h for 1TB
            soft_timeout = max(int(node_info_service.node_data_size_mb * 0.03), 7200)  # 2 hours minimum
            return (soft_timeout, None), node_info
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning("Failed to calculate decommission timeout: \n%s \nDefaulting to 6 hours", exc)
        return (6 * 60 * 60, None), {}


def _get_new_node_timeout(
    node_info_service: NodeLoadInfoService,
    timeout: int | float = None,
    tablets_enabled: bool = False,
) -> tuple[tuple[int, int | None], dict[str, Any]]:
    """Calculate timeout for adding a new node operation"""
    try:
        node_info = node_info_service.as_dict()
        if tablets_enabled:
            node_info["tablets_enabled"] = True
            return (TABLETS_SOFT_TIMEOUT, TABLETS_HARD_TIMEOUT), node_info
        # For non-tablet cases, use the passed timeout
        return (timeout, None), node_info
    except Exception as exc:  # pylint: disable=broad-except  # noqa: BLE001
        LOGGER.warning("Failed to get node info for timeout: \n%s", exc)
        return (timeout, None), {}


def _get_tablet_migration_timeout(
    node_info_service: NodeLoadInfoService,
    timeout: int | float = None,
    tablets_enabled: bool = False,
) -> tuple[tuple[int, int | None], dict[str, Any]]:
    """Calculate timeout for tablet migration operation based on node disk size.

    At worst, the node will be filled up to 90% of its disk size
    Uses node_disk_size_mb * 0.9 / expected_throughput as the base estimate,
    with soft = estimated * 2 and hard = estimated * 4.
    For vnodes, use the same estimation as decommission.
    Falls back to the caller-supplied timeout on error.
    """
    try:
        node_info = node_info_service.as_dict()
        LOGGER.debug(f"Estimating tablet migration timeout with {node_info=}")
        if tablets_enabled:
            node_info["tablets_enabled"] = True
            estimated = int(node_info_service.node_disk_size_mb * 0.9 / node_info_service.expected_throughput)
            soft_timeout = estimated * 2 + _STREAMING_OVERHEAD
            hard_timeout = estimated * 4 + _STREAMING_OVERHEAD
            return (soft_timeout, hard_timeout), node_info
        else:
            soft_timeout = max(int(node_info_service.node_disk_size_mb * 0.9 * 0.03), 7200)  # 2 hours minimum
            return (soft_timeout, None), node_info
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning(
            "Failed to calculate tablet migration timeout: \n%s \nDefaulting to caller timeout or 6 hours", exc
        )
        return (timeout or 6 * 60 * 60, None), {}


def _get_soft_timeout(
    node_info_service: NodeLoadInfoService, timeout: int | float = None
) -> tuple[int | float, dict[str, Any]]:
    # no timeout calculation - just return the timeout passed as argument along with node load info
    try:
        return timeout, node_info_service.as_dict()
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning("Failed to get node info for timeout: \n%s", exc)
        return timeout, {}


def _get_soft_timeout_no_node_info(
    node_info_service: NodeLoadInfoService, timeout: int | float = None
) -> tuple[int | float, dict[str, Any]]:
    # no timeout calculation - just return the timeout passed as argument without node load info
    return timeout, {}


def _get_query_timeout(
    node_info_service: NodeLoadInfoService, timeout: int | float = None, query: str = None
) -> tuple[int | float, dict[str, Any]]:
    timeout, stats = _get_soft_timeout(node_info_service=node_info_service, timeout=timeout)
    stats["query"] = query
    return timeout, stats


def _get_service_level_propagation_timeout(
    node_info_service: NodeLoadInfoService, timeout: int | float = None, service_level_for_test_step: str = None
) -> tuple[int | float, dict[str, Any]]:
    """service_level_for_test_step will report in ES on which step of test the timeout happened"""
    timeout, stats = _get_soft_timeout(node_info_service=node_info_service, timeout=timeout)
    stats["service_level_for_test_step"] = service_level_for_test_step
    return timeout, stats


def _get_compaction_timeout(
    node_info_service: NodeLoadInfoService, timeout: int | float = None
) -> tuple[int | float, dict[str, Any]]:
    """
    Experimentally got these timings
    Using the estimate = (MB / shards) / 25 + 120
    we ensure we always overshoot the actual duration

    Shards | MB     | Time     | Actual(s) | Est(s) | Diff(s) | Diff(%)
    -------+--------+----------+-----------+--------+---------+--------
    2      | 53084  | 00:18:05 |      1085 |   1182 |      97 |   8.9%
    4      | 51814  | 00:08:47 |       527 |    638 |     111 |  21.1%
    7      | 50698  | 00:04:25 |       265 |    410 |     145 |  54.7%

    2      | 217528 | 00:59:36 |      3576 |   4471 |     895 |  25.0%
    4      | 204093 | 00:30:53 |      1853 |   2161 |     308 |  16.6%
    7      | 187566 | 00:14:27 |       867 |   1192 |     325 |  37.5%

    2      | 340346 | 01:43:14 |      6194 |   6927 |     733 |  11.8%
    4      | 335093 | 00:49:18 |      2958 |   3471 |     513 |  17.3%
    7      | 272496 | 00:21:07 |      1267 |   1677 |     410 |  32.4%
    """
    COMPACTION_SPEED = 25  # MB/s per shard
    FIXED_OVERHEAD = 120  # seconds
    SAFETY_FACTOR = 2  # timeout is doubled to account for load, instance type, etc. variations
    try:
        data_size = node_info_service.node_data_size_mb  # MB
        shards = node_info_service.shards_count
        if data_size and shards:
            estimated = int((data_size / shards) / COMPACTION_SPEED) + FIXED_OVERHEAD
            timeout = estimated * SAFETY_FACTOR
        return timeout, node_info_service.as_dict()
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning("Failed to get node info for timeout: \n%s", exc)
        return timeout, {}


def _get_cleanup_timeout(node_info_service, timeout=None):
    """Calculate cleanup timeout based on node data size.
    After decommission/topology change, cleanup removes data that no longer
    belongs to a node. The amount of data to process is proportional to total
    node data size divided by the number of nodes (since only the reshuffled
    portion needs cleanup).

    Ref: https://scylladb.atlassian.net/browse/SCT-247

    Formula:
        timeout = max(MIN_CLEANUP_TIMEOUT, data_size_mb / throughput_mb_s * safety_factor)
    Where:
        - data_size_mb: total data on the node (from node_info_service)
        - throughput_mb_s: expected I/O throughput for the instance type
        - safety_factor: 2.0 (accounts for compaction contention, I/O variability)
        - MIN_CLEANUP_TIMEOUT: 120 seconds (floor for small datasets)
    """
    MIN_CLEANUP_TIMEOUT = 120  # seconds
    SAFETY_FACTOR = 2.0
    try:
        data_size_mb = node_info_service.node_data_size_mb
        throughput = node_info_service.expected_throughput  # MB/s
        if data_size_mb and throughput:
            calculated = (data_size_mb / throughput) * SAFETY_FACTOR
            timeout = max(MIN_CLEANUP_TIMEOUT, calculated)
        else:
            timeout = timeout or MIN_CLEANUP_TIMEOUT
        return timeout, node_info_service.as_dict()
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning("Failed to calculate cleanup timeout: \n%s \nDefaulting to %s seconds", exc, MIN_CLEANUP_TIMEOUT)
        return timeout or MIN_CLEANUP_TIMEOUT, {}


class Operations(Enum):
    """Available operations for adaptive timeouts. Each operation maps to a function to calculate timeout based on node load info.

    maps to (function, (required arguments))"""

    DECOMMISSION = ("decommission", _get_decommission_timeout, ())
    NEW_NODE = ("new_node", _get_new_node_timeout, ("timeout",))
    CREATE_INDEX = ("create_index", _get_soft_timeout, ("timeout",))
    START_SCYLLA = ("start_scylla", _get_soft_timeout, ("timeout",))
    RESTART_SCYLLA = ("restart_scylla", _get_soft_timeout, ("timeout",))
    CREATE_MV = ("create_mv", _get_soft_timeout, ("timeout",))
    MAJOR_COMPACT = ("major_compact", _get_compaction_timeout, ("timeout",))
    REPAIR = ("repair", _get_soft_timeout, ("timeout",))
    MGMT_REPAIR = ("mgmt_repair", _get_soft_timeout, ("timeout",))
    BACKUP = ("backup", _get_soft_timeout, ("timeout",))
    RESTORE = ("restore", _get_soft_timeout, ("timeout",))
    REBUILD = ("rebuild", _get_soft_timeout, ("timeout",))
    CLEANUP = ("cleanup", _get_cleanup_timeout, ("timeout",))
    REMOVE_NODE = ("remove_node", _get_soft_timeout, ("timeout",))
    SCRUB = ("scrub", _get_soft_timeout, ("timeout",))
    SOFT_TIMEOUT = ("soft_timeout", _get_soft_timeout, ("timeout",))
    FLUSH = ("flush", _get_soft_timeout, ("timeout",))
    QUERY = ("query", _get_query_timeout, ("timeout", "query"))
    SERVICE_LEVEL_PROPAGATION = (
        "service_level_propagation",
        _get_service_level_propagation_timeout,
        ("timeout", "service_level_for_test_step"),
    )
    TABLET_MIGRATION = ("tablet_migration", _get_tablet_migration_timeout, ("timeout",))
    HEALTHCHECK = ("healthcheck", _get_soft_timeout, ("timeout",))
    SSH_CONNECTIVITY = ("ssh_connectivity", _get_soft_timeout_no_node_info, ("timeout",))


class TestInfoServices:
    @staticmethod
    def get(node: "BaseNode") -> dict:  # noqa: F821
        return dict(
            n_db_nodes=len(node.parent_cluster.nodes),
        )


class TimeoutMonitor:
    """
    Monitors an operation and publishes error or critical event, when soft or hard timeouts is reached, correspondingly
    """

    def __init__(
        self,
        operation: str,
        soft_timeout: int | float,
        hard_timeout: int | float = None,
        start_time: float | None = None,
    ):
        self.operation = operation
        self.soft_timeout = soft_timeout
        self.hard_timeout = hard_timeout
        self.start_time = start_time or time.monotonic()
        self.soft_timeout_triggered = False
        self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="timeout_monitor")
        self.running = False

    def __enter__(self):
        if self.soft_timeout or self.hard_timeout:
            self.running = True
            self.executor.submit(self._monitor_timeouts)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.running = False
        try:
            self.executor.shutdown(wait=False, cancel_futures=True)
        except Exception as exc:
            LOGGER.exception("Error shutting down timeout monitor executor for %s: %s", self.operation, exc)

    def _monitor_timeouts(self):
        while self.running:
            duration = time.monotonic() - self.start_time

            if self.hard_timeout and duration > self.hard_timeout:
                HardTimeoutEvent(
                    operation=self.operation, duration=duration, hard_timeout=self.hard_timeout
                ).publish_or_dump()
                self.running = False
                return
            if not self.soft_timeout_triggered and self.soft_timeout and duration > self.soft_timeout:
                SoftTimeoutEvent(
                    operation=self.operation, duration=duration, soft_timeout=self.soft_timeout, still_running=True
                ).publish_or_dump()
                self.soft_timeout_triggered = True

            time.sleep(5.0)


@contextmanager
def adaptive_timeout(  # noqa: PLR0914
    operation: Operations,
    node: "BaseNode",  # noqa: F821
    stats_storage: AdaptiveTimeoutStore = ArgusAdaptiveTimeoutStore(),
    **kwargs,
):
    """
    Calculate timeout in seconds for given operation based on node load info and return its value.
    Upon exit, verify if timeout occurred and publish SoftTimeoutEvent if happened.
    Also store node load info and timeout value in AdaptiveTimeoutStore (ES by default) for future reference.
    Use Operation.SOFT_TIMEOUT to set timeout explicitly without calculations.
    """
    tablet_sensitive_op = operation in {Operations.DECOMMISSION, Operations.NEW_NODE, Operations.TABLET_MIGRATION}
    kwargs.setdefault("node_available", True)

    # in some situations we may want to skip tablet check, since it depends on ssh connectivity
    tablets_enabled = kwargs["node_available"] and is_tablets_feature_enabled(node)

    _, timeout_func, required_arg_names = operation.value
    args = {arg: kwargs[arg] for arg in required_arg_names}

    # if node is known to be not available, skip metrics gathering and use default timeouts
    store_metrics = node.parent_cluster.params.get("adaptive_timeout_store_metrics") and kwargs["node_available"]
    if store_metrics:
        metrics = NodeLoadInfoServices().get(node)
    else:
        metrics = {}
    if tablet_sensitive_op:
        args["tablets_enabled"] = tablets_enabled
    result = timeout_func(node_info_service=metrics, **args)
    if tablet_sensitive_op:
        (soft_timeout, hard_timeout), load_metrics = result
    else:
        soft_timeout, load_metrics = result
        hard_timeout = None
    if store_metrics:
        load_metrics.update(TestInfoServices.get(node))

    # Apply operation-based timeout scaling (opt-in via config)
    operation_factor = _get_operation_timeout_factor(node.parent_cluster.params, operation)
    soft_timeout = soft_timeout * operation_factor if soft_timeout else None
    hard_timeout = hard_timeout * operation_factor if hard_timeout else None

    start_time = time.monotonic()
    timeout_occurred = False
    monitor_ctx = (
        TimeoutMonitor(operation.name, soft_timeout, hard_timeout, start_time=start_time)
        if tablet_sensitive_op and tablets_enabled and hard_timeout
        else nullcontext()
    )

    try:
        with monitor_ctx:
            yield hard_timeout or soft_timeout
    except Exception as exc:
        exc_name = exc.__class__.__name__
        if "timeout" in exc_name.lower() or "timed out" in str(exc):
            timeout_occurred = True
        raise
    finally:
        duration = time.monotonic() - start_time
        soft_timeout_exceeded = soft_timeout and duration > soft_timeout
        if soft_timeout_exceeded:
            timeout_occurred = True
        if timeout_occurred:
            SoftTimeoutEvent(operation=operation.name, soft_timeout=soft_timeout, duration=duration).publish_or_dump()

        try:
            if load_metrics and store_metrics:
                stats_storage.store(
                    metrics=load_metrics,
                    operation=operation.name,
                    duration=duration,
                    timeout=soft_timeout,
                    timeout_occurred=timeout_occurred,
                )
        except Exception as exc:  # noqa: BLE001
            LOGGER.warning("Failed to store adaptive timeout stats: \n%s", exc, exc_info=True)
