from __future__ import annotations

import logging
import time

from yuxi.channel.extensions.qqbot.errors import classify_close_code

logger = logging.getLogger(__name__)

BACKOFF_SEQUENCE_MS = [1000, 2000, 5000, 10000, 30000, 60000]
MAX_RETRY_COUNT = 100
FAST_DISCONNECT_THRESHOLD_MS = 5000
FAST_DISCONNECT_COUNT = 3
FAST_DISCONNECT_PENALTY_MS = 60000


class ReconnectStateMachine:
    def __init__(self):
        self._retry_count = 0
        self._disconnect_timestamps: list[float] = []
        self._penalty_until: float = 0.0

    @property
    def retry_count(self) -> int:
        return self._retry_count

    def handle_close(self, code: int) -> tuple[bool, float | None]:
        should_retry, penalty = classify_close_code(code)

        if not should_retry:
            logger.info("Close code %d is not retryable", code)
            return (False, None)

        if penalty and penalty > 0:
            logger.info("Close code %d requires penalty: %dms", code, penalty)
            self._penalty_until = time.time() + penalty / 1000.0

        delay = self._calculate_delay()
        self._retry_count += 1
        self._record_disconnect()

        return (self._retry_count <= MAX_RETRY_COUNT, delay)

    def _calculate_delay(self) -> float | None:
        if time.time() < self._penalty_until:
            penalty_delay = self._penalty_until - time.time()
            logger.info("Penalty active, delay=%.1fs", penalty_delay)
            return penalty_delay

        if self._detect_fast_disconnect():
            logger.info("Fast disconnect detected, applying penalty")
            self._penalty_until = time.time() + FAST_DISCONNECT_PENALTY_MS / 1000.0
            return FAST_DISCONNECT_PENALTY_MS / 1000.0

        index = min(self._retry_count, len(BACKOFF_SEQUENCE_MS) - 1)
        delay_ms = BACKOFF_SEQUENCE_MS[index]
        return delay_ms / 1000.0

    def _record_disconnect(self) -> None:
        now = time.time()
        self._disconnect_timestamps.append(now)
        cutoff = now - FAST_DISCONNECT_THRESHOLD_MS / 1000.0
        self._disconnect_timestamps = [t for t in self._disconnect_timestamps if t > cutoff]

    def _detect_fast_disconnect(self) -> bool:
        return len(self._disconnect_timestamps) >= FAST_DISCONNECT_COUNT

    def reset(self) -> None:
        self._retry_count = 0
        self._disconnect_timestamps.clear()
        self._penalty_until = 0.0