import json
import logging
import time
from datetime import date, datetime, timedelta
from typing import Any

import boto3  # type: ignore

from ...utils import SerializedAsset
from .constants import (
    DEFAULT_USERS_LOOKBACK_DAYS,
    GLUE_DATABASE_NAME,
)
from .credentials import GlueAthenaCredentials  # type: ignore
from .spark_schema import spark_field_descriptions
from .type_parser import flatten_glue_type
from .user import parse_cloudtrail_user
from .view import build_view_ddl_row

logger = logging.getLogger(__name__)

_S3_PREFIX = "s3://"
_S3_CONSOLE_BUCKETS_URL = "https://s3.console.aws.amazon.com/s3/buckets/"
_ATHENA_API_SLEEP_SECONDS = 0.05
_ATHENA_BATCH_GET_LIMIT = 50
_CLOUDTRAIL_PROGRESS_PAGE_INTERVAL = 100
_CLOUDTRAIL_QUERY_LOOKUP_ATTRIBUTES = [
    {
        "AttributeKey": "EventName",
        "AttributeValue": "StartQueryExecution",
    }
]


def _s3_to_http(url: str) -> str:
    """Convert s3:// links to the AWS console URL."""
    if not url.startswith(_S3_PREFIX):
        return url

    s3_path = url.removeprefix(_S3_PREFIX)
    parts = s3_path.split("/", 1)
    bucket = parts[0]
    object_name = parts[1] if len(parts) > 1 else ""
    return f"{_S3_CONSOLE_BUCKETS_URL}{bucket}/?prefix={object_name}"


def _format_table_description(
    location: str,
    description: str | None,
) -> str | None:
    """Append bucket location details to a table description."""

    if not location:
        return description or None

    url = _s3_to_http(location)
    location_md = f"Bucket location: [{location}]({url})"
    if not description:
        return location_md

    return f"{description}\n{location_md}"


def _table_type(table_type: str | None) -> str:
    """Normalize Glue table types into catalog-friendly labels."""
    if not table_type:
        return "TABLE"
    if table_type.upper() == "VIRTUAL_VIEW":
        return "VIEW"
    if table_type.upper() == "EXTERNAL_TABLE":
        return "EXTERNAL"
    return "TABLE"


def _column(
    table_id: str,
    column_name: str,
    column_type: str,
    comment: str | None,
    ordinal_position: int,
) -> dict:
    """Build a catalog column payload from Glue column data"""
    return {
        "column_id": f"{table_id}.{column_name}",
        "column_name": column_name,
        # Keep Glue raw type; mapping to Catalog ColumnDataType is done at store.
        "data_type": column_type,
        "description": comment,
        "ordinal_position": ordinal_position,
        "table_id": table_id,
        "tags": [],
    }


def _parse_table(
    raw_table: dict,
    schema_id: str,
    table_arn: str,
) -> dict:
    """Convert a Glue table into a catalog table payload."""
    table_name = raw_table["Name"]
    description = raw_table.get("Description")
    storage = raw_table.get("StorageDescriptor", {})
    location = storage["Location"]
    table_type = _table_type(raw_table.get("TableType"))
    return {
        "description": _format_table_description(location, description),
        "schema_id": schema_id,
        "table_id": table_arn,
        "table_name": table_name,
        "tags": [],
        "type": table_type,
    }


