# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

# cython: language_level = 3

"""Low-level ADBC API."""

import enum
import typing
from typing import List, Tuple

import cython
from cpython.bytes cimport PyBytes_FromStringAndSize
from libc.stdint cimport int32_t, int64_t, uint8_t, uint32_t, uintptr_t
from libc.string cimport memset
from libcpp.vector cimport vector as c_vector

if typing.TYPE_CHECKING:
    from typing import Self


cdef extern from "adbc.h" nogil:
    # C ABI
    cdef struct CArrowSchema"ArrowSchema":
        pass
    cdef struct CArrowArray"ArrowArray":
        pass
    cdef struct CArrowArrayStream"ArrowArrayStream":
        pass

    # ADBC
    ctypedef uint8_t CAdbcStatusCode"AdbcStatusCode"
    cdef CAdbcStatusCode ADBC_STATUS_OK
    cdef CAdbcStatusCode ADBC_STATUS_UNKNOWN
    cdef CAdbcStatusCode ADBC_STATUS_NOT_IMPLEMENTED
    cdef CAdbcStatusCode ADBC_STATUS_NOT_FOUND
    cdef CAdbcStatusCode ADBC_STATUS_ALREADY_EXISTS
    cdef CAdbcStatusCode ADBC_STATUS_INVALID_ARGUMENT
    cdef CAdbcStatusCode ADBC_STATUS_INVALID_STATE
    cdef CAdbcStatusCode ADBC_STATUS_INVALID_DATA
    cdef CAdbcStatusCode ADBC_STATUS_INTEGRITY
    cdef CAdbcStatusCode ADBC_STATUS_INTERNAL
    cdef CAdbcStatusCode ADBC_STATUS_IO
    cdef CAdbcStatusCode ADBC_STATUS_CANCELLED
    cdef CAdbcStatusCode ADBC_STATUS_TIMEOUT
    cdef CAdbcStatusCode ADBC_STATUS_UNAUTHENTICATED
    cdef CAdbcStatusCode ADBC_STATUS_UNAUTHORIZED

    cdef const char* ADBC_OPTION_VALUE_DISABLED
    cdef const char* ADBC_OPTION_VALUE_ENABLED

    cdef const char* ADBC_CONNECTION_OPTION_AUTOCOMMIT
    cdef const char* ADBC_INGEST_OPTION_TARGET_TABLE
    cdef const char* ADBC_INGEST_OPTION_MODE
    cdef const char* ADBC_INGEST_OPTION_MODE_APPEND
    cdef const char* ADBC_INGEST_OPTION_MODE_CREATE

    cdef int ADBC_OBJECT_DEPTH_ALL
    cdef int ADBC_OBJECT_DEPTH_CATALOGS
    cdef int ADBC_OBJECT_DEPTH_DB_SCHEMAS
    cdef int ADBC_OBJECT_DEPTH_TABLES
    cdef int ADBC_OBJECT_DEPTH_COLUMNS

    cdef uint32_t ADBC_INFO_VENDOR_NAME
    cdef uint32_t ADBC_INFO_VENDOR_VERSION
    cdef uint32_t ADBC_INFO_VENDOR_ARROW_VERSION
    cdef uint32_t ADBC_INFO_DRIVER_NAME
    cdef uint32_t ADBC_INFO_DRIVER_VERSION
    cdef uint32_t ADBC_INFO_DRIVER_ARROW_VERSION

    ctypedef void (*CAdbcErrorRelease)(CAdbcError*)
    ctypedef void (*CAdbcPartitionsRelease)(CAdbcPartitions*)

    cdef struct CAdbcError"AdbcError":
        char* message
        int32_t vendor_code
        char[5] sqlstate
        CAdbcErrorRelease release

    cdef struct CAdbcDriver"AdbcDriver":
        pass

    cdef struct CAdbcDatabase"AdbcDatabase":
        void* private_data

    cdef struct CAdbcConnection"AdbcConnection":
        void* private_data

    cdef struct CAdbcStatement"AdbcStatement":
        void* private_data

    cdef struct CAdbcPartitions"AdbcPartitions":
        size_t num_partitions
        const uint8_t** partitions
        const size_t* partition_lengths
        void* private_data
        CAdbcPartitionsRelease release

    CAdbcStatusCode AdbcDatabaseNew(CAdbcDatabase* database, CAdbcError* error)
    CAdbcStatusCode AdbcDatabaseSetOption(
        CAdbcDatabase* database,
        const char* key,
        const char* value,
        CAdbcError* error)
    CAdbcStatusCode AdbcDatabaseInit(CAdbcDatabase* database, CAdbcError* error)
    CAdbcStatusCode AdbcDatabaseRelease(CAdbcDatabase* database, CAdbcError* error)

    ctypedef void (*CAdbcDriverInitFunc "AdbcDriverInitFunc")(int, void*, CAdbcError*)
    CAdbcStatusCode AdbcDriverManagerDatabaseSetInitFunc(
        CAdbcDatabase* database,
        CAdbcDriverInitFunc init_func,
        CAdbcError* error)

    CAdbcStatusCode AdbcConnectionCommit(
        CAdbcConnection* connection,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionRollback(
        CAdbcConnection* connection,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionReadPartition(
        CAdbcConnection* connection,
        const uint8_t* serialized_partition,
        size_t serialized_length,
        CArrowArrayStream* out,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionGetInfo(
        CAdbcConnection* connection,
        uint32_t* info_codes,
        size_t info_codes_length,
        CArrowArrayStream* stream,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionGetObjects(
        CAdbcConnection* connection,
        int depth,
        const char* catalog,
        const char* db_schema,
        const char* table_name,
        const char** table_type,
        const char* column_name,
        CArrowArrayStream* stream,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionGetTableSchema(
        CAdbcConnection* connection,
        const char* catalog,
        const char* db_schema,
        const char* table_name,
        CArrowSchema* schema,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionGetTableTypes(
        CAdbcConnection* connection,
        CArrowArrayStream* stream,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionInit(
        CAdbcConnection* connection,
        CAdbcDatabase* database,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionNew(
        CAdbcConnection* connection,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionSetOption(
        CAdbcConnection* connection,
        const char* key,
        const char* value,
        CAdbcError* error)
    CAdbcStatusCode AdbcConnectionRelease(
        CAdbcConnection* connection,
        CAdbcError* error)

    CAdbcStatusCode AdbcStatementBind(
        CAdbcStatement* statement,
        CArrowArray*,
        CArrowSchema*,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementBindStream(
        CAdbcStatement* statement,
        CArrowArrayStream*,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementExecutePartitions(
        CAdbcStatement* statement,
        CArrowSchema* schema, CAdbcPartitions* partitions,
        int64_t* rows_affected,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementExecuteQuery(
        CAdbcStatement* statement,
        CArrowArrayStream* out, int64_t* rows_affected,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementNew(
        CAdbcConnection* connection,
        CAdbcStatement* statement,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementPrepare(
        CAdbcStatement* statement,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementSetOption(
        CAdbcStatement* statement,
        const char* key,
        const char* value,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementSetSqlQuery(
        CAdbcStatement* statement,
        const char* query,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementSetSubstraitPlan(
        CAdbcStatement* statement,
        const uint8_t* plan,
        size_t length,
        CAdbcError* error)
    CAdbcStatusCode AdbcStatementRelease(
        CAdbcStatement* statement,
        CAdbcError* error)


cdef extern from "adbc_driver_manager.h":
    const char* CAdbcStatusCodeMessage"AdbcStatusCodeMessage"(CAdbcStatusCode code)


class AdbcStatusCode(enum.IntEnum):
    """
    A status code indicating the type of error.
    """

    OK = ADBC_STATUS_OK
    UNKNOWN = ADBC_STATUS_UNKNOWN
    NOT_IMPLEMENTED = ADBC_STATUS_NOT_IMPLEMENTED
    NOT_FOUND = ADBC_STATUS_NOT_FOUND
    ALREADY_EXISTS = ADBC_STATUS_ALREADY_EXISTS
    INVALID_ARGUMENT = ADBC_STATUS_INVALID_ARGUMENT
    INVALID_STATE = ADBC_STATUS_INVALID_STATE
    INVALID_DATA = ADBC_STATUS_INVALID_DATA
    INTEGRITY = ADBC_STATUS_INTEGRITY
    INTERNAL = ADBC_STATUS_INTERNAL
    IO = ADBC_STATUS_IO
    CANCELLED = ADBC_STATUS_CANCELLED
    TIMEOUT = ADBC_STATUS_TIMEOUT
    UNAUTHENTICATED = ADBC_STATUS_UNAUTHENTICATED
    UNAUTHORIZED = ADBC_STATUS_UNAUTHORIZED


class AdbcInfoCode(enum.IntEnum):
    VENDOR_NAME = ADBC_INFO_VENDOR_NAME
    VENDOR_VERSION = ADBC_INFO_VENDOR_VERSION
    VENDOR_ARROW_VERSION = ADBC_INFO_VENDOR_ARROW_VERSION
    DRIVER_NAME = ADBC_INFO_DRIVER_NAME
    DRIVER_VERSION = ADBC_INFO_DRIVER_VERSION
    DRIVER_ARROW_VERSION = ADBC_INFO_DRIVER_ARROW_VERSION


class Warning(UserWarning):
    """
    PEP 249-compliant base warning class.
    """


class Error(Exception):
    """
    PEP 249-compliant base exception class.

    Attributes
    ----------
    status_code : AdbcStatusCode
        The original ADBC status code.
    vendor_code : int, optional
        A vendor-specific status code if present.
    sqlstate : str, optional
        The SQLSTATE code if present.
    """

    def __init__(self, message, *, status_code, vendor_code=None, sqlstate=None):
        super().__init__(message)
        self.status_code = AdbcStatusCode(status_code)
        self.vendor_code = None
        self.sqlstate = None


class InterfaceError(Error):
    """Errors related to the database interface."""


class DatabaseError(Error):
    """Errors related to the database."""


class DataError(DatabaseError):
    """Errors related to processed data."""


class OperationalError(DatabaseError):
    """Errors related to database operation, not under user control."""


class IntegrityError(DatabaseError):
    """Errors related to relational integrity."""


class InternalError(DatabaseError):
    """Errors related to database-internal errors."""


class ProgrammingError(DatabaseError):
    """Errors related to user errors."""


class NotSupportedError(DatabaseError):
    """An operation or some functionality is not supported."""

    def __init__(self, message, *, vendor_code=None, sqlstate=None):
        super().__init__(
            message,
            status_code=AdbcStatusCode.NOT_IMPLEMENTED,
            vendor_code=vendor_code,
            sqlstate=sqlstate,
        )


INGEST_OPTION_MODE = ADBC_INGEST_OPTION_MODE.decode("utf-8")
INGEST_OPTION_MODE_APPEND = ADBC_INGEST_OPTION_MODE_APPEND.decode("utf-8")
INGEST_OPTION_MODE_CREATE = ADBC_INGEST_OPTION_MODE_CREATE.decode("utf-8")
INGEST_OPTION_TARGET_TABLE = ADBC_INGEST_OPTION_TARGET_TABLE.decode("utf-8")


cdef void check_error(CAdbcStatusCode status, CAdbcError* error) except *:
    if status == ADBC_STATUS_OK:
        return

    message = CAdbcStatusCodeMessage(status).decode("utf-8")
    vendor_code = None
    sqlstate = None

    if error != NULL:
        if error.message != NULL:
            message += ": "