import logging

from ...exceptions import NoDatabaseProvidedException
from ...utils import LocalStorage, write_summary
from ..abstract import (
    CATALOG_ASSETS,
    QUERIES_ASSETS,
    VIEWS_ASSETS,
    AssetsMapping,
    TimeFilter,
    WarehouseAsset,
    WarehouseAssetGroup,
    common_args,
    write_empty_asset,
)
from .client import GlueAthenaClient
from .constants import DEFAULT_USERS_LOOKBACK_DAYS
from .credentials import GlueAthenaCredentials

logger = logging.getLogger(__name__)

Paths = dict[str, str]

GLUE_ATHENA_ASSETS: AssetsMapping = {
    WarehouseAssetGroup.CATALOG: CATALOG_ASSETS,
    WarehouseAssetGroup.QUERY: QUERIES_ASSETS,
    WarehouseAssetGroup.ROLE: (WarehouseAsset.USER,),
    WarehouseAssetGroup.VIEW_DDL: VIEWS_ASSETS,
}


def _credentials(params: dict) -> GlueAthenaCredentials:
    """Resolve credentials for Glue/Athena access."""
    filtered = {
        key: value for key, value in params.items() if value is not None
    }
    return GlueAthenaCredentials(**filtered)


def _warn_and_write_empty_users(storage: LocalStorage, error: Exception) -> str:
    """
    Skip user extraction when CloudTrail access is unavailable.

    Glue/Athena users are inferred from query activity, so missing CloudTrail
    permissions should not prevent the rest of the extraction from completing.
    """
    logger.warning(
        "Athena user extraction skipped; continuing with other assets: %s",
        error,
    )
    return write_empty_asset(storage, WarehouseAsset.USER)


def _warn_and_write_empty_queries(
    storage: LocalStorage,
    error: Exception,
) -> str:
    """
    Skip query extraction when Athena access is unavailable.

    Query extraction is best-effort for Glue/Athena so catalog ingestion can
    still complete when workgroup permissions are missing or incomplete.
    """
    logger.warning(
        "Athena query extraction skipped; continuing with other assets: %s",
        error,
    )
    return write_empty_asset(storage, WarehouseAsset.QUERY)


class GlueAthenaExtractionProcessor:
    """Orchestrate Glue/Athena extraction into local storage."""

    def __init__(
        self,
        client: GlueAthenaClient,
        storage: LocalStorage,
    ) -> None:
        self._client = client
        self._storage = storage

    def extract_catalog(self) -> Paths:
        """Extract catalog assets from Glue."""
        catalog_locations: Paths = {}
        logger.info(
            "Starting Glue catalog extraction: databases, schemas, tables, and columns"
        )

        database_assets = self._client.databases()
        location = self._storage.put(
            WarehouseAsset.DATABASE.value, database_assets
        )
        catalog_locations[WarehouseAsset.DATABASE.value] = location
        logger.info(
            "Extracted %s databases to %s",
            len(database_assets),
            location,
        )

        schemas = self._client.schemas()
        if not schemas:
            raise NoDatabaseProvidedException

        location = self._storage.put(WarehouseAsset.SCHEMA.value, schemas)
        catalog_locations[WarehouseAsset.SCHEMA.value] = location
        logger.info(
            "Extracted %s schemas to %s",
            len(schemas),
            location,
        )
        logger.info(
            "Scanning %s Glue schemas to collect tables and columns",
            len(schemas),
        )

        tables, columns = self._client.tables_and_columns(schemas)

        location = self._storage.put(WarehouseAsset.TABLE.value, tables)
        catalog_locations[WarehouseAsset.TABLE.value] = location
        logger.info("Extracted %s tables to %s", len(tables), location)

        location = self._storage.put(WarehouseAsset.COLUMN.value, columns)
        catalog_locations[WarehouseAsset.COLUMN.value] = location
        logger.info("Extracted %s columns to %s", len(columns), location)
        return catalog_locations

    def extract_role(self, time_filter: TimeFilter | None = None) -> Paths:
        """
        Extract best-effort users inferred from Athena query activity.

        When no time filter is provided, the extraction defaults to yesterday,
        which matches the usual warehouse query ingestion behavior.
        """
        if time_filter is None:
            time_filter = TimeFilter.default()

        try:
            logger.info(
                "Extracting Athena users from CloudTrail over the last %s days up to %s",
                DEFAULT_USERS_LOOKBACK_DAYS,
                time_filter.day,
            )
            users = self._client.users(time_filter.day)
            location = self._storage.put(WarehouseAsset.USER.value, users)
            logger.info(
                "Extracted %s users from the last %s days to %s",
                len(users),
                DEFAULT_USERS_LOOKBACK_DAYS,
                location,
            )
        except Exception as error:
            # User extraction is best-effort and should not prevent the rest
            # of the Glue/Athena extraction from completing.
            location = _warn_and_write_empty_users(self._storage, error)
        return {WarehouseAsset.USER.value: location}

    def extract_query(self, time_filter: TimeFilter | None = None) -> Paths:
        """
        Extract Athena queries executed on the requested day.
        """
        if time_filter is None:
            time_filter = TimeFilter.default()

        try:
            workgroups = self._client.workgroups()
            if not workgroups:
                logger.info(
                    "Skipping Athena queries for %s: no accessible workgroups",
                    time_filter.day,
                )
                location = write_empty_asset(
                    self._storage,
                    WarehouseAsset.QUERY,
                )
                return {WarehouseAsset.QUERY.value: location}

            logger.info(
                "Extracting Athena queries for %s by scanning workgroups: %s",
                time_filter.day,
                ", ".join(workgroups),
            )
            queries = self._client.queries(
                time_filter.day,
                workgroups=workgroups,
            )
            location = self._storage.put(WarehouseAsset.QUERY.value, queries)
            logger.info("Extracted %s queries to %s", len(queries), location)
        except Exception as error:
            # Query extraction is best-effort and should not prevent the rest
            # of the Glue/Athena extraction from completing.
            location = _warn_and_write_empty_queries(self._storage, error)

        return {WarehouseAsset.QUERY.value: location}

    def extract_view_ddl(self) -> Paths:
        """
        Extract current view definitions from Glue metadata.

        This task is intentionally separate from query extraction because view
        definitions are catalog state, not day-scoped query activity.
        """
        logger.info(
            "Extracting Glue view DDL by scanning current catalog metadata"
        )
        view_ddl = self._client.view_ddl()
        location = self._storage.put(WarehouseAsset.VIEW_DDL.value, view_ddl)
        logger.info("Extracted %s view_ddl to %s", len(view_ddl), location)
        return {WarehouseAsset.VIEW_DDL.value: location}


def extract_all(**kwargs) -> None:
    """Run the Glue/Athena extraction flow."""
    output_directory, _ = common_args(kwargs)

    credentials = _credentials(kwargs)
    client = GlueAthenaClient(credentials=credentials)

    storage = LocalStorage(directory=output_directory)

    extractor = GlueAthenaExtractionProcessor(
        client=client,
        storage=storage,
    )
    skip_queries = kwargs.get("skip_queries") or False

    extractor.extract_catalog()
    extractor.extract_role()
    if not skip_queries:
        extractor.extract_query()
        extractor.extract_view_ddl()

    write_summary(
        output_directory,
        storage.stored_at_ts,
        client_name="glue_athena",
    )