def _iter_table_columns(raw_table: dict, table_arn: str) -> SerializedAsset:
    """
    Build one ordered column list from storage columns and partition keys.

    Nested types are flattened into dotted paths.
    """
    columns: SerializedAsset = []
    ordinal_position = 0
    seen_columns: set[str] = set()
    spark_descriptions = spark_field_descriptions(
        raw_table.get("Parameters", {})
    )
    column_entries = (
        ("column", raw_table.get("StorageDescriptor", {}).get("Columns", [])),
        ("partition key", raw_table.get("PartitionKeys", [])),
    )

    for column_kind, raw_columns in column_entries:
        for raw_column in raw_columns:
            column_name = raw_column.get("Name")
            if not column_name:
                continue
            column_type = raw_column["Type"]

            # Expand nested types into a flattened list
            flattened = flatten_glue_type(column_name, column_type)
            for flattened_name, flattened_type, is_root in flattened:
                if flattened_name in seen_columns:
                    # union columns and partition keys can create duplicates
                    continue

                comment = raw_column.get("Comment") if is_root else None
                # Spark comments are preferred because Glue top-level comments
                # can be truncated while Spark keeps the full text.
                description = spark_descriptions.get(flattened_name) or comment
                column = _column(
                    table_arn,
                    flattened_name,
                    flattened_type,
                    description,
                    ordinal_position,
                )
                columns.append(column)
                seen_columns.add(flattened_name)
                ordinal_position += 1

    return columns


def _warehouse_user(query_user: dict[str, str | None]) -> dict | None:
    """
    Convert a query-level identity into a warehouse user payload.

    Athena and Glue do not expose a native endpoint to enumerate users, so the
    warehouse user list is inferred from the identities present in query logs.
    Users without a stable identifier are skipped because they cannot be joined
    reliably to query activity downstream.
    """
    user_id = query_user.get("user_id")
    if not user_id:
        return None

    return {
        "user_email": query_user.get("user_email"),
        "user_id": user_id,
        "user_name": query_user.get("user_name") or user_id,
    }


def _query_user_from_cloudtrail_event(
    event: dict,
) -> tuple[str, dict] | None:
    """Parse one CloudTrail event into a query id / user pair."""
    raw_cloudtrail_event = event.get("CloudTrailEvent")
    if not raw_cloudtrail_event:
        return None

    cloudtrail_event = json.loads(raw_cloudtrail_event)
    response_elements = cloudtrail_event.get("responseElements") or {}
    query_execution_id = response_elements.get("queryExecutionId")
    if not query_execution_id:
        return None

    return query_execution_id, parse_cloudtrail_user(cloudtrail_event)


def _query_from_execution(
    query_execution: dict,
    execution_date: date,
    query_user: dict[str, str | None] | None = None,
) -> dict | None:
    """
    Convert one Athena query execution into a warehouse query payload.

    Caller code is responsible for scanning Athena history newest-first and
    keeping only the executions that match the requested day.
    """
    status = query_execution.get("Status", {})
    start_time = status.get("SubmissionDateTime")
    end_time = status.get("CompletionDateTime")
    if not start_time or not end_time:
        return None

    if start_time.date() != execution_date:
        return None

    query_user = query_user or {}
    context = query_execution.get("QueryExecutionContext", {})
    return {
        "database_name": context.get("Database"),
        "end_time": end_time.isoformat(),
        "query_id": query_execution.get("QueryExecutionId"),
        "query_text": query_execution.get("Query"),
        "start_time": start_time.isoformat(),
        "user_email": query_user.get("user_email"),
        "user_id": query_user.get("user_id"),
        "user_name": query_user.get("user_name"),
    }


def _log_cloudtrail_progress(page_count: int, event_count: int) -> None:
    """Log periodic progress while scanning CloudTrail query events."""
    if page_count % _CLOUDTRAIL_PROGRESS_PAGE_INTERVAL != 0:
        return

    logger.info(
        "CloudTrail scan in progress: %s pages, %s Athena query events",
        page_count,
        event_count,
    )


