import duckdb

from contextlib import contextmanager
from typing import Any, AnyStr, ClassVar, Iterator, Optional, Sequence
from dlt.common.destination import DestinationCapabilitiesContext

from dlt.destinations.exceptions import DatabaseTerminalException, DatabaseTransientException, DatabaseUndefinedRelation
from dlt.destinations.typing import DBApi, DBApiCursor, DBTransaction
from dlt.destinations.sql_client import SqlClientBase, raise_database_error, raise_open_connection_error

from dlt.destinations.duckdb import capabilities
from dlt.destinations.duckdb.configuration import DuckDbCredentials


class DuckDbSqlClient(SqlClientBase[duckdb.DuckDBPyConnection], DBTransaction):

    dbapi: ClassVar[DBApi] = duckdb
    capabilities: ClassVar[DestinationCapabilitiesContext] = capabilities()

    def __init__(self, dataset_name: str, credentials: DuckDbCredentials) -> None:
        super().__init__(dataset_name)
        self._conn: duckdb.DuckDBPyConnection = None
        self.credentials = credentials

    @raise_open_connection_error
    def open_connection(self) -> None:
        self._conn = self.credentials.borrow_conn(read_only=self.credentials.read_only)
        # TODO: apply config settings from credentials
        config={
            "search_path": self.fully_qualified_dataset_name(),
            "TimeZone": "UTC"
        }
        if config:
            for k, v in config.items():
                try:
                    # TODO: serialize str and ints, dbapi args do not work here
                    # TODO: enable various extensions ie. parquet
                    self._conn.execute(f"SET {k} = '{v}'")
                except (duckdb.CatalogException, duckdb.BinderException):
                    pass

    def close_connection(self) -> None:
        if self._conn:
            self.credentials.return_conn(self._conn)
            self._conn = None

    @contextmanager
    @raise_database_error
    def begin_transaction(self) -> Iterator[DBTransaction]:
        try:
            self._conn.begin()
            yield self
            self.commit_transaction()
        except Exception:
            # in some cases duckdb rollback the transaction automatically
            try:
                self.rollback_transaction()
            except DatabaseTransientException:
                pass
            raise

    @raise_database_error
    def commit_transaction(self) -> None:
        self._conn.commit()

    @raise_database_error
    def rollback_transaction(self) -> None:
        self._conn.rollback()

    @property
    def native_connection(self) -> duckdb.DuckDBPyConnection:
        return self._conn

    def has_dataset(self) -> bool:
        query = """
                SELECT 1
                    FROM INFORMATION_SCHEMA.SCHEMATA
                    WHERE schema_name = %s;
                """
        rows = self.execute_sql(query, self.fully_qualified_dataset_name(escape=False))
        return len(rows) > 0

    def create_dataset(self) -> None:
        self.execute_sql("CREATE SCHEMA %s" % self.fully_qualified_dataset_name())

    def drop_dataset(self) -> None:
        self.execute_sql("DROP SCHEMA %s CASCADE;" % self.fully_qualified_dataset_name())

    def execute_sql(self, sql: AnyStr, *args: Any, **kwargs: Any) -> Optional[Sequence[Sequence[Any]]]:
        with self.execute_query(sql, *args, **kwargs) as curr:
            if curr.description is None:
                return None
            else:
                f = curr.fetchall()
                return f

    @contextmanager
    @raise_database_error
    def execute_query(self, query: AnyStr, *args: Any, **kwargs: Any) -> Iterator[DBApiCursor]:
        assert isinstance(query, str)
        db_args = args if args else kwargs if kwargs else None
        if db_args:
            # TODO: must provide much better refactoring of params
            query = query.replace("%s", "?")
        try:
            self._conn.execute(query, db_args)
            yield self._conn  # type: ignore
        except duckdb.Error as outer:
            self.close_connection()
            self.open_connection()
            raise outer

    # def execute_fragments(self, batch: Sequence[AnyStr], *args: Any, **kwargs: Any) -> Optional[Sequence[Sequence[Any]]]:
    #     # execute in a loop to avoid rewrites
    #     results = []
    #     print(batch)
    #     for sql in batch:
    #         print(f"executing in dudckdb: {sql}")
    #         result = self.execute_sql(sql, args, kwargs)
    #         if result:
    #             results.extend(result)
    #     if results:
    #         return results
    #     else:
    #         return None

    def fully_qualified_dataset_name(self, escape: bool = True) -> str:
        return self.capabilities.escape_identifier(self.dataset_name) if escape else self.dataset_name

    @classmethod
    def _make_database_exception(cls, ex: Exception) -> Exception:
        if isinstance(ex, (duckdb.CatalogException)):  # duckdb.errors.UndefinedTable, duckdb.errors.InvalidSchemaName
            raise DatabaseUndefinedRelation(ex)  # duckdb.errors.SyntaxError, psycopg2.errors.UndefinedFunction
        elif isinstance(ex, duckdb.InvalidInputException):
            # duckdb raises TypeError on malformed query parameters
            return DatabaseTransientException(duckdb.ProgrammingError(ex))
        elif isinstance(ex, (duckdb.OperationalError, duckdb.InternalError, duckdb.SyntaxException, duckdb.ParserException)):
            term = cls._maybe_make_terminal_exception_from_data_error(ex)
            if term:
                return term
            else:
                return DatabaseTransientException(ex)
        elif isinstance(ex, (duckdb.DataError, duckdb.ProgrammingError, duckdb.IntegrityError)):
            return DatabaseTerminalException(ex)
        elif cls.is_dbapi_exception(ex):
            return DatabaseTransientException(ex)
        else:
            return ex

    @staticmethod
    def _maybe_make_terminal_exception_from_data_error(pg_ex: duckdb.Error) -> Optional[Exception]:
        return None

    @staticmethod
    def is_dbapi_exception(ex: Exception) -> bool:
        return isinstance(ex, duckdb.Error)
