# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
from textwrap import dedent
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union
from urllib.parse import unquote_plus

from sqlalchemy import exc, sql
from sqlalchemy.engine import Engine
from sqlalchemy.engine.base import Connection
from sqlalchemy.engine.default import DefaultDialect, DefaultExecutionContext
from sqlalchemy.engine.url import URL

from trino import dbapi as trino_dbapi
from trino import logging
from trino.auth import BasicAuthentication, CertificateAuthentication, JWTAuthentication
from trino.dbapi import Cursor
from trino.sqlalchemy import compiler, datatype, error

logger = logging.get_logger(__name__)


class TrinoDialect(DefaultDialect):
    name = "trino"
    driver = "rest"

    statement_compiler = compiler.TrinoSQLCompiler
    ddl_compiler = compiler.TrinoDDLCompiler
    type_compiler = compiler.TrinoTypeCompiler
    preparer = compiler.TrinoIdentifierPreparer

    # Data Type
    supports_native_enum = False
    supports_native_boolean = True
    supports_native_decimal = True

    # Column options
    supports_sequences = False
    supports_comments = True
    inline_comments = True
    supports_default_values = False

    # DDL
    supports_alter = True

    # DML
    # Queries of the form `INSERT () VALUES ()` is not supported by Trino.
    supports_empty_insert = False
    supports_multivalues_insert = True
    postfetch_lastrowid = False

    # Caching
    # Warnings are generated by SQLAlchmey if this flag is not explicitly set
    # and tests are needed before being enabled
    supports_statement_cache = False

    # Support proper ordering of CTEs in regard to an INSERT statement
    cte_follows_insert = True

    @classmethod
    def dbapi(cls):
        """
        ref: https://www.python.org/dev/peps/pep-0249/#module-interface
        """
        return trino_dbapi

    def create_connect_args(self, url: URL) -> Tuple[Sequence[Any], Mapping[str, Any]]:
        args: Sequence[Any] = list()
        kwargs: Dict[str, Any] = dict(host=url.host)

        if url.port:
            kwargs["port"] = url.port

        db_parts = (url.database or "system").split("/")
        if len(db_parts) == 1:
            kwargs["catalog"] = unquote_plus(db_parts[0])
        elif len(db_parts) == 2:
            kwargs["catalog"] = unquote_plus(db_parts[0])
            kwargs["schema"] = unquote_plus(db_parts[1])
        else:
            raise ValueError(f"Unexpected database format {url.database}")

        if url.username:
            kwargs["user"] = unquote_plus(url.username)

        if url.password:
            if not url.username:
                raise ValueError("Username is required when specify password in connection URL")
            kwargs["http_scheme"] = "https"
            kwargs["auth"] = BasicAuthentication(unquote_plus(url.username), unquote_plus(url.password))

        if "access_token" in url.query:
            kwargs["http_scheme"] = "https"
            kwargs["auth"] = JWTAuthentication(unquote_plus(url.query["access_token"]))

        if "cert" and "key" in url.query:
            kwargs["http_scheme"] = "https"
            kwargs["auth"] = CertificateAuthentication(unquote_plus(url.query['cert']), unquote_plus(url.query['key']))

        if "source" in url.query:
            kwargs["source"] = unquote_plus(url.query["source"])
        else:
            kwargs["source"] = "trino-sqlalchemy"

        if "session_properties" in url.query:
            kwargs["session_properties"] = json.loads(unquote_plus(url.query["session_properties"]))

        if "http_headers" in url.query:
            kwargs["http_headers"] = json.loads(unquote_plus(url.query["http_headers"]))

        if "extra_credential" in url.query:
            kwargs["extra_credential"] = [
                tuple(extra_credential) for extra_credential in json.loads(unquote_plus(url.query["extra_credential"]))
            ]

        if "client_tags" in url.query:
            kwargs["client_tags"] = json.loads(unquote_plus(url.query["client_tags"]))

        if "legacy_primitive_types" in url.query:
            kwargs["legacy_primitive_types"] = json.loads(unquote_plus(url.query["legacy_primitive_types"]))

        if "legacy_prepared_statements" in url.query:
            kwargs["legacy_prepared_statements"] = json.loads(unquote_plus(url.query["legacy_prepared_statements"]))

        if "verify" in url.query:
            kwargs["verify"] = json.loads(unquote_plus(url.query["verify"]))

        if "roles" in url.query:
            kwargs["roles"] = json.loads(url.query["roles"])

        return args, kwargs

    def get_columns(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]:
        if not self.has_table(connection, table_name, schema):
            raise exc.NoSuchTableError(f"schema={schema}, table={table_name}")
        return self._get_columns(connection, table_name, schema, **kw)

    def _get_columns(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]:
        schema = schema or self._get_default_schema_name(connection)
        query = dedent(
            """
            SELECT
                "column_name",
                "data_type",
                "column_default",
                UPPER("is_nullable") AS "is_nullable"
            FROM "information_schema"."columns"
            WHERE "table_schema" = :schema
              AND "table_name" = :table
            ORDER BY "ordinal_position" ASC
        """
        ).strip()
        res = connection.execute(sql.text(query), {"schema": schema, "table": table_name})
        columns = []
        for record in res:
            column = dict(
                name=record.column_name,
                type=datatype.parse_sqltype(record.data_type),
                nullable=record.is_nullable == "YES",
                default=record.column_default,
            )
            columns.append(column)
        return columns

    def get_pk_constraint(self, connection: Connection, table_name: str, schema: str = None, **kw) -> Dict[str, Any]:
        """Trino has no support for primary keys. Returns a dummy"""
        return dict(name=None, constrained_columns=[])

    def get_primary_keys(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[str]:
        pk = self.get_pk_constraint(connection, table_name, schema)
        return pk.get("constrained_columns")  # type: ignore

    def get_foreign_keys(
        self, connection: Connection, table_name: str, schema: str = None, **kw
    ) -> List[Dict[str, Any]]:
        """Trino has no support for foreign keys. Returns an empty list."""
        return []

    def get_schema_names(self, connection: Connection, **kw) -> List[str]:
        query = dedent(
            """
            SELECT "schema_name"
            FROM "information_schema"."schemata"
        """
        ).strip()
        res = connection.execute(sql.text(query))
        return [row.schema_name for row in res]

    def get_table_names(self, connection: Connection, schema: str = None, **kw) -> List[str]:
        schema = schema or self._get_default_schema_name(connection)
        if schema is None:
            raise exc.NoSuchTableError("schema is required")
        query = dedent(
            """
            SELECT "table_name"
            FROM "information_schema"."tables"
            WHERE "table_schema" = :schema
              AND "table_type" = 'BASE TABLE'
        """
        ).strip()
        res = connection.execute(sql.text(query), {"schema": schema})
        return [row.table_name for row in res]

    def get_temp_table_names(self, connection: Connection, schema: str = None, **kw) -> List[str]:
        """Trino has no support for temporary tables. Returns an empty list."""
        return []

    def get_view_names(self, connection: Connection, schema: str = None, **kw) -> List[str]:
        schema = schema or self._get_default_schema_name(connection)
        if schema is None:
            raise exc.NoSuchTableError("schema is required")

        # Querying the information_schema.views table is subpar as it compiles the view definitions.
        query = dedent(
            """
            SELECT "table_name"
            FROM "information_schema"."tables"
            WHERE "table_schema" = :schema
              AND "table_type" = 'VIEW'
        """
        ).strip()
        res = connection.execute(sql.text(query), {"schema": schema})
        return [row.table_name for row in res]

    def get_temp_view_names(self, connection: Connection, schema: str = None, **kw) -> List[str]:
        """Trino has no support for temporary views. Returns an empty list."""
        return []

    def get_view_definition(self, connection: Connection, view_name: str, schema: str = None, **kw) -> str:
        schema = schema or self._get_default_schema_name(connection)
        if schema is None:
            raise exc.NoSuchTableError("schema is required")
        query = dedent(
            """
            SELECT "view_definition"
            FROM "information_schema"."views"
            WHERE "table_schema" = :schema
              AND "table_name" = :view
        """
        ).strip()
        res = connection.execute(sql.text(query), {"schema": schema, "view": view_name})
        return res.scalar()

    def get_indexes(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]:
        if not self.has_table(connection, table_name, schema):
            raise exc.NoSuchTableError(f"schema={schema}, table={table_name}")

        partitioned_columns = self._get_columns(connection, f"{table_name}$partitions", schema, **kw)
        if not partitioned_columns:
            return []
        partition_index = dict(
            name="partition",
            column_names=[col["name"] for col in partitioned_columns],
            unique=False
        )
        return [partition_index]

    def get_sequence_names(self, connection: Connection, schema: str = None, **kw) -> List[str]:
        """Trino has no support for sequences. Returns an empty list."""
        return []

    def get_unique_constraints(
        self, connection: Connection, table_name: str, schema: str = None, **kw
    ) -> List[Dict[str, Any]]:
        """Trino has no support for unique constraints. Returns an empty list."""
        return []

    def get_check_constraints(
        self, connection: Connection, table_name: str, schema: str = None, **kw
    ) -> List[Dict[str, Any]]:
        """Trino has no support for check constraints. Returns an empty list."""
        return []

    def get_table_comment(self, connection: Connection, table_name: str, schema: str = None, **kw) -> Dict[str, Any]:
        catalog_name = self._get_default_catalog_name(connection)
        if catalog_name is None:
            raise exc.NoSuchTableError("catalog is required in connection")
        schema_name = schema or self._get_default_schema_name(connection)
        if schema_name is None:
            raise exc.NoSuchTableError("schema is required")
        query = dedent(
            """
            SELECT "comment"
            FROM "system"."metadata"."table_comments"
            WHERE "catalog_name" = :catalog_name
              AND "schema_name" = :schema_name
              AND "table_name" = :table_name
        """
        ).strip()
        try:
            res = connection.execute(
                sql.text(query),
                {"catalog_name": catalog_name, "schema_name": schema_name, "table_name": table_name}
            )
            return dict(text=res.scalar())
        except error.TrinoQueryError as e:
            if e.error_name in (
                error.PERMISSION_DENIED,
            ):
                return dict(text=None)
            raise

    def has_schema(self, connection: Connection, schema: str) -> bool:
        query = dedent(
            """
            SELECT "schema_name"
            FROM "information_schema"."schemata"
            WHERE "schema_name" = :schema
        """
        ).strip()
        res = connection.execute(sql.text(query), {"schema": schema})
        return res.first() is not None

    def has_table(self, connection: Connection, table_name: str, schema: str = None, **kw) -> bool:
        schema = schema or self._get_default_schema_name(connection)
        if schema is None:
            return False
        query = dedent(
            """
            SELECT "table_name"
            FROM "information_schema"."tables"
            WHERE "table_schema" = :schema
              AND "table_name" = :table
        """
        ).strip()
        res = connection.execute(sql.text(query), {"schema": schema, "table": table_name})
        return res.first() is not None

    def has_sequence(self, connection: Connection, sequence_name: str, schema: str = None, **kw) -> bool:
        """Trino has no support for sequence. Returns False indicate that given sequence does not exists."""
        return False

    @classmethod
    def _get_server_version_info(cls, connection: Connection) -> Any:
        def get_server_version_info(_):
            query = "SELECT version()"
            try:
                res = connection.execute(sql.text(query))
                version = res.scalar()
                return tuple([version])
            except exc.ProgrammingError as e:
                logger.debug(f"Failed to get server version: {e.orig.message}")
                return None

        # Make server_version_info lazy in order to only make HTTP calls if user explicitly requests it.
        cls.server_version_info = property(get_server_version_info, lambda instance, value: None)

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

    def _get_default_catalog_name(self, connection: Connection) -> Optional[str]:
        dbapi_connection: trino_dbapi.Connection = self._raw_connection(connection)
        return dbapi_connection.catalog

    def _get_default_schema_name(self, connection: Connection) -> Optional[str]:
        dbapi_connection: trino_dbapi.Connection = self._raw_connection(connection)
        return dbapi_connection.schema

    def do_execute(
        self, cursor: Cursor, statement: str, parameters: Tuple[Any, ...], context: DefaultExecutionContext = None
    ):
        cursor.execute(statement, parameters)

    def do_rollback(self, dbapi_connection: trino_dbapi.Connection):
        if dbapi_connection.transaction is not None:
            dbapi_connection.rollback()

    def set_isolation_level(self, dbapi_conn: trino_dbapi.Connection, level: str) -> None:
        dbapi_conn._isolation_level = trino_dbapi.IsolationLevel[level]

    def get_isolation_level(self, dbapi_conn: trino_dbapi.Connection) -> str:
        return dbapi_conn.isolation_level.name

    def get_default_isolation_level(self, dbapi_conn: trino_dbapi.Connection) -> str:
        return trino_dbapi.IsolationLevel.AUTOCOMMIT.name

    def _get_full_table(self, table_name: str, schema: str = None, quote: bool = True) -> str:
        table_part = self.identifier_preparer.quote_identifier(table_name) if quote else table_name
        if schema:
            schema_part = self.identifier_preparer.quote_identifier(schema) if quote else schema
            return f"{schema_part}.{table_part}"

        return table_part
