# [DEF:backend.core.utils.network:Module]
#
# @SEMANTICS: network, http, client, api, requests, session, authentication
# @PURPOSE: Инкапсулирует низкоуровневую HTTP-логику для взаимодействия с Superset API, включая аутентификацию, управление сессией, retry-логику и обработку ошибок.
# @LAYER: Infra
# @RELATION: DEPENDS_ON -> backend.src.core.logger
# @RELATION: DEPENDS_ON -> requests
# @PUBLIC_API: APIClient

# [SECTION: IMPORTS]
from typing import Optional, Dict, Any, List, Union, cast
import json
import io
from pathlib import Path
import requests
from requests.adapters import HTTPAdapter
import urllib3
from urllib3.util.retry import Retry
from ..logger import logger as app_logger, belief_scope
# [/SECTION]

# [DEF:SupersetAPIError:Class]
# @PURPOSE: Base exception for all Superset API related errors.
class SupersetAPIError(Exception):
    # [DEF:__init__:Function]
    # @PURPOSE: Initializes the exception with a message and context.
    # @PRE: message is a string, context is a dict.
    # @POST: Exception is initialized with context.
    def __init__(self, message: str = "Superset API error", **context: Any):
        with belief_scope("SupersetAPIError.__init__"):
            self.context = context
            super().__init__(f"[API_FAILURE] {message} | Context: {self.context}")
    # [/DEF:__init__:Function]
# [/DEF:SupersetAPIError:Class]

# [DEF:AuthenticationError:Class]
# @PURPOSE: Exception raised when authentication fails.
class AuthenticationError(SupersetAPIError):
    # [DEF:__init__:Function]
    # @PURPOSE: Initializes the authentication error.
    # @PRE: message is a string, context is a dict.
    # @POST: AuthenticationError is initialized.
    def __init__(self, message: str = "Authentication failed", **context: Any):
        with belief_scope("AuthenticationError.__init__"):
            super().__init__(message, type="authentication", **context)
    # [/DEF:__init__:Function]
# [/DEF:AuthenticationError:Class]

# [DEF:PermissionDeniedError:Class]
# @PURPOSE: Exception raised when access is denied.
class PermissionDeniedError(AuthenticationError):
    # [DEF:__init__:Function]
    # @PURPOSE: Initializes the permission denied error.
    # @PRE: message is a string, context is a dict.
    # @POST: PermissionDeniedError is initialized.
    def __init__(self, message: str = "Permission denied", **context: Any):
        with belief_scope("PermissionDeniedError.__init__"):
            super().__init__(message, **context)
    # [/DEF:__init__:Function]
# [/DEF:PermissionDeniedError:Class]

# [DEF:DashboardNotFoundError:Class]
# @PURPOSE: Exception raised when a dashboard cannot be found.
class DashboardNotFoundError(SupersetAPIError):
    # [DEF:__init__:Function]
    # @PURPOSE: Initializes the not found error with resource ID.
    # @PRE: resource_id is provided.
    # @POST: DashboardNotFoundError is initialized.
    def __init__(self, resource_id: Union[int, str], message: str = "Dashboard not found", **context: Any):
        with belief_scope("DashboardNotFoundError.__init__"):
            super().__init__(f"Dashboard '{resource_id}' {message}", subtype="not_found", resource_id=resource_id, **context)
    # [/DEF:__init__:Function]
# [/DEF:DashboardNotFoundError:Class]

# [DEF:NetworkError:Class]
# @PURPOSE: Exception raised when a network level error occurs.
class NetworkError(Exception):
    # [DEF:__init__:Function]
    # @PURPOSE: Initializes the network error.
    # @PRE: message is a string.
    # @POST: NetworkError is initialized.
    def __init__(self, message: str = "Network connection failed", **context: Any):
        with belief_scope("NetworkError.__init__"):
            self.context = context
            super().__init__(f"[NETWORK_FAILURE] {message} | Context: {self.context}")
    # [/DEF:__init__:Function]
# [/DEF:NetworkError:Class]

