from __future__ import annotations

import asyncio
import random
import time
from dataclasses import dataclass

from yuxi.utils.logging_config import logger


@dataclass
class ReconnectPolicy:
    base_delay_ms: float = 1000.0
    max_delay_ms: float = 300_000.0
    multiplier: float = 2.0
    jitter_factor: float = 0.1
    max_attempts: int = 0

    def calculate_delay(self, attempt: int) -> float:
        delay = min(self.base_delay_ms * (self.multiplier ** (attempt - 1)), self.max_delay_ms)
        jitter = delay * self.jitter_factor * random.uniform(-1, 1)
        return max(delay + jitter, 100.0)

    def should_retry(self, attempt: int) -> bool:
        return self.max_attempts <= 0 or attempt < self.max_attempts


DEFAULT_RECONNECT_POLICY = ReconnectPolicy()


@dataclass
class ReconnectState:
    attempts: int = 0
    last_attempt_at: float = 0.0
    next_delay_ms: float = 1000.0

    def record_attempt(self, policy: ReconnectPolicy) -> float:
        self.attempts += 1
        self.last_attempt_at = time.monotonic()
        self.next_delay_ms = policy.calculate_delay(self.attempts)
        return self.next_delay_ms

    def reset(self) -> None:
        self.attempts = 0
        self.last_attempt_at = 0.0
        self.next_delay_ms = 1000.0

    async def wait_and_continue(self, policy: ReconnectPolicy) -> bool:
        if not policy.should_retry(self.attempts):
            logger.error(f"Max reconnect attempts reached ({self.attempts})")
            return False
        delay_ms = self.record_attempt(policy)
        logger.info(f"Reconnect attempt {self.attempts}, waiting {delay_ms:.0f}ms")
        await asyncio.sleep(delay_ms / 1000.0)
        return True


HTTP_401_BACKOFF_POLICY = ReconnectPolicy(
    base_delay_ms=5000.0,
    max_delay_ms=600_000.0,
    multiplier=3.0,
    jitter_factor=0.2,
    max_attempts=5,
)


@dataclass
class Http401BackoffState:
    failures: int = 0
    last_failure_at: float = 0.0
    circuit_open: bool = False
    circuit_open_since: float = 0.0
    circuit_reset_after_s: float = 300.0

    def record_401(self) -> float:
        now = time.monotonic()
        if self.circuit_open:
            return max(0.0, self.circuit_open_since + self.circuit_reset_after_s - now) * 1000.0

        self.failures += 1
        self.last_failure_at = now
        if self.failures >= HTTP_401_BACKOFF_POLICY.max_attempts:
            self.circuit_open = True
            self.circuit_open_since = now
            logger.warning(
                f"[401 Backoff] Circuit breaker opened after {self.failures} failures, "
                f"will reset in {self.circuit_reset_after_s}s"
            )
            return self.circuit_reset_after_s * 1000.0

        return HTTP_401_BACKOFF_POLICY.calculate_delay(self.failures)

    def record_success(self) -> None:
        self.failures = 0
        self.circuit_open = False

    def reset(self) -> None:
        self.failures = 0
        self.last_failure_at = 0.0
        self.circuit_open = False
        self.circuit_open_since = 0.0