# #region AsyncNetworkModule [C:5] [TYPE Module] [SEMANTICS auth, superset, token, async-client]
# @defgroup Core Module group.
#
# @BRIEF Provides async Superset API client with shared auth-token cache to avoid per-request re-login.
# @LAYER Infrastructure
# @PRE  Config payloads contain a Superset base URL and authentication fields needed for login.
# @POST  Async network clients reuse cached auth tokens and expose stable async request/error translation flow.
# @SIDE_EFFECT  Performs upstream HTTP I/O and mutates process-local auth cache entries.
# @DATA_CONTRACT  Input[config: Dict[str, Any]] -> Output[authenticated async Superset HTTP interactions]
# @RELATION DEPENDS_ON -> [SupersetAuthCache]
# @INVARIANT  Async client reuses cached auth tokens per environment credentials and invalidates on 401.
# @RATIONALE Async migration of Superset network layer. Uses httpx.AsyncClient with session-level
#   cookie persistence for CSRF. Shared auth cache reduces per-request login overhead from O(n) to O(1).
#   String data passed as httpx content (not json) to avoid double-encoding of pre-serialized JSON.
# @REJECTED requests.Session — blocks async event loop. Per-request AsyncAPIClient — loses connection
#   pooling and CSRF session cookies, causing 'CSRF tokens do not match' on SQL Lab execute.
import asyncio
import io
import json
from pathlib import Path
import ssl
from typing import Any

import httpx

from ..logger import belief_scope, logger as app_logger
from .network import (
    AuthenticationError,
    DashboardNotFoundError,
    NetworkError,
    PermissionDeniedError,
    SupersetAPIError,
    SupersetAuthCache,
)


# #region AsyncAPIClient [C:4] [TYPE Class]
# @defgroup Core Module group.
# @BRIEF Async Superset API client backed by httpx.AsyncClient with shared auth cache and optional semaphore.
# @PRE  config contains base_url and auth payload.
# @POST  Client is ready for async request/authentication flow. Connection pool initialised.
# @SIDE_EFFECT  Performs HTTP I/O; acquires/releases semaphore if configured.
# @RELATION DEPENDS_ON -> [SupersetAuthCache]
# @RELATION CALLS -> [SupersetAuthCache.get]
# @RELATION CALLS -> [SupersetAuthCache.set]
# @INVARIANT  semaphore.acquire() before request and semaphore.release() in finally.
# @RATIONALE One class handles auth lifecycle, CSRF token refresh on cache-hit, and per-env
#   semaphore for connection backpressure. authenticate() was extended to re-fetch CSRF token
#   on cache-hit to set session cookie on new httpx.AsyncClient instances.
# @REJECTED Per-request client with separate httpx.AsyncClient — loses CSRF session cookie,
#   causes 'CSRF tokens do not match' on SQL Lab execute. No semaphore — no backpressure.
class AsyncAPIClient:
    DEFAULT_TIMEOUT = 30
    _auth_locks: dict[tuple[str, str, bool], asyncio.Lock] = {}
    # #region AsyncAPIClient.__init__ [TYPE Function] [C:3]
# @ingroup Core
    # @PURPOSE: Initialize async API client for one environment with optional shared semaphore.
    # @PRE  config contains base_url and auth payload.
    # @POST  Client is ready for async request/authentication flow.
    # @DATA_CONTRACT  Input[config: Dict[str, Any]] -> self._auth_cache_key[str]
    # @RELATION CALLS -> [AsyncAPIClient._normalize_base_url]
    # @RELATION DEPENDS_ON -> [SupersetAuthCache]
    def __init__(self, config: dict[str, Any], verify_ssl: bool | ssl.SSLContext = True, timeout: int = DEFAULT_TIMEOUT,
                 semaphore: asyncio.Semaphore | None = None):
        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}
        # Use system CA store via ssl.create_default_context() instead of bool True.
        # httpx defaults to certifi, ignoring SSL_CERT_FILE and corporate CA certs.
        if verify_ssl is True:
            ssl_context: bool | ssl.SSLContext = ssl.create_default_context()
        elif verify_ssl is False:
            ssl_context = False
        else:
            ssl_context = verify_ssl  # уже SSLContext
        self._client = httpx.AsyncClient(
            verify=ssl_context,
            timeout=httpx.Timeout(timeout),
            follow_redirects=True,
        )
        self._tokens: dict[str, str] = {}
        self._authenticated = False
        # Cache key по-прежнему использует bool для хеширования
        self._auth_cache_key = SupersetAuthCache.build_key(
            self.base_url,
            self.auth,
            verify_ssl if isinstance(verify_ssl, bool) else True,
        )
        # Optional shared semaphore for connection limiting (set by SupersetClientRegistry).
        self._semaphore = semaphore
    # #endregion AsyncAPIClient.__init__
    # #region AsyncAPIClient._normalize_base_url [TYPE Function] [C:1]
    # @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("/")
    # #endregion AsyncAPIClient._normalize_base_url
    # #region AsyncAPIClient._build_api_url [TYPE Function] [C:1]
    # @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}"
    # #endregion AsyncAPIClient._build_api_url
    # #region AsyncAPIClient._get_auth_lock [TYPE Function] [C:1]
    # @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
    # #endregion AsyncAPIClient._get_auth_lock
    # #region AsyncAPIClient.authenticate [TYPE Function] [C:3]
