"""Async token lifecycle manager — for use with httpx.AsyncClient and async Frame.io SDK."""

from __future__ import annotations

import asyncio
import inspect
import logging
import time
from typing import Any, Awaitable, Callable, Optional

from ._errors import (
    AuthenticationError,
    NetworkError,
    RateLimitError,
    TokenExpiredError,
)
from ._http import DEFAULT_REFRESH_BUFFER_SECONDS
from ._validation import _validate_token_import

logger = logging.getLogger("frameio.auth")


class AsyncTokenManager:
    """Manages token storage, expiry detection, and async auto-refresh.

    Args:
        refresh_func: An async callable that performs token refresh and returns
            a dict with at least ``access_token`` and ``expires_in``.
        refresh_buffer: Seconds before actual expiry to trigger proactive refresh.
        on_token_refreshed: Optional callback fired after every successful refresh.
    """

    def __init__(
        self,
        refresh_func: Callable[[], Awaitable[dict[str, Any]]],
        *,
        refresh_buffer: int = DEFAULT_REFRESH_BUFFER_SECONDS,
        on_token_refreshed: Optional[
            Callable[[dict[str, Any]], Awaitable[None] | None]
        ] = None,
    ) -> None:
        self._refresh_func = refresh_func
        self._refresh_buffer = refresh_buffer
        self._on_token_refreshed = on_token_refreshed

        self._access_token: str | None = None
        self._refresh_token: str | None = None
        self._expires_at: float = 0.0

        self._lock: asyncio.Lock | None = None  # lazy-init for Python 3.8/3.9
        self._refreshing = False
        self._refresh_done: Optional[asyncio.Event] = None
        self._refresh_error: BaseException | None = None  # propagated to waiters

    async def get_token(self) -> str:
        """Return a valid access token, refreshing if necessary.

        Pass this to the async Frame.io SDK::

            client = Frameio(token=auth.get_token)

        Returns:
            A valid access token string.
        """
        if self._is_token_valid():
            return self._access_token  # type: ignore[return-value]

        return await self._do_refresh()

    @property
    def access_token(self) -> str | None:
        """The current access token (may be expired)."""
        return self._access_token

    @property
    def refresh_token(self) -> str | None:
        """The current refresh token."""
        return self._refresh_token

    @property
    def expires_at(self) -> float:
        """Unix timestamp when the access token expires."""
        return self._expires_at

    @property
    def is_expired(self) -> bool:
        """Whether the access token is expired (or will be within the buffer)."""
        return not self._is_token_valid()

    def set_tokens(
        self,
        access_token: str,
        expires_in: int,
        refresh_token: str | None = None,
    ) -> None:
        """Manually set token data (e.g. after initial code exchange)."""
        self._access_token = access_token
        self._expires_at = time.time() + expires_in
        if refresh_token is not None:
            self._refresh_token = refresh_token

    def export_tokens(self) -> dict[str, Any]:
        """Export current token state for persistence."""
        return {
            "access_token": self._access_token,
            "refresh_token": self._refresh_token,
            "expires_at": self._expires_at,
        }

    def import_tokens(self, data: dict[str, Any]) -> None:
        """Restore token state from a previously exported dict."""
        _validate_token_import(data)
        self._access_token = data.get("access_token")
        self._refresh_token = data.get("refresh_token")
        self._expires_at = data.get("expires_at", 0.0)

    def clear(self) -> None:
        """Clear all stored tokens."""
        self._access_token = None
        self._refresh_token = None
        self._expires_at = 0.0

    def _get_lock(self) -> asyncio.Lock:
        """Lazy-initialize the asyncio.Lock on first use.

        This avoids binding to an event loop at __init__ time, which
        fails or binds to the wrong loop on Python <3.10.
        """
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    def _is_token_valid(self) -> bool:
        return (
            self._access_token is not None
            and time.time() < (self._expires_at - self._refresh_buffer)
        )

    async def _do_refresh(self) -> str:
        """Perform async refresh with lock coordination.

        Only one coroutine performs the refresh (the "owner"). Others wait
        on the event and either receive the refreshed token or re-raise the
        owner's error — they do not retry independently.
        """
        while True:
            wait_for: asyncio.Event | None = None
            async with self._get_lock():
                if self._is_token_valid():
                    return self._access_token  # type: ignore[return-value]

                # If a previous refresh failed, propagate its error
                if self._refresh_error is not None:
                    err = self._refresh_error
                    self._refresh_error = None
                    raise err

                if self._refreshing and self._refresh_done is not None:
                    wait_for = self._refresh_done
                elif not self._refreshing:
                    self._refreshing = True
                    self._refresh_done = asyncio.Event()
                    self._refresh_error = None
                    break  # owner exits the loop

            # Waiter path: wait outside the lock, then re-check
            if wait_for is not None:
                await wait_for.wait()
            # continue to top of while True to re-check token validity

        try:
            logger.debug("Token refresh triggered (async).")
            result = await self._refresh_func()
        except (
            AuthenticationError,
            TokenExpiredError,
            NetworkError,
            RateLimitError,
        ) as exc:
            async with self._get_lock():
                self._refresh_error = exc
                self._refreshing = False
                if self._refresh_done:
                    self._refresh_done.set()
                    self._refresh_done = None
            raise
        except Exception as exc:
            auth_err = AuthenticationError(
                "refresh_failed",
                f"Token refresh failed: {exc}",
            )
            async with self._get_lock():
                self._refresh_error = auth_err
                self._refreshing = False
                if self._refresh_done:
                    self._refresh_done.set()
                    self._refresh_done = None
            raise auth_err from exc

        async with self._get_lock():
            self._access_token = result["access_token"]
            self._expires_at = time.time() + result.get("expires_in", 86400)
            if "refresh_token" in result:
                self._refresh_token = result["refresh_token"]
            self._refresh_error = None
            self._refreshing = False
            if self._refresh_done:
                self._refresh_done.set()
                self._refresh_done = None

        logger.info(
            "Token refreshed (length=%d).",
            len(self._access_token) if self._access_token else 0,
        )

        if self._on_token_refreshed:
            try:
                callback_result = self._on_token_refreshed(self.export_tokens())
                if inspect.isawaitable(callback_result):
                    await callback_result
            except Exception:
                logger.exception("on_token_refreshed callback raised; continuing")

        assert self._access_token is not None
        return self._access_token
