################################################################################
#  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 pyarrow as pa
import pyarrow.compute as pc
import uuid
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple

from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.external_path_provider import ExternalPathProvider
from pypaimon.data.timestamp import Timestamp
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.manifest.schema.simple_stats import SimpleStats
from pypaimon.schema.data_types import PyarrowFieldParser
from pypaimon.table.bucket_mode import BucketMode
from pypaimon.table.row.generic_row import GenericRow


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

    def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, options: CoreOptions = None,
                 write_cols: Optional[List[str]] = None):
        from pypaimon.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_keys_fields = self.table.trimmed_primary_keys_fields
        self.trimmed_primary_keys = self.table.trimmed_primary_keys

        self.options = options
        self.target_file_size = self.options.target_file_size(self.table.is_primary_key_table)
        # POSTPONE_BUCKET uses AVRO format, otherwise default to PARQUET
        default_format = (
            CoreOptions.FILE_FORMAT_AVRO
            if self.bucket == BucketMode.POSTPONE_BUCKET.value
            else CoreOptions.FILE_FORMAT_PARQUET
        )
        self.file_format = self.options.file_format(default_format)
        self.compression = self.options.file_compression()
        self.zstd_level = self.options.file_compression_zstd_level()
        self.sequence_generator = SequenceGenerator(max_seq_number)

        self.pending_data: Optional[pa.Table] = None
        self.committed_files: List[DataFileMeta] = []
        self.write_cols = write_cols
        self.blob_as_descriptor = self.options.blob_as_descriptor()

        self.path_factory = self.table.path_factory()
        self.external_path_provider: Optional[ExternalPathProvider] = self.path_factory.create_external_path_provider(
            self.partition, self.bucket
        )
        # Store the current generated external path to preserve scheme in metadata
        self._current_external_path: Optional[str] = None

    def write(self, data: pa.RecordBatch):
        try:
            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()
        except Exception as e:
            import logging
            logger = logging.getLogger(__name__)
            logger.warning("Exception occurs when writing data. Cleaning up.", exc_info=e)
            self.abort()
            raise e

    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):
        try:
            if self.pending_data is not None and self.pending_data.num_rows > 0:
                self._write_data_to_file(self.pending_data)
        except Exception as e:
            import logging
            logger = logging.getLogger(__name__)
            logger.warning("Exception occurs when closing writer. Cleaning up.", exc_info=e)
            self.abort()
            raise e
        finally:
            self.pending_data = None
            # Note: Don't clear committed_files in close() - they should be returned by prepare_commit()

    def abort(self):
        """
        Abort all writers and clean up resources. This method should be called when an error occurs
        during writing. It deletes any files that were written and cleans up resources.
        """
        # Delete any files that were written
        for file_meta in self.committed_files:
            try:
                # Use external_path if available (contains full URL scheme), otherwise use file_path
                path_to_delete = file_meta.external_path if file_meta.external_path else file_meta.file_path
                if path_to_delete:
                    path_str = str(path_to_delete)
                    self.file_io.delete_quietly(path_str)
            except Exception as e:
                # Log but don't raise - we want to clean up as much as possible
                import logging
                logger = logging.getLogger(__name__)
                path_to_delete = file_meta.external_path if file_meta.external_path else file_meta.file_path
                logger.warning(f"Failed to delete file {path_to_delete} during abort: {e}")

        # Clean up resources
        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):
        while self.pending_data is not None:
            current_size = self.pending_data.nbytes
            if current_size <= self.target_file_size:
                break
            split_row = self._find_optimal_split_point(self.pending_data, self.target_file_size)
            if split_row <= 0:
                break
            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

    def _write_data_to_file(self, data: pa.Table):
        if data.num_rows == 0:
            return
        file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}"
        file_path = self._generate_file_path(file_name)

        is_external_path = self.external_path_provider is not None
        if is_external_path:
            # Use the stored external path from _generate_file_path to preserve scheme
            external_path_str = self._current_external_path if self._current_external_path else None
        else:
            external_path_str = None

        if self.file_format == CoreOptions.FILE_FORMAT_PARQUET:
            self.file_io.write_parquet(file_path, data, compression=self.compression, zstd_level=self.zstd_level)
        elif self.file_format == CoreOptions.FILE_FORMAT_ORC:
            self.file_io.write_orc(file_path, data, compression=self.compression, zstd_level=self.zstd_level)
        elif self.file_format == CoreOptions.FILE_FORMAT_AVRO:
            self.file_io.write_avro(file_path, data, compression=self.compression, zstd_level=self.zstd_level)
        elif self.file_format == CoreOptions.FILE_FORMAT_BLOB:
            self.file_io.write_blob(file_path, data)
        elif self.file_format == CoreOptions.FILE_FORMAT_LANCE:
            self.file_io.write_lance(file_path, data)
        else:
            raise ValueError(f"Unsupported file format: {self.file_format}")

        # min key & max key

        selected_table = data.select(self.trimmed_primary_keys)
        key_columns_batch = selected_table.to_batches()[0]
        min_key_row_batch = key_columns_batch.slice(0, 1)
        max_key_row_batch = key_columns_batch.slice(key_columns_batch.num_rows - 1, 1)
        min_key = [col.to_pylist()[0] for col in min_key_row_batch.columns]
        max_key = [col.to_pylist()[0] for col in max_key_row_batch.columns]

        # key stats & value stats
        value_stats_enabled = self.options.metadata_stats_enabled()
        if value_stats_enabled:
            stats_fields = self.table.fields if self.table.is_primary_key_table \
                else PyarrowFieldParser.to_paimon_schema(data.schema)
        else:
            stats_fields = self.table.trimmed_primary_keys_fields
        column_stats = {
            field.name: self._get_column_stats(data, field.name)
            for field in stats_fields
        }
        key_fields = self.trimmed_primary_keys_fields
        key_stats = self._collect_value_stats(data, key_fields, column_stats)
        if not all(count == 0 for count in key_stats.null_counts):
            raise RuntimeError("Primary key should not be null")

        value_fields = stats_fields if value_stats_enabled else []
        value_stats = self._collect_value_stats(data, value_fields, column_stats)

        min_seq = self.sequence_generator.start
        max_seq = self.sequence_generator.current
        self.sequence_generator.start = self.sequence_generator.current
        self.committed_files.append(DataFileMeta.create(
            file_name=file_name,
            file_size=self.file_io.get_file_size(file_path),
            row_count=data.num_rows,
            min_key=GenericRow(min_key, self.trimmed_primary_keys_fields),
            max_key=GenericRow(max_key, self.trimmed_primary_keys_fields),
            key_stats=key_stats,
            value_stats=value_stats,
            min_sequence_number=min_seq,
            max_sequence_number=max_seq,
            schema_id=self.table.table_schema.id,
            level=0,
            extra_files=[],
            creation_time=Timestamp.now(),
            delete_row_count=0,
            file_source=0,
            value_stats_cols=None if value_stats_enabled else [],
            external_path=external_path_str,  # Set external path if using external paths
            first_row_id=None,
            write_cols=self.write_cols,
            # None means all columns in the table have been written
            file_path=file_path,
        ))

    def _generate_file_path(self, file_name: str) -> str:
        if self.external_path_provider:
            external_path = self.external_path_provider.get_next_external_data_path(file_name)
            self._current_external_path = external_path
            return external_path

        bucket_path = self.path_factory.bucket_path(self.partition, self.bucket)
        return f"{bucket_path.rstrip('/')}/{file_name}"

    @staticmethod
    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.nbytes

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

        return best_split

    def _collect_value_stats(self, data: pa.Table, fields: List,
                             column_stats: Optional[Dict[str, Dict]] = None) -> SimpleStats:
        if not fields:
            return SimpleStats.empty_stats()
        
        if column_stats is None or not column_stats:
            column_stats = {
                field.name: self._get_column_stats(data, field.name)
                for field in fields
            }
        
        min_stats = [column_stats[field.name]['min_values'] for field in fields]
        max_stats = [column_stats[field.name]['max_values'] for field in fields]
        null_counts = [column_stats[field.name]['null_counts'] for field in fields]
        
        return SimpleStats(
            GenericRow(min_stats, fields),
            GenericRow(max_stats, fields),
            null_counts
        )

    @staticmethod
    def _get_column_stats(record_batch: pa.RecordBatch, column_name: str) -> Dict:
        column_array = record_batch.column(column_name)
        if column_array.null_count == len(column_array):
            return {
                "min_values": None,
                "max_values": None,
                "null_counts": column_array.null_count,
            }
        
        column_type = column_array.type
        supports_minmax = not (pa.types.is_nested(column_type) or pa.types.is_map(column_type))
        
        if not supports_minmax:
            return {
                "min_values": None,
                "max_values": None,
                "null_counts": column_array.null_count,
            }
        
        min_values = pc.min(column_array).as_py()
        max_values = pc.max(column_array).as_py()
        null_counts = column_array.null_count
        return {
            "min_values": min_values,
            "max_values": max_values,
            "null_counts": null_counts,
        }


class SequenceGenerator:
    def __init__(self, start: int = 0):
        self.start = start
        self.current = start

    def next(self) -> int:
        self.current += 1
        return self.current