# @ingroup Core
    # @PURPOSE: Authenticate against Superset and cache access/csrf tokens.
    # @POST  Client tokens are populated and reusable across requests.
    # @SIDE_EFFECT  Performs network requests to Superset authentication endpoints.
    # @DATA_CONTRACT  None -> Output[Dict[str, str]]
    # @RELATION CALLS -> [SupersetAuthCache.get]
    # @RELATION CALLS -> [SupersetAuthCache.set]
    # @RELATION CALLS -> [AsyncAPIClient._get_auth_lock]
    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)
            # Refresh CSRF token to set matching session cookie on this httpx client
            try:
                csrf_url = f"{self.api_base_url}/security/csrf_token/"
                csrf_resp = await self._client.get(
                    csrf_url,
                    headers={"Authorization": f"Bearer {cached_tokens['access_token']}"},
                )
                if csrf_resp.is_success:
                    new_csrf = csrf_resp.json().get("result")
                    if new_csrf:
                        self._tokens["csrf_token"] = new_csrf
                        SupersetAuthCache.set(self._auth_cache_key, self._tokens)
            except Exception:
                pass
            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)
                try:
                    csrf_url = f"{self.api_base_url}/security/csrf_token/"
                    csrf_resp = await self._client.get(
                        csrf_url,
                        headers={"Authorization": f"Bearer {cached_tokens['access_token']}"},
                    )
                    if csrf_resp.is_success:
                        new_csrf = csrf_resp.json().get("result")
                        if new_csrf:
                            self._tokens["csrf_token"] = new_csrf
                            SupersetAuthCache.set(self._auth_cache_key, self._tokens)
                except Exception:
                    pass
                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.reflect("Authenticated successfully.", extra={"src": "async_authenticate"})
                    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
    # #endregion AsyncAPIClient.authenticate
    # #region AsyncAPIClient.get_headers [TYPE Function] [C:3]
# @ingroup Core
    # @PURPOSE: Return authenticated Superset headers for async requests.
    # @POST  Headers include Authorization and CSRF tokens.
    # @RELATION CALLS -> [AsyncAPIClient.authenticate]
    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",
        }
    # #endregion AsyncAPIClient.get_headers
    # #region AsyncAPIClient.request [TYPE Function] [C:4]
# @ingroup Core
    # @PURPOSE: Perform one authenticated async Superset API request with optional semaphore.
    # @POST  Returns JSON payload (dict) or raw httpx.Response when raw_response=true.
    # @SIDE_EFFECT  Performs network I/O. Acquires/releases semaphore if configured.
    # @RELATION CALLS -> [AsyncAPIClient.get_headers]
    # @RELATION CALLS -> [AsyncAPIClient._handle_http_error]
    # @RELATION CALLS -> [AsyncAPIClient._handle_network_error]
    # @RELATION CALLS -> [AsyncAPIClient._is_dashboard_endpoint]
    # @RATIONALE Supports both pre-serialized JSON strings (legacy callers) and Python dicts.
    #   Strings sent via httpx content parameter to avoid double encoding by json parameter.
    #   This was the root cause of GENERIC_BACKEND_ERROR on all POST requests — httpx re-encoded
    #   pre-serialized JSON, turning the request body into a JSON string instead of a JSON object.
    # @REJECTED Passing all data through httpx json= — double-encodes strings. Using only content=
    #   breaks callers that pass dicts. Always passing dicts — requires refactoring 50+ callers.
    async def request(
        self,
        method: str,
        endpoint: str,
        params: dict[str, Any] | None = None,
        data: Any = None,
        headers: dict[str, str] | None = None,
        raw_response: bool = False,
        timeout: int | None = None,
        allow_redirects: bool = True,
    ) -> Any:
        with belief_scope("AsyncAPIClient.request"):
            # Acquire semaphore if configured (global per-env limit).
            if self._semaphore is not None:
                await self._semaphore.acquire()
            try:
                url = self._build_api_url(endpoint)
                req_headers = {**headers} if headers else {}
                auth_headers = await self.get_headers()
                req_headers.update(auth_headers)

                # Support both pre-serialized JSON strings (legacy convention)
                # and Python dicts/list. Strings are sent as raw body via `content`
                # to avoid double-encoding by httpx's `json` parameter.
                request_kwargs: dict[str, Any] = {
                    "method": method,
                    "url": url,
                    "params": params,
                    "headers": req_headers,
                    "timeout": httpx.Timeout(timeout or self.request_settings.get("timeout", self.DEFAULT_TIMEOUT)),
                    "follow_redirects": allow_redirects,
                }
                if isinstance(data, str):
                    request_kwargs["content"] = data.encode("utf-8")
                elif data is not None:
                    request_kwargs["json"] = data

                response = await self._client.request(**request_kwargs)

                if response.status_code == 401:
                    if method.upper() in ("GET", "HEAD", "OPTIONS"):
                        SupersetAuthCache.invalidate(self._auth_cache_key)
                        self._authenticated = False
                        auth_headers = await self.get_headers()
                        req_headers.update(auth_headers)
                        retry_kwargs = {**request_kwargs, "headers": req_headers}
                        response = await self._client.request(**retry_kwargs)
                    else:
                        raise AuthenticationError()

                if raw_response:
                    return response

                try:
                    return response.json()
                except Exception:
                    raise SupersetAPIError(
                        f"Failed to parse response JSON: {response.text[:500]}",
                        status_code=response.status_code,
                    )
            except httpx.HTTPStatusError as exc:
                self._handle_http_error(exc, endpoint)
            except httpx.HTTPError as exc:
                url = self._build_api_url(endpoint)
                self._handle_network_error(exc, url)
            finally:
                if self._semaphore is not None:
                    self._semaphore.release()
    # #endregion AsyncAPIClient.request

    # #region AsyncAPIClient._handle_http_error [TYPE Function] [C:3]
