# #region DatasetMapperModule [TYPE Module] [SEMANTICS dataset, superset, dataset-mapper, sqllab]
#
# @BRIEF Этот модуль отвечает за обновление метаданных (verbose_name) в датасетах Superset,
# извлекая их из Superset SQL Lab (любая БД, подключённая к Superset) или XLSX-файлов.
# @LAYER Domain
# @RELATION DEPENDS_ON -> backend.core.superset_client
# @RELATION DEPENDS_ON -> [EXT:Library:pandas]
# @PUBLIC_API  DatasetMapper
from typing import Any

import pandas as pd  # type: ignore

from ..logger import belief_scope, logger as app_logger


# #region DatasetMapper [TYPE Class]
# @BRIEF Класс для маппинга и обновления verbose_name в датасетах Superset.
class DatasetMapper:
    # #region __init__ [TYPE Function]
    # @PURPOSE: Initializes the mapper.
    # @POST  Объект DatasetMapper инициализирован.
    def __init__(self):
        pass
    # #endregion __init__

    # #region get_sqllab_mappings [TYPE Function]
    # @PURPOSE: Извлекает маппинги column_name -> verbose_name через SQL Lab Superset.
    # @PRE  sqllab_executor должен быть инициализирован с database_id.
    # @PRE  dataset_id должен существовать в Superset.
    # @POST  Возвращается словарь column_name -> verbose_name из результатов SQL-запроса.
    #   Если sql_query не указан, строится запрос к information_schema.columns.
    # @PARAM  client (SupersetClient) - Авторизованный клиент Superset.
    # @PARAM  dataset_id (int) - ID датасета (для получения table_name/schema).
    # @PARAM  sqllab_executor (SupersetSqlLabExecutor) - Исполнитель SQL Lab.
    # @PARAM  database_id (int) - ID базы данных в Superset.
    # @PARAM  sql_query (Optional[str]) - Произвольный SQL-запрос (должен вернуть column_name + verbose_name).
    # @RETURN  Dict[str, str] - Словарь column_name -> verbose_name.
    def get_sqllab_mappings(
        self,
        client: Any,
        dataset_id: int,
        sqllab_executor: Any,
        database_id: int,
        sql_query: str | None = None,
    ) -> dict[str, str]:
        with belief_scope("Fetch mappings via SQL Lab"):
            # Получаем датасет из Superset, чтобы узнать table_name и schema
            dataset_response = client.get_dataset(dataset_id)
            dataset_data = dataset_response.get("result", dataset_response)
            table_name = dataset_data.get("table_name")
            table_schema = dataset_data.get("schema") or "public"
            database_name = dataset_data.get("database", {}).get("database_name", "")

            app_logger.info(
                "[get_sqllab_mappings] Dataset %d: %s.%s (db: %s)",
                dataset_id, table_schema, table_name, database_name
            )

            # Если SQL-запрос не указан, строим дефолтный для information_schema
            if not sql_query:
                sql_query = (
                    "SELECT column_name, pg_catalog.col_description(\n"
                    "    (quote_ident(table_schema) || '.' || quote_ident(table_name))::regclass::oid,\n"
                    "    ordinal_position\n"
                    ") AS verbose_name\n"
                    f"FROM information_schema.columns\n"
                    f"WHERE table_schema = '{table_schema}' AND table_name = '{table_name}';\n"
                )
                app_logger.info("[get_sqllab_mappings] Using default information_schema query for %s.%s", table_schema, table_name)
            else:
                app_logger.info("[get_sqllab_mappings] Using custom SQL query")

            # Выполняем запрос через SQL Lab
            app_logger.info("[get_sqllab_mappings] Executing SQL Lab query on database %d...", database_id)
            result = sqllab_executor.execute_and_poll(sql_query, database_id)
            rows = result.get("results") or result.get("data") or result.get("result", [])
            app_logger.info("[get_sqllab_mappings] SQL Lab returned %d rows", len(rows))

            mappings: dict[str, str] = {}
            for row in rows:
                col_name = row.get("column_name")
                verbose_name = row.get("verbose_name")
                if col_name and verbose_name:
                    mappings[col_name] = str(verbose_name)

            app_logger.info("[get_sqllab_mappings] Extracted %d mappings", len(mappings))
            return mappings
    # #endregion get_sqllab_mappings

    # #region load_excel_mappings [TYPE Function]
    # @PURPOSE: Загружает маппинги column_name -> verbose_name из XLSX файла.
    # @PRE  file_path должен указывать на существующий XLSX файл.
    # @POST  Возвращается словарь с маппингами из файла.
    # @THROW: Exception - При ошибках чтения файла или парсинга.
    # @PARAM  file_path (str) - Путь к XLSX файлу.
    # @RETURN  Dict[str, str] - Словарь с маппингами.
    def load_excel_mappings(self, file_path: str) -> dict[str, str]:
        with belief_scope("Load mappings from Excel"):
            app_logger.info("[load_excel_mappings][Enter] Loading mappings from %s.", file_path)
            try:
                df = pd.read_excel(file_path)
                mappings = df.set_index('column_name')['verbose_name'].to_dict()
                app_logger.info("[load_excel_mappings][Success] Loaded %d mappings.", len(mappings))
                return mappings
            except Exception as e:
                app_logger.error("[load_excel_mappings][Failure] %s", e, exc_info=True)
                raise
    # #endregion load_excel_mappings

    # #region run_mapping [TYPE Function]
    # @PURPOSE: Основная функция для выполнения маппинга и обновления verbose_name датасета в Superset.
    # @PRE  superset_client должен быть авторизован.
    # @PRE  dataset_id должен быть существующим ID в Superset.
    # @POST  Если найдены изменения, датасет в Superset обновлен через API.
    # @RELATION CALLS -> [EXT:method:DatasetMapper.get_sqllab_mappings]
    # @RELATION CALLS -> [EXT:method:DatasetMapper.load_excel_mappings]
    # @RELATION CALLS -> [EXT:method:SupersetClient.get_dataset]
    # @RELATION CALLS -> [EXT:method:SupersetClient.update_dataset]
    # @PARAM  superset_client (Any) - Клиент Superset.
    # @PARAM  dataset_id (int) - ID датасета для обновления.
    # @PARAM  source (str) - Источник данных ('sqllab' или 'excel').
    # @PARAM  sqllab_executor (Optional[Any]) - Исполнитель SQL Lab (для sqllab source).
    # @PARAM  database_id (Optional[int]) - ID базы данных Superset (для sqllab source).
    # @PARAM  sql_query (Optional[str]) - Произвольный SQL-запрос (для sqllab source).
    # @PARAM  excel_path (Optional[str]) - Путь к XLSX файлу.
    def run_mapping(
        self,
        superset_client: Any,
        dataset_id: int,
        source: str,
        sqllab_executor: Any = None,
        database_id: int | None = None,
        sql_query: str | None = None,
        excel_path: str | None = None,
    ):
        with belief_scope(f"Run dataset mapping for ID {dataset_id}"):
            app_logger.info("[run_mapping][Enter] Starting dataset mapping for ID %d from source '%s'.", dataset_id, source)
            mappings: dict[str, str] = {}

            try:
                if source == "sqllab":
                    assert sqllab_executor and database_id, "sqllab_executor and database_id are required for sqllab source."
                    mappings.update(self.get_sqllab_mappings(superset_client, dataset_id, sqllab_executor, database_id, sql_query))
                elif source == "excel":
                    assert excel_path, "excel_path is required."
                    mappings.update(self.load_excel_mappings(excel_path))
                else:
                    app_logger.error("[run_mapping][Failure] Invalid source: %s.", source)
                    return

                dataset_response = superset_client.get_dataset(dataset_id)
                dataset_data = dataset_response['result']

                original_columns = dataset_data.get('columns', [])
                updated_columns = []
                changes_made = False
                for column in original_columns:
                    col_name = column.get('column_name')

                    new_column = {
                        "column_name": col_name,
                        "id": column.get("id"),
                        "advanced_data_type": column.get("advanced_data_type"),
                        "description": column.get("description"),
                        "expression": column.get("expression"),
                        "extra": column.get("extra"),
                        "filterable": column.get("filterable"),
                        "groupby": column.get("groupby"),
                        "is_active": column.get("is_active"),
                        "is_dttm": column.get("is_dttm"),
                        "python_date_format": column.get("python_date_format"),
                        "type": column.get("type"),
                        "uuid": column.get("uuid"),
                        "verbose_name": column.get("verbose_name"),
                    }

                    new_column = {k: v for k, v in new_column.items() if v is not None}
                    if col_name in mappings:
                        mapping_value = mappings[col_name]
                        if isinstance(mapping_value, str) and new_column.get('verbose_name') != mapping_value:
                            new_column['verbose_name'] = mapping_value
                            changes_made = True

                    updated_columns.append(new_column)

                updated_metrics = []
                for metric in dataset_data.get("metrics", []):
                    new_metric = {
                        "id": metric.get("id"),
                        "metric_name": metric.get("metric_name"),
                        "expression": metric.get("expression"),
                        "verbose_name": metric.get("verbose_name"),
                        "description": metric.get("description"),
                        "d3format": metric.get("d3format"),
                        "currency": metric.get("currency"),
                        "extra": metric.get("extra"),
                        "warning_text": metric.get("warning_text"),
                        "metric_type": metric.get("metric_type"),
                        "uuid": metric.get("uuid"),
                    }
                    updated_metrics.append({k: v for k, v in new_metric.items() if v is not None})

                if changes_made:
                    payload_for_update = {
                        "database_id": dataset_data.get("database", {}).get("id"),
                        "table_name": dataset_data.get("table_name"),
                        "schema": dataset_data.get("schema"),
                        "columns": updated_columns,
                        "owners": [owner["id"] for owner in dataset_data.get("owners", [])],
                        "metrics": updated_metrics,
                        "extra": dataset_data.get("extra"),
                        "description": dataset_data.get("description"),
                        "sql": dataset_data.get("sql"),
                        "cache_timeout": dataset_data.get("cache_timeout"),
                        "catalog": dataset_data.get("catalog"),
                        "default_endpoint": dataset_data.get("default_endpoint"),
                        "external_url": dataset_data.get("external_url"),
                        "fetch_values_predicate": dataset_data.get("fetch_values_predicate"),
                        "filter_select_enabled": dataset_data.get("filter_select_enabled"),
                        "is_managed_externally": dataset_data.get("is_managed_externally"),
                        "is_sqllab_view": dataset_data.get("is_sqllab_view"),
                        "main_dttm_col": dataset_data.get("main_dttm_col"),
                        "normalize_columns": dataset_data.get("normalize_columns"),
                        "offset": dataset_data.get("offset"),
                        "template_params": dataset_data.get("template_params"),
                    }

                    payload_for_update = {k: v for k, v in payload_for_update.items() if v is not None}
                    superset_client.update_dataset(dataset_id, payload_for_update)
                    app_logger.info("[run_mapping][Success] Dataset %d columns' verbose_name updated.", dataset_id)
                else:
                    app_logger.info("[run_mapping][State] No changes in columns' verbose_name, skipping update.")
            except (AssertionError, FileNotFoundError, Exception) as e:
                app_logger.error("[run_mapping][Failure] %s", e, exc_info=True)
                return
    # #endregion run_mapping
# #endregion DatasetMapper
# #endregion DatasetMapperModule
