# [DEF:backend.src.core.utils.async_network:Module]
#
# @TIER: CRITICAL
# @SEMANTICS: network, httpx, async, superset, authentication, cache
# @PURPOSE: Provides async Superset API client with shared auth-token cache to avoid per-request re-login.
# @LAYER: Infra
# @RELATION: DEPENDS_ON -> backend.src.core.utils.network.SupersetAuthCache
# @INVARIANT: Async client reuses cached auth tokens per environment credentials and invalidates on 401.

# [SECTION: IMPORTS]
from typing import Optional, Dict, Any, Union
import asyncio

import httpx

from ..logger import logger as app_logger, belief_scope
from .network import (
    AuthenticationError,
    DashboardNotFoundError,
    NetworkError,
    PermissionDeniedError,
    SupersetAPIError,
    SupersetAuthCache,
)
# [/SECTION]


# [DEF:AsyncAPIClient:Class]
# @PURPOSE: Async Superset API client backed by httpx.AsyncClient with shared auth cache.
class AsyncAPIClient:
    DEFAULT_TIMEOUT = 30
    _auth_locks: Dict[tuple[str, str, bool], asyncio.Lock] = {}

    # [DEF:__init__:Function]
    # @PURPOSE: Initialize async API client for one environment.
    # @PRE: config contains base_url and auth payload.
    # @POST: Client is ready for async request/authentication flow.
    def __init__(self, config: Dict[str, Any], verify_ssl: bool = True, timeout: int = DEFAULT_TIMEOUT):
        self.base_url: str = self._normalize_base_url(config.get("base_url", ""))
        self.api_base_url: str = f"{self.base_url}/api/v1"
        self.auth = config.get("auth")
        self.request_settings = {"verify_ssl": verify_ssl, "timeout": timeout}
        self._client = httpx.AsyncClient(
            verify=verify_ssl,
            timeout=httpx.Timeout(timeout),
            follow_redirects=True,
        )
        self._tokens: Dict[str, str] = {}
        self._authenticated = False
        self._auth_cache_key = SupersetAuthCache.build_key(
            self.base_url,
            self.auth,
            verify_ssl,
        )

    # [/DEF:__init__:Function]

    # [DEF:_normalize_base_url:Function]
    # @PURPOSE: Normalize base URL for Superset API root construction.
    # @POST: Returns canonical base URL without trailing slash and duplicate /api/v1 suffix.
    def _normalize_base_url(self, raw_url: str) -> str:
        normalized = str(raw_url or "").strip().rstrip("/")
        if normalized.lower().endswith("/api/v1"):
            normalized = normalized[:-len("/api/v1")]
        return normalized.rstrip("/")
    # [/DEF:_normalize_base_url:Function]

    # [DEF:_build_api_url:Function]
    # @PURPOSE: Build full API URL from relative Superset endpoint.
    # @POST: Returns absolute URL for upstream request.
    def _build_api_url(self, endpoint: str) -> str:
        normalized_endpoint = str(endpoint or "").strip()
        if normalized_endpoint.startswith("http://") or normalized_endpoint.startswith("https://"):
            return normalized_endpoint
        if not normalized_endpoint.startswith("/"):
            normalized_endpoint = f"/{normalized_endpoint}"
        if normalized_endpoint.startswith("/api/v1/") or normalized_endpoint == "/api/v1":
            return f"{self.base_url}{normalized_endpoint}"
        return f"{self.api_base_url}{normalized_endpoint}"
    # [/DEF:_build_api_url:Function]

    # [DEF:_get_auth_lock:Function]
    # @PURPOSE: Return per-cache-key async lock to serialize fresh login attempts.
    # @POST: Returns stable asyncio.Lock instance.
    @classmethod
    def _get_auth_lock(cls, cache_key: tuple[str, str, bool]) -> asyncio.Lock:
        existing_lock = cls._auth_locks.get(cache_key)
        if existing_lock is not None:
            return existing_lock
        created_lock = asyncio.Lock()
        cls._auth_locks[cache_key] = created_lock
        return created_lock
    # [/DEF:_get_auth_lock:Function]

    # [DEF:authenticate:Function]
    # @PURPOSE: Authenticate against Superset and cache access/csrf tokens.
    # @POST: Client tokens are populated and reusable across requests.
    async def authenticate(self) -> Dict[str, str]:
        cached_tokens = SupersetAuthCache.get(self._auth_cache_key)
        if cached_tokens and cached_tokens.get("access_token") and cached_tokens.get("csrf_token"):
            self._tokens = cached_tokens
            self._authenticated = True
            app_logger.info("[async_authenticate][CacheHit] Reusing cached Superset auth tokens for %s", self.base_url)
            return self._tokens

        auth_lock = self._get_auth_lock(self._auth_cache_key)
        async with auth_lock:
            cached_tokens = SupersetAuthCache.get(self._auth_cache_key)
            if cached_tokens and cached_tokens.get("access_token") and cached_tokens.get("csrf_token"):
                self._tokens = cached_tokens
                self._authenticated = True
                app_logger.info("[async_authenticate][CacheHitAfterWait] Reusing cached Superset auth tokens for %s", self.base_url)
                return self._tokens

            with belief_scope("AsyncAPIClient.authenticate"):
                app_logger.info("[async_authenticate][Enter] Authenticating to %s", self.base_url)
                try:
                    login_url = f"{self.api_base_url}/security/login"
                    response = await self._client.post(login_url, json=self.auth)
                    response.raise_for_status()
                    access_token = response.json()["access_token"]

                    csrf_url = f"{self.api_base_url}/security/csrf_token/"
                    csrf_response = await self._client.get(
                        csrf_url,
                        headers={"Authorization": f"Bearer {access_token}"},
                    )
                    csrf_response.raise_for_status()

                    self._tokens = {
                        "access_token": access_token,
                        "csrf_token": csrf_response.json()["result"],
                    }
                    self._authenticated = True
                    SupersetAuthCache.set(self._auth_cache_key, self._tokens)
                    app_logger.info("[async_authenticate][Exit] Authenticated successfully.")
                    return self._tokens
                except httpx.HTTPStatusError as exc:
                    SupersetAuthCache.invalidate(self._auth_cache_key)
                    status_code = exc.response.status_code if exc.response is not None else None
                    if status_code in [502, 503, 504]:
                        raise NetworkError(
                            f"Environment unavailable during authentication (Status {status_code})",
                            status_code=status_code,
                        ) from exc
                    raise AuthenticationError(f"Authentication failed: {exc}") from exc
                except (httpx.HTTPError, KeyError) as exc:
                    SupersetAuthCache.invalidate(self._auth_cache_key)
                    raise NetworkError(f"Network or parsing error during authentication: {exc}") from exc
    # [/DEF:authenticate:Function]

    # [DEF:get_headers:Function]
    # @PURPOSE: Return authenticated Superset headers for async requests.
    # @POST: Headers include Authorization and CSRF tokens.
    async def get_headers(self) -> Dict[str, str]:
        if not self._authenticated:
            await self.authenticate()
        return {
            "Authorization": f"Bearer {self._tokens['access_token']}",
            "X-CSRFToken": self._tokens.get("csrf_token", ""),
            "Referer": self.base_url,
            "Content-Type": "application/json",
        }
    # [/DEF:get_headers:Function]

    # [DEF:request:Function]
    # @PURPOSE: Perform one authenticated async Superset API request.
    # @POST: Returns JSON payload or raw httpx.Response when raw_response=true.
    async def request(
        self,
        method: str,
        endpoint: str,
        headers: Optional[Dict[str, str]] = None,
        raw_response: bool = False,
        **kwargs,
    ) -> Union[httpx.Response, Dict[str, Any]]:
        full_url = self._build_api_url(endpoint)
        request_headers = await self.get_headers()
        if headers:
            request_headers.update(headers)
        if "allow_redirects" in kwargs and "follow_redirects" not in kwargs:
            kwargs["follow_redirects"] = bool(kwargs.pop("allow_redirects"))

        try:
            response = await self._client.request(method, full_url, headers=request_headers, **kwargs)
            response.raise_for_status()
            return response if raw_response else response.json()
        except httpx.HTTPStatusError as exc:
            if exc.response is not None and exc.response.status_code == 401:
                self._authenticated = False
                self._tokens = {}
                SupersetAuthCache.invalidate(self._auth_cache_key)
            self._handle_http_error(exc, endpoint)
        except httpx.HTTPError as exc:
            self._handle_network_error(exc, full_url)
    # [/DEF:request:Function]

    # [DEF:_handle_http_error:Function]
    # @PURPOSE: Translate upstream HTTP errors into stable domain exceptions.
    # @POST: Raises domain-specific exception for caller flow control.
    def _handle_http_error(self, exc: httpx.HTTPStatusError, endpoint: str) -> None:
        with belief_scope("AsyncAPIClient._handle_http_error"):
            status_code = exc.response.status_code
        if status_code in [502, 503, 504]:
            raise NetworkError(f"Environment unavailable (Status {status_code})", status_code=status_code) from exc
        if status_code == 404:
            raise DashboardNotFoundError(endpoint) from exc
        if status_code == 403:
            raise PermissionDeniedError() from exc
        if status_code == 401:
            raise AuthenticationError() from exc
        raise SupersetAPIError(f"API Error {status_code}: {exc.response.text}") from exc
    # [/DEF:_handle_http_error:Function]

    # [DEF:_handle_network_error:Function]
    # @PURPOSE: Translate generic httpx errors into NetworkError.
    # @POST: Raises NetworkError with URL context.
    def _handle_network_error(self, exc: httpx.HTTPError, url: str) -> None:
        with belief_scope("AsyncAPIClient._handle_network_error"):
            if isinstance(exc, httpx.TimeoutException):
                message = "Request timeout"
            elif isinstance(exc, httpx.ConnectError):
                message = "Connection error"
            else:
                message = f"Unknown network error: {exc}"
            raise NetworkError(message, url=url) from exc
    # [/DEF:_handle_network_error:Function]

    # [DEF:aclose:Function]
    # @PURPOSE: Close underlying httpx client.
    # @POST: Client resources are released.
    async def aclose(self) -> None:
        await self._client.aclose()
    # [/DEF:aclose:Function]
# [/DEF:AsyncAPIClient:Class]

# [/DEF:backend.src.core.utils.async_network:Module]