# @ingroup Core
    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:
            if self._is_dashboard_endpoint(endpoint):
                raise DashboardNotFoundError(endpoint) from exc
            raise SupersetAPIError(
                f"API resource not found at endpoint '{endpoint}'",
                status_code=status_code,
                endpoint=endpoint,
                subtype="not_found",
            ) 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
    # #endregion AsyncAPIClient._handle_http_error
    # #region AsyncAPIClient._is_dashboard_endpoint [TYPE Function] [C:2]
# @ingroup Core
    # @PURPOSE: Determine whether an API endpoint represents a dashboard resource for 404 translation.
    # @POST  Returns true only for dashboard-specific endpoints.
    def _is_dashboard_endpoint(self, endpoint: str) -> bool:
        normalized_endpoint = str(endpoint or "").strip().lower()
        if not normalized_endpoint:
            return False
        if normalized_endpoint.startswith("http://") or normalized_endpoint.startswith("https://"):
            try:
                normalized_endpoint = "/" + normalized_endpoint.split("/api/v1", 1)[1].lstrip("/")
            except IndexError:
                return False
        if normalized_endpoint.startswith("/api/v1/"):
            normalized_endpoint = normalized_endpoint[len("/api/v1"):]
        return normalized_endpoint.startswith("/dashboard/") or normalized_endpoint == "/dashboard"
    # #endregion AsyncAPIClient._is_dashboard_endpoint
    # #region AsyncAPIClient._handle_network_error [TYPE Function] [C:3]
# @ingroup Core
    # @PURPOSE: Translate generic httpx errors into NetworkError.
    # @POST  Raises NetworkError with URL context.
    # @DATA_CONTRACT  Input[httpx.HTTPError] -> NetworkError
    # @RELATION DEPENDS_ON -> [NetworkError]
    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
    # #endregion AsyncAPIClient._handle_network_error
    # #region AsyncAPIClient.fetch_paginated_count [TYPE Function] [C:3]
# @ingroup Core
    # @PURPOSE: Fetch total count of items for a paginated endpoint.
    # @POST  Returns integer count from API response.
    # @RELATION CALLS -> [AsyncAPIClient.request]
    async def fetch_paginated_count(
        self,
        endpoint: str,
        query_params: dict[str, Any],
        count_field: str = "count",
    ) -> int:
        with belief_scope("AsyncAPIClient.fetch_paginated_count"):
            response = await self.request(
                method="GET",
                endpoint=endpoint,
                params={"q": json.dumps(query_params)},
            )
            return response.get(count_field, 0) if isinstance(response, dict) else 0
    # #endregion AsyncAPIClient.fetch_paginated_count

    # #region AsyncAPIClient.fetch_paginated_data [TYPE Function] [C:3]
