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

import pyarrow as pa
from typing import Tuple, Optional, List
from pathlib import Path
from abc import ABC, abstractmethod

from pypaimon.pynative.common.core_option import CoreOptions
from pypaimon.pynative.row.binary_row import BinaryRow
from pypaimon.pynative.table.data_file_meta import DataFileMeta


class DataWriter(ABC):
    """Base class for data writers that handle PyArrow tables directly."""

    def __init__(self, table, partition: Tuple, bucket: int):
        from pypaimon.pynative.table.file_store_table import FileStoreTable

        self.table: FileStoreTable = table
        self.partition = partition
        self.bucket = bucket

        self.file_io = self.table.file_io
        self.trimmed_primary_key_fields = self.table.table_schema.get_trimmed_primary_key_fields()

        options = self.table.options
        self.target_file_size = 256 * 1024 * 1024
        self.file_format = options.get(CoreOptions.FILE_FORMAT, CoreOptions.FILE_FORMAT_PARQUET)
        self.compression = options.get(CoreOptions.FILE_COMPRESSION, "zstd")

        self.pending_data: Optional[pa.RecordBatch] = None
        self.committed_files: List[DataFileMeta] = []

    def write(self, data: pa.RecordBatch):
        processed_data = self._process_data(data)

        if self.pending_data is None:
            self.pending_data = processed_data
        else:
            self.pending_data = self._merge_data(self.pending_data, processed_data)

        self._check_and_roll_if_needed()

    def prepare_commit(self) -> List[DataFileMeta]:
        if self.pending_data is not None and self.pending_data.num_rows > 0:
            self._write_data_to_file(self.pending_data)
            self.pending_data = None

        return self.committed_files.copy()

    def close(self):
        self.pending_data = None
        self.committed_files.clear()

    @abstractmethod
    def _process_data(self, data: pa.RecordBatch) -> pa.RecordBatch:
        """Process incoming data (e.g., add system fields, sort). Must be implemented by subclasses."""

    @abstractmethod
    def _merge_data(self, existing_data: pa.RecordBatch, new_data: pa.RecordBatch) -> pa.RecordBatch:
        """Merge existing data with new data. Must be implemented by subclasses."""

    def _check_and_roll_if_needed(self):
        if self.pending_data is None:
            return

        current_size = self.pending_data.get_total_buffer_size()
        if current_size > self.target_file_size:
            split_row = _find_optimal_split_point(self.pending_data, self.target_file_size)
            if split_row > 0:
                data_to_write = self.pending_data.slice(0, split_row)
                remaining_data = self.pending_data.slice(split_row)

                self._write_data_to_file(data_to_write)
                self.pending_data = remaining_data
                self._check_and_roll_if_needed()

    def _write_data_to_file(self, data: pa.RecordBatch):
        if data.num_rows == 0:
            return
        file_name = f"data-{uuid.uuid4()}.{self.file_format}"
        file_path = self._generate_file_path(file_name)
        try:
            if self.file_format == CoreOptions.FILE_FORMAT_PARQUET:
                self.file_io.write_parquet(file_path, data, compression=self.compression)
            elif self.file_format == CoreOptions.FILE_FORMAT_ORC:
                self.file_io.write_orc(file_path, data, compression=self.compression)
            elif self.file_format == CoreOptions.FILE_FORMAT_AVRO:
                self.file_io.write_avro(file_path, data, compression=self.compression)
            else:
                raise ValueError(f"Unsupported file format: {self.file_format}")

            key_columns_batch = data.select(self.trimmed_primary_key_fields)
            min_key_data = key_columns_batch.slice(0, 1).to_pylist()[0]
            max_key_data = key_columns_batch.slice(key_columns_batch.num_rows - 1, 1).to_pylist()[0]
            self.committed_files.append(DataFileMeta(
                file_name=file_name,
                file_size=self.file_io.get_file_size(file_path),
                row_count=data.num_rows,
                min_key=BinaryRow(min_key_data, self.trimmed_primary_key_fields),
                max_key=BinaryRow(max_key_data, self.trimmed_primary_key_fields),
                key_stats=None,  # TODO
                value_stats=None,
                min_sequence_number=0,
                max_sequence_number=0,
                schema_id=0,
                level=0,
                extra_files=None,
                file_path=str(file_path),
            ))

        except Exception as e:
            raise RuntimeError(f"Failed to write {self.file_format} file {file_path}: {e}") from e

    def _generate_file_path(self, file_name: str) -> Path:
        path_builder = self.table.table_path

        for i, field_name in enumerate(self.table.partition_keys):
            path_builder = path_builder / (field_name + "=" + self.partition[i])
        path_builder = path_builder / ("bucket-" + str(self.bucket)) / file_name

        return path_builder


def _find_optimal_split_point(data: pa.RecordBatch, target_size: int) -> int:
    total_rows = data.num_rows
    if total_rows <= 1:
        return 0

    left, right = 1, total_rows
    best_split = 0

    while left <= right:
        mid = (left + right) // 2
        slice_data = data.slice(0, mid)
        slice_size = slice_data.get_total_buffer_size()

        if slice_size <= target_size:
            best_split = mid
            left = mid + 1
        else:
            right = mid - 1

    return best_split