# -*- coding: utf-8 -*-
from __future__ import annotations

import contextlib
import re
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Mapping,
    MutableMapping,
    Optional,
    Pattern,
    Tuple,
    Type,
    Union,
    cast,
)

import botocore
from sqlalchemy import exc, schema, text, types, util
from sqlalchemy.engine import Engine, reflection
from sqlalchemy.engine.default import DefaultDialect
from sqlalchemy.engine.interfaces import ExecutionContext
from sqlalchemy.sql.compiler import (
    DDLCompiler,
    GenericTypeCompiler,
    IdentifierPreparer,
    SQLCompiler,
)

import pyathena
from pyathena.sqlalchemy.compiler import (
    AthenaDDLCompiler,
    AthenaStatementCompiler,
    AthenaTypeCompiler,
)
from pyathena.sqlalchemy.preparer import AthenaDMLIdentifierPreparer
from pyathena.sqlalchemy.types import TINYINT, AthenaDate, AthenaStruct, AthenaTimestamp
from pyathena.sqlalchemy.util import _HashableDict
from pyathena.util import strtobool

if TYPE_CHECKING:
    from types import ModuleType

    from sqlalchemy import (
        URL,
        ClauseElement,
        Connection,
        PoolProxiedConnection,
    )
    from sqlalchemy.engine.interfaces import (
        ReflectedForeignKeyConstraint,
        ReflectedIndex,
        ReflectedPrimaryKeyConstraint,
    )
    from sqlalchemy.sql.schema import SchemaItem


ischema_names: Dict[str, Type[Any]] = {
    "boolean": types.BOOLEAN,
    "float": types.FLOAT,
    # TODO: types.DOUBLE is not defined in SQLAlchemy 1.4.
    "double": types.FLOAT,
    "real": types.FLOAT,
    "tinyint": TINYINT,
    "smallint": types.SMALLINT,
    "integer": types.INTEGER,
    "int": types.INTEGER,
    "bigint": types.BIGINT,
    "decimal": types.DECIMAL,
    "char": types.CHAR,
    "varchar": types.VARCHAR,
    "string": types.String,
    "date": types.DATE,
    "timestamp": types.TIMESTAMP,
    "binary": types.BINARY,
    "varbinary": types.BINARY,
    "array": types.String,
    "map": types.String,
    "struct": AthenaStruct,
    "row": AthenaStruct,
    "json": types.String,
}


