from __future__ import annotations

import asyncio
import datetime
from enum import StrEnum

from yuxi.utils.logging_config import logger


class ConnectionState(StrEnum):
    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    QR_PENDING = "qr_pending"
    AUTHORIZED = "authorized"
    ACTIVE = "active"
    RECONNECTING = "reconnecting"
    LOGGED_OUT = "logged_out"


class ConnectionController:
    def __init__(self):
        self._state: ConnectionState = ConnectionState.DISCONNECTED
        self._state_changed_at: float | None = None
        self._last_transition: dict[str, str] = {}
        self._reconnect_count: int = 0
        self._max_reconnects: int = 10
        self._base_delay: float = 1.0
        self._max_delay: float = 60.0
        self._state_listeners: list = []

    @property
    def state(self) -> ConnectionState:
        return self._state

    @property
    def reconnect_count(self) -> int:
        return self._reconnect_count

    def transition(self, new_state: ConnectionState) -> bool:
        if new_state == self._state:
            return False
        old = self._state
        self._state = new_state
        self._state_changed_at = asyncio.get_event_loop().time()
        self._last_transition = {
            "from": old.value,
            "to": new_state.value,
            "at": datetime.datetime.now().isoformat(),
        }
        logger.info(f"ConnectionController: {old.value} → {new_state.value}")
        if new_state == ConnectionState.RECONNECTING:
            self._reconnect_count += 1
        elif new_state == ConnectionState.ACTIVE:
            self._reconnect_count = 0
        for listener in self._state_listeners:
            try:
                if asyncio.iscoroutinefunction(listener):
                    asyncio.ensure_future(listener(old, new_state))
                else:
                    listener(old, new_state)
            except Exception:
                logger.exception("ConnectionController listener failed")
        return True

    def on_state_change(self, listener) -> None:
        self._state_listeners.append(listener)

    def should_retry(self) -> bool:
        return self._reconnect_count < self._max_reconnects

    def next_delay(self) -> float:
        capped = min(self._reconnect_count, 10)
        delay = self._base_delay * (2**capped)
        return min(delay, self._max_delay)

    def reset_reconnects(self) -> None:
        self._reconnect_count = 0

    def is_stable(self) -> bool:
        return self._state in (
            ConnectionState.ACTIVE,
            ConnectionState.AUTHORIZED,
        )