# 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.

# distutils: language = c++

from pyarrow.includes.common cimport *
from pyarrow.includes.libarrow cimport *


cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
    cdef char* CTracingServerMiddlewareName\
        " arrow::flight::TracingServerMiddleware::kMiddlewareName"

    cdef cppclass CActionType" arrow::flight::ActionType":
        c_string type
        c_string description
        bint operator==(CActionType)
        CResult[c_string] SerializeToString()

        @staticmethod
        CResult[CActionType] Deserialize(const c_string& serialized)

    cdef cppclass CAction" arrow::flight::Action":
        c_string type
        shared_ptr[CBuffer] body
        bint operator==(CAction)
        CResult[c_string] SerializeToString()
        c_string ToString()

        @staticmethod
        CResult[CAction] Deserialize(const c_string& serialized)

    cdef cppclass CFlightResult" arrow::flight::Result":
        CFlightResult()
        CFlightResult(CFlightResult)
        shared_ptr[CBuffer] body
        bint operator==(CFlightResult)
        CResult[c_string] SerializeToString()
        c_string ToString()

        @staticmethod
        CResult[CFlightResult] Deserialize(const c_string& serialized)

    cdef cppclass CBasicAuth" arrow::flight::BasicAuth":
        CBasicAuth()
        CBasicAuth(CBuffer)
        CBasicAuth(CBasicAuth)
        c_string username
        c_string password
        bint operator==(CBasicAuth)
        CResult[c_string] SerializeToString()
        c_string ToString()

        @staticmethod
        CResult[CBasicAuth] Deserialize(const c_string& serialized)

    cdef cppclass CResultStream" arrow::flight::ResultStream":
        CResult[unique_ptr[CFlightResult]] Next()

    cdef cppclass CDescriptorType \
            " arrow::flight::FlightDescriptor::DescriptorType":
        bint operator==(CDescriptorType)

    CDescriptorType CDescriptorTypeUnknown\
        " arrow::flight::FlightDescriptor::UNKNOWN"
    CDescriptorType CDescriptorTypePath\
        " arrow::flight::FlightDescriptor::PATH"
    CDescriptorType CDescriptorTypeCmd\
        " arrow::flight::FlightDescriptor::CMD"

    cdef cppclass CFlightDescriptor" arrow::flight::FlightDescriptor":
        CDescriptorType type
        c_string cmd
        vector[c_string] path
        bint operator==(CFlightDescriptor)
        CResult[c_string] SerializeToString()
        c_string ToString()

        @staticmethod
        CResult[CFlightDescriptor] Deserialize(const c_string& serialized)

    cdef cppclass CTicket" arrow::flight::Ticket":
        CTicket()
        c_string ticket
        bint operator==(CTicket)
        CResult[c_string] SerializeToString()
        c_string ToString()

        @staticmethod
        CResult[CTicket] Deserialize(const c_string& serialized)

    cdef cppclass CCriteria" arrow::flight::Criteria":
        CCriteria()
        c_string expression
        bint operator==(CCriteria)
        CResult[c_string] SerializeToString()

        @staticmethod
        CResult[CCriteria] Deserialize(const c_string& serialized)

    cdef cppclass CLocation" arrow::flight::Location":
        CLocation()
        c_string ToString()
        c_bool Equals(const CLocation& other)

        @staticmethod
        CResult[CLocation] Parse(const c_string& uri_string)

        @staticmethod
        CResult[CLocation] ForGrpcTcp(const c_string& host, int port)

        @staticmethod
        CResult[CLocation] ForGrpcTls(const c_string& host, int port)

        @staticmethod
        CResult[CLocation] ForGrpcUnix(const c_string& path)

    cdef cppclass CFlightEndpoint" arrow::flight::FlightEndpoint":
        CFlightEndpoint()

        CTicket ticket
        vector[CLocation] locations

        bint operator==(CFlightEndpoint)
        CResult[c_string] SerializeToString()
        c_string ToString()

        @staticmethod
        CResult[CFlightEndpoint] Deserialize(const c_string& serialized)

    cdef cppclass CFlightInfo" arrow::flight::FlightInfo":
        CFlightInfo(CFlightInfo info)
        int64_t total_records()
        int64_t total_bytes()
        CResult[shared_ptr[CSchema]] GetSchema(CDictionaryMemo* memo)
        CFlightDescriptor& descriptor()
        const vector[CFlightEndpoint]& endpoints()
        CResult[c_string] SerializeToString()
        c_string ToString()
        bint operator==(CFlightInfo)

        @staticmethod
        CResult[unique_ptr[CFlightInfo]] Deserialize(
            const c_string& serialized)

    cdef cppclass CSchemaResult" arrow::flight::SchemaResult":
        CSchemaResult()
        CSchemaResult(CSchemaResult result)
        CResult[shared_ptr[CSchema]] GetSchema(CDictionaryMemo* memo)
        bint operator==(CSchemaResult)
        CResult[c_string] SerializeToString()
        c_string ToString()

        @staticmethod
        CResult[CSchemaResult] Deserialize(const c_string& serialized)

    cdef cppclass CFlightListing" arrow::flight::FlightListing":
        CResult[unique_ptr[CFlightInfo]] Next()

    cdef cppclass CSimpleFlightListing" arrow::flight::SimpleFlightListing":
        # This doesn't work with Cython >= 3
        # CSimpleFlightListing(vector[CFlightInfo]&& info)
        CSimpleFlightListing(const vector[CFlightInfo]& info)

    cdef cppclass CFlightPayload" arrow::flight::FlightPayload":
        shared_ptr[CBuffer] descriptor
        shared_ptr[CBuffer] app_metadata
        CIpcPayload ipc_message

    cdef cppclass CFlightDataStream" arrow::flight::FlightDataStream":
        shared_ptr[CSchema] schema()
        CResult[CFlightPayload] Next()

    cdef cppclass CFlightStreamChunk" arrow::flight::FlightStreamChunk":
        CFlightStreamChunk()
        shared_ptr[CRecordBatch] data
        shared_ptr[CBuffer] app_metadata

    cdef cppclass CMetadataRecordBatchReader \
            " arrow::flight::MetadataRecordBatchReader":
        CResult[shared_ptr[CSchema]] GetSchema()
        CResult[CFlightStreamChunk] Next()
        CResult[shared_ptr[CTable]] ToTable()

    CResult[shared_ptr[CRecordBatchReader]] MakeRecordBatchReader\
        " arrow::flight::MakeRecordBatchReader"(
            shared_ptr[CMetadataRecordBatchReader])

    cdef cppclass CMetadataRecordBatchWriter \
            " arrow::flight::MetadataRecordBatchWriter"(CRecordBatchWriter):
        CStatus Begin(shared_ptr[CSchema] schema,
                      const CIpcWriteOptions& options)
        CStatus WriteMetadata(shared_ptr[CBuffer] app_metadata)
        CStatus WriteWithMetadata(const CRecordBatch& batch,
                                  shared_ptr[CBuffer] app_metadata)

    cdef cppclass CFlightStreamReader \
            " arrow::flight::FlightStreamReader"(CMetadataRecordBatchReader):
        void Cancel()
        CResult[shared_ptr[CTable]] ToTableWithStopToken" ToTable"\
            (const CStopToken& stop_token)

    cdef cppclass CFlightMessageReader \
            " arrow::flight::FlightMessageReader"(CMetadataRecordBatchReader):
        CFlightDescriptor& descriptor()

    cdef cppclass CFlightMessageWriter \
            " arrow::flight::FlightMessageWriter"(CMetadataRecordBatchWriter):
        pass

    cdef cppclass CFlightStreamWriter \
            " arrow::flight::FlightStreamWriter"(CMetadataRecordBatchWriter):
        CStatus DoneWriting()

    cdef cppclass CRecordBatchStream \
            " arrow::flight::RecordBatchStream"(CFlightDataStream):
        CRecordBatchStream(shared_ptr[CRecordBatchReader]& reader,
                           const CIpcWriteOptions& options)

    cdef cppclass CFlightMetadataReader" arrow::flight::FlightMetadataReader":
        CStatus ReadMetadata(shared_ptr[CBuffer]* out)

    cdef cppclass CFlightMetadataWriter" arrow::flight::FlightMetadataWriter":
        CStatus WriteMetadata(const CBuffer& message)

    cdef cppclass CServerAuthReader" arrow::flight::ServerAuthReader":
        CStatus Read(c_string* token)

    cdef cppclass CServerAuthSender" arrow::flight::ServerAuthSender":
        CStatus Write(c_string& token)

    cdef cppclass CClientAuthReader" arrow::flight::ClientAuthReader":
        CStatus Read(c_string* token)

    cdef cppclass CClientAuthSender" arrow::flight::ClientAuthSender":
        CStatus Write(c_string& token)

    cdef cppclass CServerAuthHandler" arrow::flight::ServerAuthHandler":
        pass

    cdef cppclass CClientAuthHandler" arrow::flight::ClientAuthHandler":
        pass

    cdef cppclass CServerCallContext" arrow::flight::ServerCallContext":
        c_string& peer_identity()
        c_string& peer()
        c_bool is_cancelled()
        void AddHeader(const c_string& key, const c_string& value)
        void AddTrailer(const c_string& key, const c_string& value)
        CServerMiddleware* GetMiddleware(const c_string& key)

    cdef cppclass CTimeoutDuration" arrow::flight::TimeoutDuration":
        CTimeoutDuration(double)

    cdef cppclass CFlightCallOptions" arrow::flight::FlightCallOptions":
        CFlightCallOptions()
        CTimeoutDuration timeout
        CIpcWriteOptions write_options
        CIpcReadOptions read_options
        vector[pair[c_string, c_string]] headers
        CStopToken stop_token

    cdef cppclass CCertKeyPair" arrow::flight::CertKeyPair":
        CCertKeyPair()
        c_string pem_cert
        c_string pem_key

    cdef cppclass CFlightMethod" arrow::flight::FlightMethod":
        bint operator==(CFlightMethod)

    CFlightMethod CFlightMethodInvalid\
        " arrow::flight::FlightMethod::Invalid"
    CFlightMethod CFlightMethodHandshake\
        " arrow::flight::FlightMethod::Handshake"
    CFlightMethod CFlightMethodListFlights\
        " arrow::flight::FlightMethod::ListFlights"
    CFlightMethod CFlightMethodGetFlightInfo\
        " arrow::flight::FlightMethod::GetFlightInfo"
    CFlightMethod CFlightMethodGetSchema\
        " arrow::flight::FlightMethod::GetSchema"
    CFlightMethod CFlightMethodDoGet\
        " arrow::flight::FlightMethod::DoGet"
    CFlightMethod CFlightMethodDoPut\
        " arrow::flight::FlightMethod::DoPut"
    CFlightMethod CFlightMethodDoAction\
        " arrow::flight::FlightMethod::DoAction"
    CFlightMethod CFlightMethodListActions\
        " arrow::flight::FlightMethod::ListActions"
    CFlightMethod CFlightMethodDoExchange\
        " arrow::flight::FlightMethod::DoExchange"

    cdef cppclass CCallInfo" arrow::flight::CallInfo":
        CFlightMethod method

    # This is really std::unordered_multimap, but Cython has no
    # bindings for it, so treat it as an opaque class and bind the
    # methods we need
    cdef cppclass CCallHeaders" arrow::flight::CallHeaders":
        cppclass const_iterator:
            pair[c_string, c_string] operator*()
            # For Cython < 3
            const_iterator operator++()
            # For Cython >= 3
            const_iterator operator++(int)
            bint operator==(const_iterator)
            bint operator!=(const_iterator)
        const_iterator cbegin()
        const_iterator cend()

    cdef cppclass CAddCallHeaders" arrow::flight::AddCallHeaders":
        void AddHeader(const c_string& key, const c_string& value)

    cdef cppclass CServerMiddleware" arrow::flight::ServerMiddleware":
        c_string name()

    cdef cppclass CServerMiddlewareFactory\
            " arrow::flight::ServerMiddlewareFactory":
        pass

    cdef cppclass CClientMiddleware" arrow::flight::ClientMiddleware":
        pass

    cdef cppclass CClientMiddlewareFactory\
            " arrow::flight::ClientMiddlewareFactory":
        pass

    cpdef cppclass CTracingServerMiddlewareTraceKey\
            " arrow::flight::TracingServerMiddleware::TraceKey":
        CTracingServerMiddlewareTraceKey()
        c_string key
        c_string value