# [DEF:APIClient:Class]
# @PURPOSE: Инкапсулирует HTTP-логику для работы с API, включая сессии, аутентификацию, и обработку запросов.
class APIClient:
    DEFAULT_TIMEOUT = 30

    # [DEF:__init__:Function]
    # @PURPOSE: Инициализирует API клиент с конфигурацией, сессией и логгером.
    # @PARAM: config (Dict[str, Any]) - Конфигурация.
    # @PARAM: verify_ssl (bool) - Проверять ли SSL.
    # @PARAM: timeout (int) - Таймаут запросов.
    # @PRE: config must contain 'base_url' and 'auth'.
    # @POST: APIClient instance is initialized with a session.
    def __init__(self, config: Dict[str, Any], verify_ssl: bool = True, timeout: int = DEFAULT_TIMEOUT):
        with belief_scope("__init__"):
            app_logger.info("[APIClient.__init__][Entry] Initializing APIClient.")
        self.base_url: str = config.get("base_url", "")
        self.auth = config.get("auth")
        self.request_settings = {"verify_ssl": verify_ssl, "timeout": timeout}
        self.session = self._init_session()
        self._tokens: Dict[str, str] = {}
        self._authenticated = False
        app_logger.info("[APIClient.__init__][Exit] APIClient initialized.")
    # [/DEF:__init__:Function]

    # [DEF:_init_session:Function]
    # @PURPOSE: Создает и настраивает `requests.Session` с retry-логикой.
    # @PRE: self.request_settings must be initialized.
    # @POST: Returns a configured requests.Session instance.
    # @RETURN: requests.Session - Настроенная сессия.
    def _init_session(self) -> requests.Session:
        with belief_scope("_init_session"):
            session = requests.Session()
        retries = Retry(total=3, backoff_factor=0.5, status_forcelist=[500, 502, 503, 504])
        adapter = HTTPAdapter(max_retries=retries)
        session.mount('http://', adapter)
        session.mount('https://', adapter)
        if not self.request_settings["verify_ssl"]:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
            app_logger.warning("[_init_session][State] SSL verification disabled.")
        session.verify = self.request_settings["verify_ssl"]
        return session
    # [/DEF:_init_session:Function]

    # [DEF:authenticate:Function]
    # @PURPOSE: Выполняет аутентификацию в Superset API и получает access и CSRF токены.
    # @PRE: self.auth and self.base_url must be valid.
    # @POST: `self._tokens` заполнен, `self._authenticated` установлен в `True`.
    # @RETURN: Dict[str, str] - Словарь с токенами.
    # @THROW: AuthenticationError, NetworkError - при ошибках.
    def authenticate(self) -> Dict[str, str]:
        with belief_scope("authenticate"):
            app_logger.info("[authenticate][Enter] Authenticating to %s", self.base_url)
            try:
                login_url = f"{self.base_url}/security/login"
                # Log the payload keys and values (masking password)
                masked_auth = {k: ("******" if k == "password" else v) for k, v in self.auth.items()}
                app_logger.info(f"[authenticate][Debug] Login URL: {login_url}")
                app_logger.info(f"[authenticate][Debug] Auth payload: {masked_auth}")
                
                response = self.session.post(login_url, json=self.auth, timeout=self.request_settings["timeout"])
                
                if response.status_code != 200:
                    app_logger.error(f"[authenticate][Error] Status: {response.status_code}, Response: {response.text}")
                
                response.raise_for_status()
                access_token = response.json()["access_token"]
                
                csrf_url = f"{self.base_url}/security/csrf_token/"
                csrf_response = self.session.get(csrf_url, headers={"Authorization": f"Bearer {access_token}"}, timeout=self.request_settings["timeout"])
                csrf_response.raise_for_status()
                
                self._tokens = {"access_token": access_token, "csrf_token": csrf_response.json()["result"]}
                self._authenticated = True
                app_logger.info("[authenticate][Exit] Authenticated successfully.")
                return self._tokens
            except requests.exceptions.HTTPError as e:
                status_code = e.response.status_code if e.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 e
                raise AuthenticationError(f"Authentication failed: {e}") from e
            except (requests.exceptions.RequestException, KeyError) as e:
                raise NetworkError(f"Network or parsing error during authentication: {e}") from e
    # [/DEF:authenticate:Function]

    @property
    # [DEF:headers:Function]
    # @PURPOSE: Возвращает HTTP-заголовки для аутентифицированных запросов.
    # @PRE: APIClient is initialized and authenticated or can be authenticated.
    # @POST: Returns headers including auth tokens.
    def headers(self) -> Dict[str, str]:
        with belief_scope("headers"):
            if not self._authenticated: 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:headers:Function]

    # [DEF:request:Function]
    # @PURPOSE: Выполняет универсальный HTTP-запрос к API.
    # @PARAM: method (str) - HTTP метод.
    # @PARAM: endpoint (str) - API эндпоинт.
    # @PARAM: headers (Optional[Dict]) - Дополнительные заголовки.
    # @PARAM: raw_response (bool) - Возвращать ли сырой ответ.
    # @PRE: method and endpoint must be strings.
    # @POST: Returns response content or raw Response object.
    # @RETURN: `requests.Response` если `raw_response=True`, иначе `dict`.
    # @THROW: SupersetAPIError, NetworkError и их подклассы.
    def request(self, method: str, endpoint: str, headers: Optional[Dict] = None, raw_response: bool = False, **kwargs) -> Union[requests.Response, Dict[str, Any]]:
        with belief_scope("request"):
            full_url = f"{self.base_url}{endpoint}"
        _headers = self.headers.copy()
        if headers: _headers.update(headers)
        
        try:
            response = self.session.request(method, full_url, headers=_headers, **kwargs)
            response.raise_for_status()
            return response if raw_response else response.json()
        except requests.exceptions.HTTPError as e:
            self._handle_http_error(e, endpoint)
        except requests.exceptions.RequestException as e:
            self._handle_network_error(e, full_url)
    # [/DEF:request:Function]

    # [DEF:_handle_http_error:Function]
    # @PURPOSE: (Helper) Преобразует HTTP ошибки в кастомные исключения.
    # @PARAM: e (requests.exceptions.HTTPError) - Ошибка.
    # @PARAM: endpoint (str) - Эндпоинт.
    # @PRE: e must be a valid HTTPError with a response.
    # @POST: Raises a specific SupersetAPIError or subclass.
    def _handle_http_error(self, e: requests.exceptions.HTTPError, endpoint: str):
        with belief_scope("_handle_http_error"):
            status_code = e.response.status_code
        if status_code == 502 or status_code == 503 or status_code == 504:
            raise NetworkError(f"Environment unavailable (Status {status_code})", status_code=status_code) from e
        if status_code == 404: raise DashboardNotFoundError(endpoint) from e
        if status_code == 403: raise PermissionDeniedError() from e
        if status_code == 401: raise AuthenticationError() from e
        raise SupersetAPIError(f"API Error {status_code}: {e.response.text}") from e
    # [/DEF:_handle_http_error:Function]

    # [DEF:_handle_network_error:Function]
    # @PURPOSE: (Helper) Преобразует сетевые ошибки в `NetworkError`.
    # @PARAM: e (requests.exceptions.RequestException) - Ошибка.
    # @PARAM: url (str) - URL.
    # @PRE: e must be a RequestException.
    # @POST: Raises a NetworkError.
    def _handle_network_error(self, e: requests.exceptions.RequestException, url: str):
        with belief_scope("_handle_network_error"):
            if isinstance(e, requests.exceptions.Timeout): msg = "Request timeout"
            elif isinstance(e, requests.exceptions.ConnectionError): msg = "Connection error"
            else: msg = f"Unknown network error: {e}"
            raise NetworkError(msg, url=url) from e
    # [/DEF:_handle_network_error:Function]

    # [DEF:upload_file:Function]
    # @PURPOSE: Загружает файл на сервер через multipart/form-data.
    # @PARAM: endpoint (str) - Эндпоинт.
    # @PARAM: file_info (Dict[str, Any]) - Информация о файле.
    # @PARAM: extra_data (Optional[Dict]) - Дополнительные данные.
    # @PARAM: timeout (Optional[int]) - Таймаут.
    # @PRE: file_info must contain 'file_obj' and 'file_name'.
    # @POST: File is uploaded and response returned.
    # @RETURN: Ответ API в виде словаря.
    # @THROW: SupersetAPIError, NetworkError, TypeError.
    def upload_file(self, endpoint: str, file_info: Dict[str, Any], extra_data: Optional[Dict] = None, timeout: Optional[int] = None) -> Dict:
        with belief_scope("upload_file"):
            full_url = f"{self.base_url}{endpoint}"
        _headers = self.headers.copy(); _headers.pop('Content-Type', None)
        
        file_obj, file_name, form_field = file_info.get("file_obj"), file_info.get("file_name"), file_info.get("form_field", "file")
        
        files_payload = {}
        if isinstance(file_obj, (str, Path)):
            with open(file_obj, 'rb') as f:
                files_payload = {form_field: (file_name, f.read(), 'application/x-zip-compressed')}
        elif isinstance(file_obj, io.BytesIO):
            files_payload = {form_field: (file_name, file_obj.getvalue(), 'application/x-zip-compressed')}
        else:
            raise TypeError(f"Unsupported file_obj type: {type(file_obj)}")
            
        return self._perform_upload(full_url, files_payload, extra_data, _headers, timeout)
    # [/DEF:upload_file:Function]

    # [DEF:_perform_upload:Function]
    # @PURPOSE: (Helper) Выполняет POST запрос с файлом.
    # @PARAM: url (str) - URL.
    # @PARAM: files (Dict) - Файлы.
    # @PARAM: data (Optional[Dict]) - Данные.
    # @PARAM: headers (Dict) - Заголовки.
    # @PARAM: timeout (Optional[int]) - Таймаут.
    # @PRE: url, files, and headers must be provided.
    # @POST: POST request is performed and JSON response returned.
    # @RETURN: Dict - Ответ.
    def _perform_upload(self, url: str, files: Dict, data: Optional[Dict], headers: Dict, timeout: Optional[int]) -> Dict:
        with belief_scope("_perform_upload"):
            try:
                response = self.session.post(url, files=files, data=data or {}, headers=headers, timeout=timeout or self.request_settings["timeout"])
                response.raise_for_status()
                if response.status_code == 200:
                    try:
                        return response.json()
                    except Exception as json_e:
                        app_logger.debug(f"[_perform_upload][Debug] Response is not valid JSON: {response.text[:200]}...")
                        raise SupersetAPIError(f"API error during upload: Response is not valid JSON: {json_e}") from json_e
                return response.json()
            except requests.exceptions.HTTPError as e:
                raise SupersetAPIError(f"API error during upload: {e.response.text}") from e
            except requests.exceptions.RequestException as e:
                raise NetworkError(f"Network error during upload: {e}", url=url) from e
    # [/DEF:_perform_upload:Function]

    # [DEF:fetch_paginated_count:Function]
    # @PURPOSE: Получает общее количество элементов для пагинации.
    # @PARAM: endpoint (str) - Эндпоинт.
    # @PARAM: query_params (Dict) - Параметры запроса.
    # @PARAM: count_field (str) - Поле с количеством.
    # @PRE: query_params must be a dictionary.
    # @POST: Returns total count of items.
    # @RETURN: int - Количество.
    def fetch_paginated_count(self, endpoint: str, query_params: Dict, count_field: str = "count") -> int:
        with belief_scope("fetch_paginated_count"):
            response_json = cast(Dict[str, Any], self.request("GET", endpoint, params={"q": json.dumps(query_params)}))
        return response_json.get(count_field, 0)
    # [/DEF:fetch_paginated_count:Function]

    # [DEF:fetch_paginated_data:Function]
    # @PURPOSE: Автоматически собирает данные со всех страниц пагинированного эндпоинта.
    # @PARAM: endpoint (str) - Эндпоинт.
    # @PARAM: pagination_options (Dict[str, Any]) - Опции пагинации.
    # @PRE: pagination_options must contain 'base_query', 'total_count', 'results_field'.
    # @POST: Returns all items across all pages.
    # @RETURN: List[Any] - Список данных.
    def fetch_paginated_data(self, endpoint: str, pagination_options: Dict[str, Any]) -> List[Any]:
        with belief_scope("fetch_paginated_data"):
            base_query, total_count = pagination_options["base_query"], pagination_options["total_count"]
        results_field, page_size = pagination_options["results_field"], base_query.get('page_size')
        assert page_size and page_size > 0, "'page_size' must be a positive number."
        
        results = []
        for page in range((total_count + page_size - 1) // page_size):
            query = {**base_query, 'page': page}
            response_json = cast(Dict[str, Any], self.request("GET", endpoint, params={"q": json.dumps(query)}))
            results.extend(response_json.get(results_field, []))
        return results
    # [/DEF:fetch_paginated_data:Function]

# [/DEF:APIClient:Class]

# [/DEF:backend.core.utils.network:Module]