class GlueAthenaClient:
    """
    Fetch Glue catalog metadata together with lightweight CloudTrail enrichments.
    """

    def __init__(
        self,
        credentials: GlueAthenaCredentials,
    ) -> None:
        self._region = credentials.aws_region
        self._account_id = credentials.aws_account_id

        session = boto3.session.Session(
            region_name=credentials.aws_region,
            aws_access_key_id=credentials.access_key_id,
            aws_secret_access_key=credentials.access_key_secret,
        )
        self._athena_client = session.client("athena")
        self._cloudtrail_client = session.client("cloudtrail")
        self._glue_client = session.client("glue")

    def databases(self) -> SerializedAsset:
        """
        Return the single AWS Data Catalog database asset.

        AWS exposes Glue through the built-in `awsdatacatalog` database:
        https://docs.aws.amazon.com/redshift/latest/mgmt/query-editor-v2-glue.html
        """
        return [
            {
                "database_id": GLUE_DATABASE_NAME,
                "database_name": GLUE_DATABASE_NAME,
            }
        ]

    @property
    def _glue_arn_prefix(self) -> str:
        return f"arn:aws:glue:{self._region}:{self._account_id}"

    def _resource_arn(self, resource: str, identifier: str) -> str:
        return f"{self._glue_arn_prefix}:{resource}/{identifier}"

    def schemas(self) -> SerializedAsset:
        """
        Return schema assets by calling Glue `get_databases()`.

        Naming is intentionally normalized to schemas here, even though the
        Glue endpoint is named "databases".
        """
        schemas: SerializedAsset = []
        paginator = self._glue_client.get_paginator("get_databases")
        for page in paginator.paginate():
            # Glue's `get_databases` endpoint returns logical schemas.
            for raw_schema in page.get("DatabaseList", []):
                schema_name = raw_schema["Name"]
                schema_arn = self._resource_arn("database", schema_name)
                schemas.append(
                    {
                        "database_id": GLUE_DATABASE_NAME,
                        "schema_id": f"{GLUE_DATABASE_NAME}.{schema_arn}",
                        "schema_name": schema_name,
                        "description": raw_schema.get("Description"),
                        "tags": [],
                    }
                )
        return schemas

    def tables_and_columns(
        self,
        schemas: SerializedAsset,
    ) -> tuple[SerializedAsset, SerializedAsset]:
        """Collect tables and columns for each schema."""
        tables: SerializedAsset = []
        columns: SerializedAsset = []

        paginator = self._glue_client.get_paginator("get_tables")
        total_schemas = len(schemas)
        for index, schema in enumerate(schemas, start=1):
            schema_name = schema["schema_name"]
            schema_id = schema["schema_id"]
            logger.info(
                f"Scanning Glue schema {index}/{total_schemas} "
                f"for tables and columns: {schema_name}"
            )

            page_iterator = paginator.paginate(DatabaseName=schema_name)
            for page in page_iterator:
                for raw_table in page.get("TableList", []):
                    table_arn = self._resource_arn(
                        "table",
                        f"{schema_name}/{raw_table['Name']}",
                    )
                    table = _parse_table(raw_table, schema_id, table_arn)
                    tables.append(table)

                    table_columns = _iter_table_columns(raw_table, table_arn)
                    columns.extend(table_columns)

        return tables, columns

    def view_ddl(self) -> SerializedAsset:
        """
        Retrieve all Glue/Athena view definitions from the catalog.
        """
        view_ddl: SerializedAsset = []
        schemas = self.schemas()
        if not schemas:
            return view_ddl

        paginator = self._glue_client.get_paginator("get_tables")
        total_schemas = len(schemas)
        for index, schema in enumerate(schemas, start=1):
            schema_name = schema["schema_name"]
            logger.info(
                f"Scanning Glue schema {index}/{total_schemas} "
                f"for view DDL: {schema_name}"
            )

            page_iterator = paginator.paginate(DatabaseName=schema_name)
            for page in page_iterator:
                for raw_table in page.get("TableList", []):
                    view_ddl_row = build_view_ddl_row(schema_name, raw_table)
                    if not view_ddl_row:
                        continue

                    view_ddl.append(view_ddl_row)

        return view_ddl

    def _iter_workgroup_pages(self):
        """Yield Athena workgroup pages, with a manual pagination fallback."""
        if self._athena_client.can_paginate("list_work_groups"):
            paginator = self._athena_client.get_paginator("list_work_groups")
            return paginator.paginate()

        # Some client runtimes expose the Athena operation without a paginator,
        # so we keep a manual fallback
        next_token: str | None = None
        pages: list[dict] = []
        while True:
            kwargs: dict[str, str] = {}
            if next_token:
                kwargs["NextToken"] = next_token

            response = self._athena_client.list_work_groups(**kwargs)
            pages.append(response)

            next_token = response.get("NextToken")
            if not next_token:
                break

            # Add a small pause to avoid burst throttling while paging.
            time.sleep(_ATHENA_API_SLEEP_SECONDS)

        return pages

    def workgroups(self) -> list[str]:
        """
        Return Athena workgroups visible to the current credentials.
        """
        workgroups: list[str] = []

        for page in self._iter_workgroup_pages():
            for workgroup in page.get("WorkGroups", []):
                if workgroup.get("State") != "ENABLED":
                    continue

                name = workgroup.get("Name")
                if name:
                    workgroups.append(name)

        return workgroups

    def _iter_cloudtrail_query_events(
        self,
        start_time: datetime,
        end_time: datetime,
    ) -> SerializedAsset:
        """
        Retrieve paginated CloudTrail `StartQueryExecution` events for a time
        window.
        """
        events: SerializedAsset = []
        next_token: str | None = None
        page_count = 0
        event_count = 0
        base_kwargs: dict[str, Any] = {
            "StartTime": start_time,
            "EndTime": end_time,
            "LookupAttributes": _CLOUDTRAIL_QUERY_LOOKUP_ATTRIBUTES,
        }

        while True:
            kwargs = dict(base_kwargs)
            if next_token:
                kwargs["NextToken"] = next_token

            response = self._cloudtrail_client.lookup_events(**kwargs)
            page_events = response.get("Events", [])
            events.extend(page_events)
            page_count += 1
            event_count += len(page_events)
            _log_cloudtrail_progress(page_count, event_count)

            next_token = response.get("NextToken")
            if not next_token:
                break

            # Add a small pause between CloudTrail pages for the same reason
            # as Athena: avoid burst throttling without slowing tests too much.
            time.sleep(_ATHENA_API_SLEEP_SECONDS)

        return events

    def _query_users(
        self,
        start_time: datetime,
        end_time: datetime,
    ) -> dict[str, dict]:
        """
        Retrieve user information for Athena queries from CloudTrail.
        """
        # The dict key is the Athena query execution id used to join the
        # CloudTrail user information back onto the Athena query payload.
        users_by_query_id: dict[str, dict] = {}
        for event in self._iter_cloudtrail_query_events(start_time, end_time):
            parsed_event = _query_user_from_cloudtrail_event(event)
            if not parsed_event:
                continue

            query_execution_id, query_user = parsed_event
            users_by_query_id[query_execution_id] = query_user

        return users_by_query_id

    def _workgroup_queries(
        self,
        workgroup: str,
        execution_date: date,
        users_by_query_id: dict[str, dict] | None = None,
    ) -> SerializedAsset:
        """
        Return queries for one Athena workgroup on the target day.

        Athena query history has no server-side date filter. We therefore scan
        `list_query_executions()` page by page and stop once we reach
        executions older than the requested day.

        The early stop below assumes Athena pagination is reverse
        chronological overall, not just within a single page. This matches our
        observed behavior, but is not explicitly guaranteed in the AWS docs.
        """
        queries: SerializedAsset = []
        next_token: str | None = None

        while True:
            kwargs: dict[str, str | int] = {
                "MaxResults": _ATHENA_BATCH_GET_LIMIT,
                "WorkGroup": workgroup,
            }
            if next_token:
                kwargs["NextToken"] = next_token

            response = self._athena_client.list_query_executions(**kwargs)
            query_ids = response.get("QueryExecutionIds", [])
            if not query_ids:
                break

            details = self._athena_client.batch_get_query_execution(
                QueryExecutionIds=query_ids
            )
            executions = sorted(
                details.get("QueryExecutions", []),
                key=lambda query: (
                    submission_date.timestamp()
                    if (
                        submission_date := query.get("Status", {}).get(
                            "SubmissionDateTime"
                        )
                    )
                    else float("-inf")
                ),
                reverse=True,
            )

            stop_workgroup = False
            for query_execution in executions:
                start_time = query_execution.get("Status", {}).get(
                    "SubmissionDateTime"
                )
                if not start_time:
                    continue

                execution_day = start_time.date()
                if execution_day > execution_date:
                    continue
                if execution_day < execution_date:
                    stop_workgroup = True
                    break

                query_id = query_execution.get("QueryExecutionId")
                query_user = None
                if query_id and users_by_query_id:
                    query_user = users_by_query_id.get(query_id)

                query = _query_from_execution(
                    query_execution,
                    execution_date,
                    query_user=query_user,
                )
                if query:
                    queries.append(query)

            # Stop early once we observe queries older than the requested day.
            # This relies on Athena returning paginated query history in
            # reverse chronological order across pages.
            if stop_workgroup:
                break

            next_token = response.get("NextToken")
            if not next_token:
                break

            # Add a small pause to avoid burst throttling while paging.
            time.sleep(_ATHENA_API_SLEEP_SECONDS)

        return queries

    def users(
        self,
        execution_date: date,
    ) -> SerializedAsset:
        """
        Build a best-effort user list from the identities seen in recent
        Athena query activity.

        Athena and Glue do not provide a dedicated "list users" endpoint for
        this use case. To keep the extractor reusable, users are inferred from
        the CloudTrail `StartQueryExecution` events already used to enrich
        query history, then deduplicated by stable user id.

        The extraction uses a rolling window instead of a single day to reduce
        snapshot yo-yo effects when source users are inferred from query logs.
        """
        end_time = datetime.combine(
            execution_date + timedelta(days=1),
            datetime.min.time(),
        )
        start_time = end_time - timedelta(days=DEFAULT_USERS_LOOKBACK_DAYS)
        users_by_id: dict[str, dict] = {}

        for query_user in self._query_users(start_time, end_time).values():
            user = _warehouse_user(query_user)
            if not user:
                continue

            user_id = user["user_id"]
            existing_user = users_by_id.get(user_id)
            if not existing_user:
                users_by_id[user_id] = user
                continue

            existing_name = existing_user.get("user_name")
            incoming_name = user.get("user_name")
            if (
                incoming_name
                and existing_name in (None, user_id)
                and incoming_name != user_id
            ):
                existing_user["user_name"] = user["user_name"]
            if not existing_user.get("user_email") and user.get("user_email"):
                existing_user["user_email"] = user["user_email"]

        return list(users_by_id.values())

    def queries(
        self,
        execution_date: date,
        workgroups: list[str] | None = None,
    ) -> SerializedAsset:
        """
        Return Athena queries executed on the requested day.

        When `workgroups` is not provided, all accessible enabled workgroups
        are scanned.
        """
        target_workgroups = (
            self.workgroups() if workgroups is None else workgroups
        )
        if not target_workgroups:
            return []

        start_time = datetime.combine(execution_date, datetime.min.time())
        end_time = start_time + timedelta(days=1)
        try:
            users_by_query_id = self._query_users(start_time, end_time)
        except Exception as error:
            logger.warning(
                "Athena query extraction could not enrich users from CloudTrail: %s",
                error,
            )
            users_by_query_id = {}

        queries: SerializedAsset = []
        total_workgroups = len(target_workgroups)
        for index, workgroup in enumerate(target_workgroups, start=1):
            logger.info(
                "Scanning Athena workgroup %s/%s for queries: %s",
                index,
                total_workgroups,
                workgroup,
            )
            queries.extend(
                self._workgroup_queries(
                    workgroup,
                    execution_date,
                    users_by_query_id,
                )
            )

        return queries