class AthenaDialect(DefaultDialect):
    name: str = "awsathena"
    preparer: Type[IdentifierPreparer] = AthenaDMLIdentifierPreparer
    statement_compiler: Type[SQLCompiler] = AthenaStatementCompiler
    ddl_compiler: Type[DDLCompiler] = AthenaDDLCompiler
    type_compiler: Type[GenericTypeCompiler] = AthenaTypeCompiler
    default_paramstyle: str = pyathena.paramstyle
    cte_follows_insert: bool = True
    supports_alter: bool = False
    supports_pk_autoincrement: Optional[bool] = False
    supports_default_values: bool = False
    supports_empty_insert: bool = False
    supports_multivalues_insert: bool = True
    supports_native_decimal: bool = True
    supports_native_boolean: bool = True
    supports_unicode_statements: Optional[bool] = True
    supports_unicode_binds: Optional[bool] = True
    supports_statement_cache: bool = True
    returns_unicode_strings: Optional[bool] = True
    description_encoding: Optional[bool] = None
    postfetch_lastrowid: bool = False
    construct_arguments: Optional[
        List[Tuple[Type[Union["SchemaItem", "ClauseElement"]], Mapping[str, Any]]]
    ] = [
        (
            schema.Table,
            {
                "location": None,
                "compression": None,
                "row_format": None,
                "file_format": None,
                "serdeproperties": None,
                "tblproperties": None,
                "bucket_count": None,
            },
        ),
        (
            schema.Column,
            {
                "partition": False,
                "partition_transform": None,
                "partition_transform_bucket_count": None,
                "partition_transform_truncate_length": None,
                "cluster": False,
            },
        ),
    ]

    colspecs = {
        types.DATE: AthenaDate,
        types.DATETIME: AthenaTimestamp,
        types.TIMESTAMP: AthenaTimestamp,
    }

    ischema_names: Dict[str, Type[Any]] = ischema_names

    _connect_options: Dict[str, Any] = {}  # type: ignore
    _pattern_column_type: Pattern[str] = re.compile(r"^([a-zA-Z]+)(?:$|[\(|<](.+)[\)|>]$)")

    def __init__(self, json_deserializer=None, json_serializer=None, **kwargs):
        DefaultDialect.__init__(self, **kwargs)
        self._json_deserializer = json_deserializer
        self._json_serializer = json_serializer

    @classmethod
    def import_dbapi(cls) -> "ModuleType":
        return pyathena

    @classmethod
    def dbapi(cls) -> "ModuleType":  # type: ignore
        return pyathena

    def _raw_connection(self, connection: Union[Engine, "Connection"]) -> "PoolProxiedConnection":
        if isinstance(connection, Engine):
            return connection.raw_connection()
        return connection.connection

    def create_connect_args(self, url: "URL") -> Tuple[Tuple[str], MutableMapping[str, Any]]:
        # Connection string format:
        #   awsathena+rest://
        #   {aws_access_key_id}:{aws_secret_access_key}@athena.{region_name}.amazonaws.com:443/
        #   {schema_name}?s3_staging_dir={s3_staging_dir}&...
        self._connect_options = self._create_connect_args(url)
        return cast(Tuple[str], ()), self._connect_options

    def _create_connect_args(self, url: "URL") -> Dict[str, Any]:
        opts: Dict[str, Any] = {
            "aws_access_key_id": url.username if url.username else None,
            "aws_secret_access_key": url.password if url.password else None,
            "region_name": re.sub(
                r"^athena\.([a-z0-9-]+)\.amazonaws\.(com|com.cn)$", r"\1", url.host
            )
            if url.host
            else None,
            "schema_name": url.database if url.database else "default",
        }
        opts.update(url.query)
        if "verify" in opts:
            verify = opts["verify"]
            # If a ValueError occurs, it is probably the file name of the CA certificate being used.
            with contextlib.suppress(ValueError):
                verify = bool(strtobool(verify))
            opts.update({"verify": verify})
        if "duration_seconds" in opts:
            opts.update({"duration_seconds": int(opts["duration_seconds"])})
        if "poll_interval" in opts:
            opts.update({"poll_interval": float(opts["poll_interval"])})
        if "kill_on_interrupt" in opts:
            opts.update({"kill_on_interrupt": bool(strtobool(opts["kill_on_interrupt"]))})
        if "result_reuse_enable" in opts:
            opts.update({"result_reuse_enable": bool(strtobool(opts["result_reuse_enable"]))})
        if "result_reuse_minutes" in opts:
            opts.update({"result_reuse_minutes": int(opts["result_reuse_minutes"])})
        return opts

    @reflection.cache
    def _get_schemas(self, connection, **kw):
        raw_connection = self._raw_connection(connection)
        catalog = raw_connection.catalog_name  # type: ignore
        with raw_connection.driver_connection.cursor() as cursor:  # type: ignore
            try:
                return cursor.list_databases(catalog)
            except pyathena.error.OperationalError as e:
                cause = e.__cause__
                if (
                    isinstance(cause, botocore.exceptions.ClientError)
                    and cause.response["Error"]["Code"] == "InvalidRequestException"
                ):
                    return []
                raise

    @reflection.cache
    def _get_table(self, connection, table_name: str, schema: Optional[str] = None, **kw):
        raw_connection = self._raw_connection(connection)
        schema = schema if schema else raw_connection.schema_name  # type: ignore
        with raw_connection.driver_connection.cursor() as cursor:  # type: ignore
            try:
                return cursor.get_table_metadata(table_name, schema_name=schema, logging_=False)
            except pyathena.error.OperationalError as e:
                cause = e.__cause__
                if (
                    isinstance(cause, botocore.exceptions.ClientError)
                    and cause.response["Error"]["Code"] == "MetadataException"
                ):
                    raise exc.NoSuchTableError(table_name) from e
                raise

    @reflection.cache
    def _get_tables(self, connection, schema: Optional[str] = None, **kw):
        raw_connection = self._raw_connection(connection)
        schema = schema if schema else raw_connection.schema_name  # type: ignore
        with raw_connection.driver_connection.cursor() as cursor:  # type: ignore
            return cursor.list_table_metadata(schema_name=schema)

    def get_schema_names(self, connection, **kw):
        schemas = self._get_schemas(connection, **kw)
        return [s.name for s in schemas]

    def get_table_names(self, connection: "Connection", schema: Optional[str] = None, **kw):
        # Tables created by Athena are always classified as `EXTERNAL_TABLE`,
        # but Athena can also query tables classified as `MANAGED_TABLE`, `EXTERNAL`, or `customer`.
        # Managed Tables are created by default when creating tables via Spark when
        # Glue has been enabled as the Hive Metastore for Elastic Map Reduce (EMR) clusters.
        # With Athena Federation, tables in the database that are connected to Athena via lambda
        # function, is classified as `EXTERNAL` and fully queryable
        tables = self._get_tables(connection, schema, **kw)
        return [
            t.name
            for t in tables
            if t.table_type in ["EXTERNAL_TABLE", "MANAGED_TABLE", "EXTERNAL", "customer"]
        ]

    def get_view_names(self, connection: "Connection", schema: Optional[str] = None, **kw):
        tables = self._get_tables(connection, schema, **kw)
        return [t.name for t in tables if t.table_type == "VIRTUAL_VIEW"]

    def get_table_comment(
        self, connection: "Connection", table_name: str, schema: Optional[str] = None, **kw
    ):
        metadata = self._get_table(connection, table_name, schema=schema, **kw)
        return {"text": metadata.comment}

    def get_table_options(
        self, connection: "Connection", table_name: str, schema: Optional[str] = None, **kw
    ):
        metadata = self._get_table(connection, table_name, schema=schema, **kw)
        # TODO The metadata retrieved from the API does not seem to include bucketing information.
        return {
            "awsathena_location": metadata.location,
            "awsathena_compression": metadata.compression,
            "awsathena_row_format": metadata.row_format,
            "awsathena_file_format": metadata.file_format,
            "awsathena_serdeproperties": _HashableDict(metadata.serde_properties),
            "awsathena_tblproperties": _HashableDict(metadata.table_properties),
        }

    def has_table(
        self, connection: "Connection", table_name: str, schema: Optional[str] = None, **kw
    ):
        try:
            columns = self.get_columns(connection, table_name, schema)
            return bool(columns)
        except exc.NoSuchTableError:
            return False

    @reflection.cache
    def get_view_definition(
        self, connection: Connection, view_name: str, schema: Optional[str] = None, **kw
    ):
        raw_connection = self._raw_connection(connection)
        schema = schema if schema else raw_connection.schema_name  # type: ignore
        query = f"""SHOW CREATE VIEW "{schema}"."{view_name}";"""
        try:
            res = connection.scalars(text(query))
        except exc.OperationalError as e:
            raise exc.NoSuchTableError(f"{schema}.{view_name}") from e
        else:
            return "\n".join(res)

    @reflection.cache
    def get_columns(
        self, connection: "Connection", table_name: str, schema: Optional[str] = None, **kw
    ):
        metadata = self._get_table(connection, table_name, schema=schema, **kw)
        columns = [
            {
                "name": c.name,
                "type": self._get_column_type(c.type),
                "nullable": True,
                "default": None,
                "autoincrement": False,
                "comment": c.comment,
                "dialect_options": {"awsathena_partition": None},
            }
            for c in metadata.columns
        ]
        columns += [
            {
                "name": c.name,
                "type": self._get_column_type(c.type),
                "nullable": True,
                "default": None,
                "autoincrement": False,
                "comment": c.comment,
                "dialect_options": {"awsathena_partition": True},
            }
            for c in metadata.partition_keys
        ]
        return columns

    def _get_column_type(self, type_: str):
        match = self._pattern_column_type.match(type_)
        if match:
            name = match.group(1).lower()
            length = match.group(2)
        else:
            name = type_.lower()
            length = None

        if name in self.ischema_names:
            col_type = self.ischema_names[name]
        else:
            util.warn(f"Did not recognize type '{type_}'")
            col_type = types.NullType

        args = []
        if length:
            if col_type is types.DECIMAL:
                precision, scale = length.split(",")
                args = [int(precision), int(scale)]
            elif col_type is types.CHAR or col_type is types.VARCHAR:
                args = [int(length)]

        return col_type(*args)

    def get_foreign_keys(
        self, connection: "Connection", table_name: str, schema: Optional[str] = None, **kw
    ) -> List["ReflectedForeignKeyConstraint"]:
        # Athena has no support for foreign keys.
        return []  # pragma: no cover

    def get_pk_constraint(
        self, connection: "Connection", table_name: str, schema: Optional[str] = None, **kw
    ) -> "ReflectedPrimaryKeyConstraint":
        # Athena has no support for primary keys.
        return {"name": None, "constrained_columns": []}  # pragma: no cover

    def get_indexes(
        self, connection: "Connection", table_name: str, schema: Optional[str] = None, **kw
    ) -> List["ReflectedIndex"]:
        # Athena has no support for indexes.
        return []  # pragma: no cover

    def do_execute(self, cursor, statement, parameters, context=None):
        on_start_query_execution = None
        if isinstance(context, ExecutionContext):
            execution_options = context.execution_options
            if execution_options is not None:
                on_start_query_execution = execution_options.get("on_start_query_execution")

        if on_start_query_execution is not None:
            cursor.execute(statement, parameters, on_start_query_execution=on_start_query_execution)
        else:
            cursor.execute(statement, parameters)

    def do_rollback(self, dbapi_connection: "PoolProxiedConnection") -> None:
        # No transactions for Athena
        pass  # pragma: no cover

    def _check_unicode_returns(
        self, connection: "Connection", additional_tests: Optional[List[Any]] = None
    ) -> bool:
        # Requests gives back Unicode strings
        return True  # pragma: no cover

    def _check_unicode_description(self, connection: "Connection") -> bool:
        # Requests gives back Unicode strings
        return True  # pragma: no cover
