from __future__ import annotations

import logging
import os
import threading
import time
from collections.abc import Sequence
from pathlib import Path

from azure.core.credentials import TokenCredential
from azure.identity import DefaultAzureCredential

from prompt_auto_lab.core.store_dir import get_store_dir

logger = logging.getLogger(__name__)


def _parse_int_env(name: str, default: int) -> int:
    value = os.getenv(name)
    if value is None:
        return default
    try:
        return int(value)
    except ValueError:
        logger.warning("Invalid %s value %r; using %s", name, value, default)
        return default


def _parse_bool_env(name: str, default: bool) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in ("1", "true", "yes", "on", "y")


_DEFAULT_REFRESH_MARGIN_S = _parse_int_env("AZURE_TOKEN_REFRESH_MARGIN_S", 120)
_DEFAULT_REFRESH_INTERVAL_S = _parse_int_env("AZURE_TOKEN_REFRESH_INTERVAL_S", 300)
_REFRESH_LOG_ENABLED = _parse_bool_env("AZURE_TOKEN_REFRESH_LOG", False)


def _resolve_refresh_log_path() -> Path:
    value = os.getenv("AZURE_TOKEN_REFRESH_LOG_PATH")
    if value:
        return Path(value)
    return get_store_dir() / "auth_token.log"


def _ensure_refresh_logger() -> None:
    if not _REFRESH_LOG_ENABLED:
        return
    for handler in logger.handlers:
        if getattr(handler, "_autopromptlab_auth_token", False):
            return
    log_path = _resolve_refresh_log_path()
    log_path.parent.mkdir(parents=True, exist_ok=True)
    handler = logging.FileHandler(log_path, encoding="utf-8")
    handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s"))
    handler.setLevel(logging.DEBUG)
    handler._autopromptlab_auth_token = True
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)


class AzureAuthTokenProvider:
    def __init__(
        self,
        *,
        scopes: Sequence[str] | str,
        credential: TokenCredential | None = None,
        refresh_margin_s: int | None = None,
        refresh_interval_s: int | None = None,
        auto_refresh: bool = True,
    ) -> None:
        if isinstance(scopes, str):
            scopes = (scopes,)
        self._scopes = tuple(scopes)
        self._credential = credential or DefaultAzureCredential()
        self._refresh_margin_s = max(
            0,
            int(_DEFAULT_REFRESH_MARGIN_S if refresh_margin_s is None else refresh_margin_s),
        )
        self._refresh_interval_s = max(
            0,
            int(_DEFAULT_REFRESH_INTERVAL_S if refresh_interval_s is None else refresh_interval_s),
        )
        self._lock = threading.Lock()
        self._token: str | None = None
        self._expires_on: int | None = None
        self._stop_event = threading.Event()
        self._thread: threading.Thread | None = None

        if auto_refresh:
            self.start_auto_refresh()

    def __call__(self) -> str:
        return self.get_token()

    def get_token(self) -> str:
        self._refresh_if_needed(time.time(), force=False)
        if self._token is None:
            raise RuntimeError("Failed to acquire Azure AD token.")
        return self._token

    def start_auto_refresh(self) -> None:
        if self._refresh_interval_s <= 0:
            return
        if self._thread and self._thread.is_alive():
            return
        self._stop_event.clear()
        self._thread = threading.Thread(
            target=self._refresh_loop,
            name=f"AzureTokenRefresher({','.join(self._scopes)})",
            daemon=True,
        )
        self._thread.start()

    def stop_auto_refresh(self) -> None:
        if self._thread is None:
            return
        self._stop_event.set()
        self._thread.join(timeout=2)
        self._thread = None

    def _needs_refresh(self, now_s: float) -> bool:
        if self._token is None or self._expires_on is None:
            return True
        return now_s >= (self._expires_on - self._refresh_margin_s)

    def _refresh_if_needed(self, now_s: float, *, force: bool) -> None:
        if not force and not self._needs_refresh(now_s):
            return
        with self._lock:
            now_s = time.time()
            if not force and not self._needs_refresh(now_s):
                return
            token = self._credential.get_token(*self._scopes)
            self._token = token.token
            self._expires_on = int(token.expires_on)
            if _REFRESH_LOG_ENABLED:
                _ensure_refresh_logger()
                logger.debug(
                    "Azure token refreshed; scopes=%s expires_on=%s",
                    self._scopes,
                    self._expires_on,
                )

    def _refresh_loop(self) -> None:
        while not self._stop_event.is_set():
            try:
                self._refresh_if_needed(time.time(), force=False)
            except Exception as exc:
                logger.warning("Azure token refresh failed: %s", exc)
            self._stop_event.wait(self._refresh_interval_s)


_default_credential: DefaultAzureCredential | None = None
_provider_lock = threading.Lock()
_providers: dict[tuple[str, ...], AzureAuthTokenProvider] = {}


def _get_default_credential() -> DefaultAzureCredential:
    global _default_credential
    if _default_credential is None:
        _default_credential = DefaultAzureCredential()
    return _default_credential


def get_azure_token_provider(
    scope: str | Sequence[str],
    *,
    refresh_margin_s: int | None = None,
    refresh_interval_s: int | None = None,
    auto_refresh: bool = True,
) -> AzureAuthTokenProvider:
    scopes = (scope,) if isinstance(scope, str) else tuple(scope)
    with _provider_lock:
        provider = _providers.get(scopes)
        if provider is None:
            provider = AzureAuthTokenProvider(
                scopes=scopes,
                credential=_get_default_credential(),
                refresh_margin_s=refresh_margin_s,
                refresh_interval_s=refresh_interval_s,
                auto_refresh=auto_refresh,
            )
            _providers[scopes] = provider
    if auto_refresh:
        provider.start_auto_refresh()
    return provider


def stop_all_token_providers() -> None:
    with _provider_lock:
        providers = list(_providers.values())
    for provider in providers:
        provider.stop_auto_refresh()
