# <GRACE_MODULE id="superset_tool.client" name="client.py">
#   @SEMANTICS: superset, api, client, rest, http, dashboard, dataset, import, export
#   @PURPOSE: Предоставляет высокоуровневый клиент для взаимодействия с Superset REST API, инкапсулируя логику запросов, обработку ошибок и пагинацию.
#   @DEPENDS_ON: superset_tool.models -> Использует SupersetConfig для конфигурации.
#   @DEPENDS_ON: superset_tool.exceptions -> Выбрасывает специализированные исключения.
#   @DEPENDS_ON: superset_tool.utils -> Использует утилиты для сети, логгирования и работы с файлами.

# <IMPORTS>
import json
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union
from requests import Response
from superset_tool.models import SupersetConfig
from superset_tool.exceptions import ExportError, InvalidZipFormatError
from superset_tool.utils.fileio import get_filename_from_headers
from superset_tool.utils.logger import SupersetLogger
from superset_tool.utils.network import APIClient
# </IMPORTS>

# --- Начало кода модуля ---

# <ANCHOR id="SupersetClient" type="Class">
#   @PURPOSE: Класс-обёртка над Superset REST API, предоставляющий методы для работы с дашбордами и датасетами.
#   @RELATION: CREATES_INSTANCE_OF -> APIClient
#   @RELATION: USES -> SupersetConfig
class SupersetClient:
    def __init__(self, config: SupersetConfig, logger: Optional[SupersetLogger] = None):
        # <ANCHOR id="SupersetClient.__init__" type="Function">
        #   @PURPOSE: Инициализирует клиент, проверяет конфигурацию и создает сетевой клиент.
        #   @PARAM: config: SupersetConfig - Конфигурация подключения.
        #   @PARAM: logger: Optional[SupersetLogger] - Экземпляр логгера.
        #   @POST: Атрибуты `logger`, `config`, и `network` созданы.
        self.logger = logger or SupersetLogger(name="SupersetClient")
        self.logger.info("[SupersetClient.__init__][Enter] Initializing SupersetClient.")
        self._validate_config(config)
        self.config = config
        self.network = APIClient(
            config=config.dict(),
            verify_ssl=config.verify_ssl,
            timeout=config.timeout,
            logger=self.logger,
        )
        self.delete_before_reimport: bool = False
        self.logger.info("[SupersetClient.__init__][Exit] SupersetClient initialized.")
        # </ANCHOR id="SupersetClient.__init__">

    # <ANCHOR id="SupersetClient._validate_config" type="Function">
    #   @PURPOSE: Проверяет, что переданный объект конфигурации имеет корректный тип.
    #   @PARAM: config: SupersetConfig - Объект для проверки.
    #   @THROW: TypeError - Если `config` не является экземпляром `SupersetConfig`.
    def _validate_config(self, config: SupersetConfig) -> None:
        self.logger.debug("[_validate_config][Enter] Validating SupersetConfig.")
        assert isinstance(config, SupersetConfig), "Конфигурация должна быть экземпляром SupersetConfig"
        self.logger.debug("[_validate_config][Exit] Config is valid.")
    # </ANCHOR id="SupersetClient._validate_config">

    @property
    def headers(self) -> dict:
        # <ANCHOR id="SupersetClient.headers" type="Property">
        #   @PURPOSE: Возвращает базовые HTTP-заголовки, используемые сетевым клиентом.
        return self.network.headers
        # </ANCHOR id="SupersetClient.headers">

    # <ANCHOR id="SupersetClient.get_dashboards" type="Function">
    #   @PURPOSE: Получает полный список дашбордов, автоматически обрабатывая пагинацию.
    #   @PARAM: query: Optional[Dict] - Дополнительные параметры запроса для API.
    #   @RETURN: Tuple[int, List[Dict]] - Кортеж (общее количество, список дашбордов).
    #   @RELATION: CALLS -> self._fetch_total_object_count
    #   @RELATION: CALLS -> self._fetch_all_pages
    def get_dashboards(self, query: Optional[Dict] = None) -> Tuple[int, List[Dict]]:
        self.logger.info("[get_dashboards][Enter] Fetching dashboards.")
        validated_query = self._validate_query_params(query)
        total_count = self._fetch_total_object_count(endpoint="/dashboard/")
        paginated_data = self._fetch_all_pages(
            endpoint="/dashboard/",
            pagination_options={"base_query": validated_query, "total_count": total_count, "results_field": "result"},
        )
        self.logger.info("[get_dashboards][Exit] Found %d dashboards.", total_count)
        return total_count, paginated_data
    # </ANCHOR id="SupersetClient.get_dashboards">

    # <ANCHOR id="SupersetClient.export_dashboard" type="Function">
    #   @PURPOSE: Экспортирует дашборд в виде ZIP-архива.
    #   @PARAM: dashboard_id: int - ID дашборда для экспорта.
    #   @RETURN: Tuple[bytes, str] - Бинарное содержимое ZIP-архива и имя файла.
    #   @THROW: ExportError - Если экспорт завершился неудачей.
    #   @RELATION: CALLS -> self.network.request
    def export_dashboard(self, dashboard_id: int) -> Tuple[bytes, str]:
        self.logger.info("[export_dashboard][Enter] Exporting dashboard %s.", dashboard_id)
        response = self.network.request(
            method="GET",
            endpoint="/dashboard/export/",
            params={"q": json.dumps([dashboard_id])},
            stream=True,
            raw_response=True,
        )
        self._validate_export_response(response, dashboard_id)
        filename = self._resolve_export_filename(response, dashboard_id)
        self.logger.info("[export_dashboard][Exit] Exported dashboard %s to %s.", dashboard_id, filename)
        return response.content, filename
    # </ANCHOR id="SupersetClient.export_dashboard">

    # <ANCHOR id="SupersetClient.import_dashboard" type="Function">
    #   @PURPOSE: Импортирует дашборд из ZIP-файла с возможностью автоматического удаления и повторной попытки при ошибке.
    #   @PARAM: file_name: Union[str, Path] - Путь к ZIP-архиву.
    #   @PARAM: dash_id: Optional[int] - ID дашборда для удаления при сбое.
    #   @PARAM: dash_slug: Optional[str] - Slug дашборда для поиска ID, если ID не предоставлен.
    #   @RETURN: Dict - Ответ API в случае успеха.
    #   @RELATION: CALLS -> self._do_import
    #   @RELATION: CALLS -> self.delete_dashboard
    #   @RELATION: CALLS -> self.get_dashboards
    def import_dashboard(self, file_name: Union[str, Path], dash_id: Optional[int] = None, dash_slug: Optional[str] = None) -> Dict:
        file_path = str(file_name)
        self._validate_import_file(file_path)
        try:
            return self._do_import(file_path)
        except Exception as exc:
            self.logger.error("[import_dashboard][Failure] First import attempt failed: %s", exc, exc_info=True)
            if not self.delete_before_reimport:
                raise

            target_id = self._resolve_target_id_for_delete(dash_id, dash_slug)
            if target_id is None:
                self.logger.error("[import_dashboard][Failure] No ID available for delete-retry.")
                raise

            self.delete_dashboard(target_id)
            self.logger.info("[import_dashboard][State] Deleted dashboard ID %s, retrying import.", target_id)
            return self._do_import(file_path)
    # </ANCHOR id="SupersetClient.import_dashboard">

    # <ANCHOR id="SupersetClient._resolve_target_id_for_delete" type="Function">
    #   @PURPOSE: Определяет ID дашборда для удаления, используя ID или slug.
    #   @INTERNAL
    def _resolve_target_id_for_delete(self, dash_id: Optional[int], dash_slug: Optional[str]) -> Optional[int]:
        if dash_id is not None:
            return dash_id
        if dash_slug is not None:
            self.logger.debug("[_resolve_target_id_for_delete][State] Resolving ID by slug '%s'.", dash_slug)
            try:
                _, candidates = self.get_dashboards(query={"filters": [{"col": "slug", "op": "eq", "value": dash_slug}]})
                if candidates:
                    target_id = candidates[0]["id"]
                    self.logger.debug("[_resolve_target_id_for_delete][Success] Resolved slug to ID %s.", target_id)
                    return target_id
            except Exception as e:
                self.logger.warning("[_resolve_target_id_for_delete][Warning] Could not resolve slug '%s' to ID: %s", dash_slug, e)
        return None
    # </ANCHOR id="SupersetClient._resolve_target_id_for_delete">

    # <ANCHOR id="SupersetClient._do_import" type="Function">
    #   @PURPOSE: Выполняет один запрос на импорт без обработки исключений.
    #   @INTERNAL
    def _do_import(self, file_name: Union[str, Path]) -> Dict:
        return self.network.upload_file(
            endpoint="/dashboard/import/",
            file_info={"file_obj": Path(file_name), "file_name": Path(file_name).name, "form_field": "formData"},
            extra_data={"overwrite": "true"},
            timeout=self.config.timeout * 2,
        )
    # </ANCHOR id="SupersetClient._do_import">

    # <ANCHOR id="SupersetClient.delete_dashboard" type="Function">
    #   @PURPOSE: Удаляет дашборд по его ID или slug.
    #   @PARAM: dashboard_id: Union[int, str] - ID или slug дашборда.
    #   @RELATION: CALLS -> self.network.request
    def delete_dashboard(self, dashboard_id: Union[int, str]) -> None:
        self.logger.info("[delete_dashboard][Enter] Deleting dashboard %s.", dashboard_id)
        response = self.network.request(method="DELETE", endpoint=f"/dashboard/{dashboard_id}")
        if response.get("result", True) is not False:
            self.logger.info("[delete_dashboard][Success] Dashboard %s deleted.", dashboard_id)
        else:
            self.logger.warning("[delete_dashboard][Warning] Unexpected response while deleting %s: %s", dashboard_id, response)
    # </ANCHOR id="SupersetClient.delete_dashboard">

    # <ANCHOR id="SupersetClient._extract_dashboard_id_from_zip" type="Function">
    #   @PURPOSE: Извлекает ID дашборда из `metadata.yaml` внутри ZIP-архива.
    #   @INTERNAL
    def _extract_dashboard_id_from_zip(self, file_name: Union[str, Path]) -> Optional[int]:
        try:
            import yaml
            with zipfile.ZipFile(file_name, "r") as zf:
                for name in zf.namelist():
                    if name.endswith("metadata.yaml"):
                        with zf.open(name) as meta_file:
                            meta = yaml.safe_load(meta_file)
                            dash_id = meta.get("dashboard_uuid") or meta.get("dashboard_id")
                            if dash_id: return int(dash_id)
        except Exception as exc:
            self.logger.error("[_extract_dashboard_id_from_zip][Failure] %s", exc, exc_info=True)
        return None
    # </ANCHOR id="SupersetClient._extract_dashboard_id_from_zip">

    # <ANCHOR id="SupersetClient._extract_dashboard_slug_from_zip" type="Function">
    #   @PURPOSE: Извлекает slug дашборда из `metadata.yaml` внутри ZIP-архива.
    #   @INTERNAL
    def _extract_dashboard_slug_from_zip(self, file_name: Union[str, Path]) -> Optional[str]:
        try:
            import yaml
            with zipfile.ZipFile(file_name, "r") as zf:
                for name in zf.namelist():
                    if name.endswith("metadata.yaml"):
                        with zf.open(name) as meta_file:
                            meta = yaml.safe_load(meta_file)
                            if slug := meta.get("slug"):
                                return str(slug)
        except Exception as exc:
            self.logger.error("[_extract_dashboard_slug_from_zip][Failure] %s", exc, exc_info=True)
        return None
    # </ANCHOR id="SupersetClient._extract_dashboard_slug_from_zip">

    # <ANCHOR id="SupersetClient._validate_export_response" type="Function">
    #   @PURPOSE: Проверяет, что HTTP-ответ на экспорт является валидным ZIP-архивом.
    #   @INTERNAL
    #   @THROW: ExportError - Если ответ не является ZIP-архивом или пуст.
    def _validate_export_response(self, response: Response, dashboard_id: int) -> None:
        content_type = response.headers.get("Content-Type", "")
        if "application/zip" not in content_type:
            raise ExportError(f"Получен не ZIP-архив (Content-Type: {content_type})")
        if not response.content:
            raise ExportError("Получены пустые данные при экспорте")
    # </ANCHOR id="SupersetClient._validate_export_response">

    # <ANCHOR id="SupersetClient._resolve_export_filename" type="Function">
    #   @PURPOSE: Определяет имя файла для экспорта из заголовков или генерирует его.
    #   @INTERNAL
    def _resolve_export_filename(self, response: Response, dashboard_id: int) -> str:
        filename = get_filename_from_headers(response.headers)
        if not filename:
            from datetime import datetime
            timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
            filename = f"dashboard_export_{dashboard_id}_{timestamp}.zip"
            self.logger.warning("[_resolve_export_filename][Warning] Generated filename: %s", filename)
        return filename
    # </ANCHOR id="SupersetClient._resolve_export_filename">

    # <ANCHOR id="SupersetClient._validate_query_params" type="Function">
    #   @PURPOSE: Формирует корректный набор параметров запроса с пагинацией.
    #   @INTERNAL
    def _validate_query_params(self, query: Optional[Dict]) -> Dict:
        base_query = {"columns": ["slug", "id", "changed_on_utc", "dashboard_title", "published"], "page": 0, "page_size": 1000}
        return {**base_query, **(query or {})}
    # </ANCHOR id="SupersetClient._validate_query_params">

    # <ANCHOR id="SupersetClient._fetch_total_object_count" type="Function">
    #   @PURPOSE: Получает общее количество объектов по указанному эндпоинту для пагинации.
    #   @INTERNAL
    def _fetch_total_object_count(self, endpoint: str) -> int:
        return self.network.fetch_paginated_count(
            endpoint=endpoint,
            query_params={"page": 0, "page_size": 1},
            count_field="count",
        )
    # </ANCHOR id="SupersetClient._fetch_total_object_count">

    # <ANCHOR id="SupersetClient._fetch_all_pages" type="Function">
    #   @PURPOSE: Итерируется по всем страницам пагинированного API и собирает все данные.
    #   @INTERNAL
    def _fetch_all_pages(self, endpoint: str, pagination_options: Dict) -> List[Dict]:
        return self.network.fetch_paginated_data(endpoint=endpoint, pagination_options=pagination_options)
    # </ANCHOR id="SupersetClient._fetch_all_pages">

    # <ANCHOR id="SupersetClient._validate_import_file" type="Function">
    #   @PURPOSE: Проверяет, что файл существует, является ZIP-архивом и содержит `metadata.yaml`.
    #   @INTERNAL
    #   @THROW: FileNotFoundError - Если файл не найден.
    #   @THROW: InvalidZipFormatError - Если файл не является ZIP или не содержит `metadata.yaml`.
    def _validate_import_file(self, zip_path: Union[str, Path]) -> None:
        path = Path(zip_path)
        assert path.exists(), f"Файл {zip_path} не существует"
        assert zipfile.is_zipfile(path), f"Файл {zip_path} не является ZIP-архивом"
        with zipfile.ZipFile(path, "r") as zf:
            assert any(n.endswith("metadata.yaml") for n in zf.namelist()), f"Архив {zip_path} не содержит 'metadata.yaml'"
    # </ANCHOR id="SupersetClient._validate_import_file">

    # <ANCHOR id="SupersetClient.get_datasets" type="Function">
    #   @PURPOSE: Получает полный список датасетов, автоматически обрабатывая пагинацию.
    #   @PARAM: query: Optional[Dict] - Дополнительные параметры запроса для API.
    #   @RETURN: Tuple[int, List[Dict]] - Кортеж (общее количество, список датасетов).
    #   @RELATION: CALLS -> self._fetch_total_object_count
    #   @RELATION: CALLS -> self._fetch_all_pages
    def get_datasets(self, query: Optional[Dict] = None) -> Tuple[int, List[Dict]]:
        self.logger.info("[get_datasets][Enter] Fetching datasets.")
        validated_query = self._validate_query_params(query)
        total_count = self._fetch_total_object_count(endpoint="/dataset/")
        paginated_data = self._fetch_all_pages(
            endpoint="/dataset/",
            pagination_options={"base_query": validated_query, "total_count": total_count, "results_field": "result"},
        )
        self.logger.info("[get_datasets][Exit] Found %d datasets.", total_count)
        return total_count, paginated_data
    # </ANCHOR id="SupersetClient.get_datasets">

    # <ANCHOR id="SupersetClient.get_dataset" type="Function">
    #   @PURPOSE: Получает информацию о конкретном датасете по его ID.
    #   @PARAM: dataset_id: int - ID датасета.
    #   @RETURN: Dict - Словарь с информацией о датасете.
    #   @RELATION: CALLS -> self.network.request
    def get_dataset(self, dataset_id: int) -> Dict:
        self.logger.info("[get_dataset][Enter] Fetching dataset %s.", dataset_id)
        response = self.network.request(method="GET", endpoint=f"/dataset/{dataset_id}")
        self.logger.info("[get_dataset][Exit] Got dataset %s.", dataset_id)
        return response
    # </ANCHOR id="SupersetClient.get_dataset">

    # <ANCHOR id="SupersetClient.update_dataset" type="Function">
    #   @PURPOSE: Обновляет данные датасета по его ID.
    #   @PARAM: dataset_id: int - ID датасета для обновления.
    #   @PARAM: data: Dict - Словарь с данными для обновления.
    #   @RETURN: Dict - Ответ API.
    #   @RELATION: CALLS -> self.network.request
    def update_dataset(self, dataset_id: int, data: Dict) -> Dict:
        self.logger.info("[update_dataset][Enter] Updating dataset %s.", dataset_id)
        response = self.network.request(
            method="PUT",
            endpoint=f"/dataset/{dataset_id}",
            data=json.dumps(data),
            headers={'Content-Type': 'application/json'}
        )
        self.logger.info("[update_dataset][Exit] Updated dataset %s.", dataset_id)
        return response
    # </ANCHOR id="SupersetClient.update_dataset">

# </ANCHOR id="SupersetClient">

# --- Конец кода модуля ---

# </GRACE_MODULE id="superset_tool.client">