# 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 logging
import uuid
import pyarrow as pa
from typing import Optional, Tuple, Dict

from pypaimon.execution.python.write.format_file_writer import (
    write_blob_file
)
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.data.timestamp import Timestamp
from pypaimon.table.row.blob import (
    BlobConsumer,
    video_payload_descriptor,
)
from pypaimon.execution.python.write.writer.append_only_data_writer import AppendOnlyDataWriter
from pypaimon.execution.python.write.writer.blob_file_writer import BlobFileWriter
from pypaimon.execution.python.write.writer.video_group import VideoGroupRollingPolicy

logger = logging.getLogger(__name__)


class BlobWriter(AppendOnlyDataWriter):

    def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, blob_column: str,
                 options: Dict[str, str] = None, blob_consumer: Optional[BlobConsumer] = None,
                 video: bool = False, uri_reader_factory=None):
        super().__init__(table, partition, bucket, max_seq_number,
                         options, write_cols=[blob_column])

        self.video = video
        if video and blob_consumer is not None:
            raise ValueError(
                f"BlobConsumer is not supported for video frame field '{blob_column}'."
            )
        self.file_format = (
            CoreOptions.FILE_FORMAT_VIDEO
            if video
            else CoreOptions.FILE_FORMAT_BLOB
        )

        # Store blob column name for use in metadata creation
        self.blob_column = blob_column

        options = self.table.options
        self.blob_target_file_size = CoreOptions.blob_target_file_size(options)
        # Use the effective options (constructor-provided), consistent with the other accessors.
        self.blob_copy_buffer_size = self.options.blob_copy_buffer_size()

        self._blob_consumer = blob_consumer
        self._uri_reader_factory = uri_reader_factory
        self.current_writer: Optional[BlobFileWriter] = None
        self.current_file_path: Optional[str] = None
        self.record_count = 0

        self.file_uuid = str(uuid.uuid4())
        self.file_count = 0
        self._video_group_policy = VideoGroupRollingPolicy() if video else None

        logger.info(f"Initialized BlobWriter with blob file format, blob_target_file_size={self.blob_target_file_size}")

    def _check_and_roll_if_needed(self):
        # Rolling here is driven by the size of the external blobs, which the
        # buffered descriptors say nothing about, so there is no cheap count to
        # check first: every write drains the buffer.
        pending = self._buffer.take()
        if pending is None:
            return

        # Always write blob rows one-by-one so rolling uses actual blob bytes size rather than
        # in-memory serialized descriptor size.
        for i in range(pending.num_rows):
            row_data = pending.slice(i, 1)
            next_group = video_payload_descriptor(row_data.column(0)[0])
            self._roll_before_video_group(next_group)
            self._write_row_to_file(row_data)
            self.record_count += 1
            if self._video_group_policy is not None:
                self._video_group_policy.record(next_group)

            if self.rolling_file():
                if (
                    self._video_group_policy is None
                    or not self._video_group_policy.defer_roll()
                ):
                    self.close_current_writer()

    def _write_row_to_file(self, row_data: pa.Table):
        """Write a single row to the current blob file. Opens a new file if needed."""
        if row_data.num_rows == 0:
            return

        if self.current_writer is None:
            self.open_current_writer()

        self.current_writer.write_row(row_data)
        # This ensures each row has a unique sequence number for data versioning and consistency
        self.sequence_generator.next()

    def write_blob(self, value, arrow_type=pa.large_binary()):
        next_group = video_payload_descriptor(value)
        self._roll_before_video_group(next_group)
        if self.current_writer is None:
            self.open_current_writer()

        self.current_writer.write_blob(self.blob_column, arrow_type, value)
        self.sequence_generator.next()
        self.record_count += 1
        if self._video_group_policy is not None:
            self._video_group_policy.record(next_group)

        if self.rolling_file():
            if (
                self._video_group_policy is None
                or not self._video_group_policy.defer_roll()
            ):
                self.close_current_writer()

    def open_current_writer(self):
        file_name = (f"{CoreOptions.data_file_prefix(self.options)}"
                     f"{self.file_uuid}-{self.file_count}.{self.file_format}")
        self.file_count += 1  # Increment counter for next file
        file_path = self._generate_file_path(file_name)
        self.current_file_path = file_path
        self.current_writer = BlobFileWriter(
            self.file_io,
            file_path,
            blob_consumer=self._blob_consumer,
            copy_buffer_size=self.blob_copy_buffer_size,
            video=self.video,
            uri_reader_factory=self._uri_reader_factory,
        )

    def rolling_file(self) -> bool:
        if self.current_writer is None:
            return False

        return (
            self.current_writer.row_count >= self.target_file_row_num
            or self.current_writer.reach_target_size(self.blob_target_file_size)
        )

    def should_roll_before_group(self, row_count: int) -> bool:
        return (
            self._video_group_policy is not None
            and self.current_writer is not None
            and (
                self._video_group_policy.pending_roll
                or self.current_writer.row_count + row_count
                > self.target_file_row_num
            )
        )

    def close_current_writer(self):
        """Close current writer and create metadata."""
        if self.current_writer is None:
            return

        file_size = self.current_writer.close()
        file_name = self.current_file_path.split('/')[-1]
        row_count = self.current_writer.row_count

        # Determine if this is an external path
        is_external_path = self.external_path_provider is not None
        external_path_str = self.current_file_path if is_external_path else None

        self._add_file_metadata(file_name, self.current_file_path, row_count, file_size, external_path_str)

        self.current_writer = None
        self.current_file_path = None
        if self._video_group_policy is not None:
            self._video_group_policy.reset()

    def _roll_before_video_group(self, next_group):
        if (
            self._video_group_policy is not None
            and self.current_writer is not None
            and self._video_group_policy.should_roll_before(next_group)
        ):
            self.close_current_writer()

    def _write_data_to_file(self, data):
        """
        Keep a fallback path for direct blob table writes while preserving the writer uuid+counter
        naming behavior.
        """
        if data.num_rows == 0:
            return

        # This ensures each row gets a unique sequence number, matching the behavior expected
        for _ in range(data.num_rows):
            self.sequence_generator.next()

        file_name = f"data-{self.file_uuid}-{self.file_count}.{self.file_format}"
        self.file_count += 1
        file_path = self._generate_file_path(file_name)

        write_blob_file(self.file_io, file_path, data)

        file_size = self.file_io.get_file_size(file_path)

        is_external_path = self.external_path_provider is not None
        external_path_str = file_path if is_external_path else None

        # Reuse _add_file_metadata for consistency (blob table is append-only, no primary keys)
        self._add_file_metadata(file_name, file_path, data, file_size, external_path_str)

    def _add_file_metadata(self, file_name: str, file_path: str, data_or_row_count, file_size: int,
                           external_path: Optional[str] = None):
        """Add file metadata to committed_files."""
        from pypaimon.manifest.schema.data_file_meta import DataFileMeta
        from pypaimon.manifest.schema.simple_stats import SimpleStats
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.schema.data_types import PyarrowFieldParser

        # Handle both Table and row_count
        if isinstance(data_or_row_count, pa.Table):
            data = data_or_row_count
            row_count = data.num_rows
            data_fields = PyarrowFieldParser.to_paimon_schema(data.schema)
            # Compute statistics across all batches, not just the first one
            # This ensures correct min/max/null_counts when data has multiple batches
            column_stats = {
                field.name: self._get_column_stats(data, field.name)
                for field in data_fields
            }
            min_value_stats = [column_stats[field.name]['min_values'] for field in data_fields]
            max_value_stats = [column_stats[field.name]['max_values'] for field in data_fields]
            value_null_counts = [column_stats[field.name]['null_counts'] for field in data_fields]
        else:
            # row_count only (from BlobFileWriter)
            row_count = data_or_row_count
            data_fields = [self.table.field_dict[self.blob_column]]
            min_value_stats = [None]
            max_value_stats = [None]
            value_null_counts = [0]

        min_seq = self.sequence_generator.current - row_count
        max_seq = self.sequence_generator.current - 1
        self.sequence_generator.start = self.sequence_generator.current

        self.committed_files.append(DataFileMeta.create(
            file_name=file_name,
            file_size=file_size,
            row_count=row_count,
            min_key=GenericRow([], []),
            max_key=GenericRow([], []),
            key_stats=SimpleStats(GenericRow([], []), GenericRow([], []), []),
            value_stats=SimpleStats(
                GenericRow(min_value_stats, data_fields),
                GenericRow(max_value_stats, data_fields),
                value_null_counts),
            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().to_millis_timestamp(),
            delete_row_count=0,
            file_source=0,  # FileSource.APPEND = 0
            value_stats_cols=None,
            external_path=external_path,
            first_row_id=None,
            write_cols=self.write_cols,
            file_path=file_path,
        ))

    def prepare_commit(self):
        """Prepare commit, preserving blob files in write/rolling order."""
        # Close current file if open.
        if self.current_writer is not None:
            self.close_current_writer()

        # Call parent to flush anything left in the buffer.
        return super().prepare_commit()

    def close(self):
        """Close current writer if open."""
        # Close current file if open.
        if self.current_writer is not None:
            self.close_current_writer()

        # Call parent to flush anything left in the buffer.
        super().close()

    def delete_file_upon_abort(self) -> bool:
        return self._blob_consumer is None

    def abort(self):
        if self.current_writer is not None:
            try:
                self.current_writer.abort()
            except Exception as e:
                logger.warning(f"Error aborting blob writer: {e}", exc_info=e)
            self.current_writer = None
            self.current_file_path = None
        if self._video_group_policy is not None:
            self._video_group_policy.reset()
        if not self.delete_file_upon_abort():
            self._buffer.reset()
            self.committed_files.clear()
        else:
            super().abort()

    @staticmethod
    def _get_column_stats(data_or_batch, column_name: str):
        """
        Compute column statistics for a column in a Table or RecordBatch.
        """
        # Handle both Table and RecordBatch
        if isinstance(data_or_batch, pa.Table):
            column_array = data_or_batch.column(column_name)
        else:
            column_array = data_or_batch.column(column_name)

        # For blob data, don't generate min/max values
        return {
            "min_values": None,
            "max_values": None,
            "null_counts": column_array.null_count,
        }