# 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.
#
import pandas
import pyarrow as pa

from tsfile.row_record import RowRecord
from tsfile.schema import TableSchema as TableSchemaPy
from tsfile.schema import TimeseriesSchema as TimeseriesSchemaPy, DeviceSchema as DeviceSchemaPy
from tsfile.tablet import Tablet as TabletPy
from libc.string cimport memset
from libc.stdint cimport uintptr_t
from .tsfile_cpp cimport *
from .tsfile_py_cpp cimport *

#cython: language_level=3

cdef class TsFileWriterPy:
    cdef TsFileWriter writer

    def __init__(self, pathname: str, memory_threshold: int = 128 * 1024 * 1024):
        self.writer = tsfile_writer_new_c(pathname, memory_threshold)

    def register_timeseries(self, device_name : str, timeseries_schema : TimeseriesSchemaPy):
        """
        Register a timeseries with tsfile writer.
        device_name: device name of the timeseries
        timeseries_schema: measurement's name/datatype/encoding/compressor
        """
        cdef TimeseriesSchema * c_schema = to_c_timeseries_schema(timeseries_schema)
        cdef ErrorCode errno
        try:
            errno = tsfile_writer_register_timeseries_py_cpp(self.writer, device_name, c_schema)
            check_error(errno)
        finally:
            free_c_timeseries_schema(c_schema)

    def register_device(self, device_schema : DeviceSchemaPy):
        """
        Register a device with tsfile writer.
        device_schema: the device definition, including device_name, some measurements' schema.
        """
        cdef DeviceSchema * device_schema_c = to_c_device_schema(device_schema)
        cdef ErrorCode errno
        try:
            errno = tsfile_writer_register_device_py_cpp(self.writer, device_schema_c)
            check_error(errno)
        finally:
            free_c_device_schema(device_schema_c)

    def register_table(self, table_schema : TableSchemaPy):
        """
        Register a table with tsfile writer.
        table_schema: the table definition, include table_name, columns' schema.
        """
        cdef TableSchema * c_schema = to_c_table_schema(table_schema)
        cdef ErrorCode errno
        try:
            errno = tsfile_writer_register_table_py_cpp(self.writer, c_schema)
            check_error(errno)
        finally:
            free_c_table_schema(c_schema)

    def write_tablet(self, tablet : TabletPy):
        """
        Write a tablet into tsfile with tsfile writer.
        tablet: data collection to be inserted
        Currently used for writing data into a device.
        """
        cdef Tablet ctablet = to_c_tablet(tablet)
        cdef ErrorCode errno
        try:
            errno = _tsfile_writer_write_tablet(self.writer, ctablet)
            check_error(errno)
        finally:
            free_c_tablet(ctablet)

    def write_dataframe(self, target_table: str, dataframe: pandas.DataFrame, tableschema: TableSchemaPy):
        cdef Tablet ctablet = dataframe_to_c_tablet(target_table, dataframe, tableschema)
        cdef ErrorCode errno
        try:
            errno = _tsfile_writer_write_table(self.writer, ctablet)
            check_error(errno)
        finally:
            free_c_tablet(ctablet)

    def write_row_record(self, record : RowRecord):
        """
        Write a record into tsfile with tsfile writer.
        record: timestamp and data collection
        Currently used for writing a row data into a device.
        """
        cdef TsRecord record_c = to_c_record(record)
        cdef ErrorCode errno
        try:
            errno = _tsfile_writer_write_ts_record(self.writer, record_c)
            check_error(errno)
        finally:
            free_c_row_record(record_c)

    def write_table(self, tablet : TabletPy):
        """
        Write a tablet data into a table in tsfile writer.
        Currently used for writing data into table.
        """
        cdef Tablet ctablet = to_c_tablet(tablet)
        cdef ErrorCode errno
        try:
            errno = _tsfile_writer_write_table(self.writer, ctablet)
            check_error(errno)
        finally:
            free_c_tablet(ctablet)

    def write_arrow_batch(self, table_name: str, data, time_col_index: int = -1):
        """
        Write an Arrow RecordBatch or Table into tsfile using Arrow C Data
        Interface for efficient batch writing without Python-level row loops.
        table_name: target table name (must be registered)
        data: pyarrow.RecordBatch or pyarrow.Table
        time_col_index: index of the time column in the Arrow schema.
            >= 0: use the specified column as the time column.
            <  0: auto-detect by Arrow timestamp type (default).
        """
        if isinstance(data, pa.Table):
            data = data.combine_chunks().to_batches()
            if not data:
                return
            data = data[0]

        cdef ArrowArray arrow_array
        cdef ArrowSchema arrow_schema
        cdef ErrorCode errno
        memset(&arrow_array, 0, sizeof(ArrowArray))
        memset(&arrow_schema, 0, sizeof(ArrowSchema))

        cdef uintptr_t array_ptr = <uintptr_t>&arrow_array
        cdef uintptr_t schema_ptr = <uintptr_t>&arrow_schema
        data._export_to_c(array_ptr, schema_ptr)

        cdef bytes tname = table_name.lower().encode('utf-8')
        try:
            errno = _tsfile_writer_write_arrow_table(
                self.writer, tname, &arrow_array, &arrow_schema, time_col_index)
            check_error(errno)
        finally:
            if arrow_array.release != NULL:
                arrow_array.release(&arrow_array)
            if arrow_schema.release != NULL:
                arrow_schema.release(&arrow_schema)

    cpdef close(self):
        """
        Flush data and Close tsfile writer.
        """
        cdef ErrorCode errno
        if self.writer == NULL:
            return
        errno = _tsfile_writer_close(self.writer)
        check_error(errno)
        self.writer = NULL

    cpdef flush(self):
        """
        Flush data in memory to tsfile.
        Writer will not be closed and can still be used
        """
        cdef ErrorCode errno
        if self.writer == NULL:
            return
        errno = _tsfile_writer_flush(self.writer)
        check_error(errno)

    def __dealloc__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