# @ingroup Core
    # @PURPOSE: Fetch all pages of a paginated endpoint with automatic page iteration.
    # @POST  Returns combined list of all items across pages.
    # @SIDE_EFFECT  Performs multiple sequential GET requests to iterate pages.
    # @RELATION CALLS -> [AsyncAPIClient.request]
    async def fetch_paginated_data(
        self,
        endpoint: str,
        pagination_options: dict[str, Any],
    ) -> list[Any]:
        with belief_scope("AsyncAPIClient.fetch_paginated_data"):
            base_query = pagination_options["base_query"]
            total_count = pagination_options.get("total_count")
            results_field = pagination_options["results_field"]
            count_field = pagination_options.get("count_field", "count")
            page_size = base_query.get("page_size", 1000)

            results: list[Any] = []

            # Fetch first page
            query = {**base_query, "page": 0}
            response_json = await self.request(
                method="GET",
                endpoint=endpoint,
                params={"q": json.dumps(query)},
            )
            if isinstance(response_json, dict):
                first_page_results = response_json.get(results_field, [])
                results.extend(first_page_results)
                if total_count is None:
                    total_count = response_json.get(count_field, len(first_page_results))
            else:
                return results

            # Fetch remaining pages
            total_pages = max(0, (total_count + page_size - 1) // page_size)
            for page in range(1, total_pages):
                query = {**base_query, "page": page}
                page_response = await self.request(
                    method="GET",
                    endpoint=endpoint,
                    params={"q": json.dumps(query)},
                )
                if isinstance(page_response, dict):
                    results.extend(page_response.get(results_field, []))

            return results
    # #endregion AsyncAPIClient.fetch_paginated_data

    # #region AsyncAPIClient.upload_file [TYPE Function] [C:4]
# @ingroup Core
    # @PURPOSE: Upload a file via multipart/form-data POST with optional extra data fields.
    # @PRE  endpoint is valid relative Superset API endpoint. file_info contains file_obj and file_name.
    # @POST  Returns parsed JSON API response.
    # @SIDE_EFFECT  Performs multipart HTTP upload via httpx.AsyncClient.
    # @RELATION CALLS -> [AsyncAPIClient.get_headers]
    # @RATIONALE Async migration of requests.post to httpx.AsyncClient.post with multipart files.
    #   Uses httpx files= parameter for multipart upload with file_name and Content-Type per file.
    # @REJECTED requests.post with files= — blocks async event loop. Direct file read -> httpx content
    #   — loses multipart boundary handling.
    async def upload_file(
        self,
        endpoint: str,
        file_info: dict[str, Any],
        extra_data: dict[str, Any] | None = None,
        timeout: int | None = None,
    ) -> dict:
        with belief_scope("AsyncAPIClient.upload_file"):
            url = self._build_api_url(endpoint)
            req_headers = await self.get_headers()
            req_headers.pop("Content-Type", None)

            file_obj = file_info.get("file_obj")
            file_name = file_info.get("file_name")
            form_field = file_info.get("form_field", "file")

            if isinstance(file_obj, (str, Path)):
                file_path = Path(file_obj)
                if not file_path.exists():
                    raise FileNotFoundError(f"Upload file does not exist: {file_obj}")
                files = {
                    form_field: (file_name, file_path.read_bytes(), "application/x-zip-compressed"),
                }
            elif isinstance(file_obj, io.BytesIO):
                files = {
                    form_field: (file_name, file_obj.getvalue(), "application/x-zip-compressed"),
                }
            else:
                raise TypeError(f"Unsupported file_obj type: {type(file_obj)}")

            try:
                response = await self._client.post(
                    url,
                    data=extra_data or {},
                    files=files,
                    headers=req_headers,
                    timeout=httpx.Timeout(timeout or self.request_settings.get("timeout", self.DEFAULT_TIMEOUT)),
                )
                response.raise_for_status()
                if response.status_code == 200:
                    try:
                        return response.json()
                    except Exception as json_e:
                        raise SupersetAPIError(
                            f"API error during upload: Response is not valid JSON: {json_e}"
                        ) from json_e
                return response.json()
            except httpx.HTTPStatusError as exc:
                raise SupersetAPIError(
                    f"API error during upload: {exc.response.text}"
                ) from exc
            except httpx.HTTPError as exc:
                raise NetworkError(f"Network error during upload: {exc}", url=url) from exc
    # #endregion AsyncAPIClient.upload_file

    # #region AsyncAPIClient.aclose [TYPE Function] [C:3]
# @ingroup Core
    # @PURPOSE: Close underlying httpx client.
    # @POST  Client resources are released.
    # @SIDE_EFFECT  Closes network connections.
    # @RELATION DEPENDS_ON -> [AsyncAPIClient.__init__]
    async def aclose(self) -> None:
        await self._client.aclose()
    # #endregion AsyncAPIClient.aclose
# #endregion AsyncAPIClient
# #endregion AsyncNetworkModule
