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

"""Optional Rust data writer behind PyPaimon's batch and stream builders."""

import pyarrow as pa

from pypaimon.common.options.core_options import ChangelogProducer
from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.execution.native.runtime import native_runtime_available
from pypaimon.schema.arrow_schema import arrow_schemas_compatible
from pypaimon.schema.data_types import PyarrowFieldParser, is_blob_file_field
from pypaimon.table.bucket_mode import BucketMode
from pypaimon.write.commit_identifier import validate_commit_identifier
from pypaimon.execution.native.write.native_commit import (
    create_native_write_table, from_native_commit_messages,
)
from pypaimon.write.row_utils import (
    require_columns, row_to_named_values, row_values_to_arrow_table, validate_write_columns,
)


def create_native_write(table, commit_user, static_partition=None, stream=False,
                        *, fixed_bucket=False, bucket_plan=None, restore_snapshot_id=None):
    """Return a native writer if the table can use the filesystem write path."""
    # DV tables flush normal level-0 files using the configured merge engine.
    # Batch visibility is controlled by deletion-vectors.merge-on-read.
    if (not native_runtime_available()
            or table.bucket_mode() not in (BucketMode.HASH_FIXED, BucketMode.HASH_DYNAMIC,
                                           BucketMode.BUCKET_UNAWARE,
                                           BucketMode.CROSS_PARTITION,
                                           BucketMode.POSTPONE_MODE)
            # The format matters only when this write emits INPUT changelogs.
            # Batch overwrites suppress them because they cannot be committed.
            or (table.options.changelog_file_format() not in (None, 'parquet')
                and table.options.changelog_producer() == ChangelogProducer.INPUT
                and (stream or static_partition is None))
            or table.options.file_format() != 'parquet'
            or table.options.video_frame_fields()):
        return None
    native_table = create_native_write_table(table)
    if native_table is None:
        return None
    if fixed_bucket:
        builder = native_table.new_postpone_fixed_bucket_write_builder().with_commit_user(commit_user)
        if bucket_plan is not None:
            builder = builder.with_bucket_plan(bucket_plan.to_arrow(table))
        if static_partition is not None:
            builder = builder.with_overwrite(static_partition)
        return NativePostponeFixedBucketTableWrite(
            table, commit_user, static_partition, bucket_plan, builder.new_write())
    if stream:
        builder = native_table.new_stream_write_builder().with_commit_user(commit_user)
    else:
        builder = native_table.new_batch_write_builder().with_commit_user(commit_user)
        if static_partition is not None:
            builder = builder.with_overwrite(static_partition)
    if restore_snapshot_id is not None:
        builder.with_restore_snapshot(restore_snapshot_id)
    return NativeTableWrite(table, commit_user, static_partition, stream,
                            builder.new_write(), restore_snapshot_id)


class NativeTableWrite:
    """Adapt Python inputs and commit messages to one Rust writer."""

    def __init__(self, table, commit_user, static_partition, stream, native_writer,
                 restore_snapshot_id=None):
        self.table = table
        self.commit_user = commit_user
        self.static_partition = static_partition
        self.stream = stream
        self._native_writer = native_writer
        self._written = False
        self._batch_committed = False
        self._table_schema = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
        self._schema = self._table_schema
        self._write_cols = None
        self._blob_consumer = None
        self._blob_uri_reader_factory = None
        self._restore_snapshot_id = restore_snapshot_id
        self._blob_rows = None
        self._defer_blob_input = False
        if any(is_blob_file_field(field) for field in table.table_schema.fields):
            from pypaimon.execution.native.write.native_blob_rows import NativeBlobRows
            self._blob_rows = NativeBlobRows()
            self._native_writer.with_blob_uri_reader_factory(self._blob_rows)

    def with_write_type(self, write_cols):
        names = validate_write_columns(write_cols)
        if self.table.is_primary_key_table:
            if len(names) == len(self._table_schema.names) and set(names) == set(self._table_schema.names):
                # Python treats a complete PK write type as the unchanged table
                # schema. It does not configure a PK column-subset writer.
                return self
            raise ValueError('Column-subset primary-key writes require the Python executor')
        self._configure_native('with_write_type', names)
        self._schema = pa.schema([self._table_schema.field(name) for name in names])
        self._write_cols = names
        return self

    def with_blob_consumer(self, blob_consumer):
        if blob_consumer is None:
            self._configure_native('with_blob_consumer', None)
            self._blob_consumer = None
            return self
        if not callable(blob_consumer):
            raise TypeError('blob_consumer must be callable')
        from pypaimon.table.row.blob import BlobDescriptor

        def consume(field_name, encoded):
            descriptor = None if encoded is None else BlobDescriptor.deserialize(encoded)
            return blob_consumer(field_name, descriptor)

        self._configure_native('with_blob_consumer', consume)
        self._blob_consumer = blob_consumer
        return self

    def with_blob_uri_reader_factory(self, uri_reader_factory):
        if uri_reader_factory is not None and not callable(getattr(uri_reader_factory, 'create', None)):
            raise TypeError('URI reader factory must provide callable create')
        if self._blob_rows is not None:
            # Only core knows whether RowKind filtering accepted any rows or a
            # postpone writer buffered data. Validate before changing the live
            # bridge's fallback, including when clearing the application factory.
            self._configure_native('with_blob_uri_reader_factory', self._blob_rows)
            self._blob_rows.fallback = uri_reader_factory
        else:
            self._configure_native('with_blob_uri_reader_factory', uri_reader_factory)
        self._blob_uri_reader_factory = uri_reader_factory
        return self

    def _configure_native(self, method, value):
        try:
            getattr(self._native_writer, method)(value)
        except ValueError as error:
            # Core owns written state: an ignored RowKind can leave it writable
            # even after this facade accepted a batch. Preserve the existing
            # Python state-error type without changing unrelated input errors.
            message = method + ' must be called before any write operation'
            if message in str(error):
                raise RuntimeError(message) from error
            raise

    def write_arrow(self, data, bucket=None):
        data = self._prepare_native_arrow_data(data)
        if isinstance(data, pa.RecordBatch):
            return self.write_arrow_batch(data, bucket)
        for batch in data.to_batches():
            self.write_arrow_batch(batch, bucket)

    def write_arrow_batch(self, data, bucket=None):
        data = self._prepare_native_arrow_data(data)
        self._write_native_batch(data, bucket)

    def _write_native_batch(self, data, bucket=None):
        if data.num_rows:
            # A failed native write may already have produced files. Never
            # retry that batch through Python after this point.
            self._written = True
        self._native_writer.write_arrow(data, bucket)

    def _prepare_native_arrow_data(self, data):
        # Core resolves row kinds and filters ignored events before inspecting
        # payload values. This facade checks only the logical Arrow schema.
        compatible = dict(check_top_level_nullability=False, allow_binary_compatibility=True,
                          check_nested_nullability=False)
        names = data.schema.names
        if len(set(names)) == len(names) and all(name in self._table_schema.names for name in names):
            expected = pa.schema([self._table_schema.field(name) for name in names])
            if arrow_schemas_compatible(data.schema, expected, **compatible):
                # Core validates the selected write type and routing columns.
                # Python checks only the declared types of transported fields.
                return data
        raise ValueError(
            "Input schema isn't consistent with table schema and write cols. "
            f"Input schema is: {data.schema} Table schema is: {self._schema} "
            f"Write cols is: {self._write_cols}")

    def write_pandas(self, dataframe):
        from pypaimon.write.input_adapter import _pandas_write_batch
        self.write_arrow_batch(_pandas_write_batch(self.table, dataframe, self._schema, self._table_schema))

    def roll_before_group_if_needed(self, row_count):
        if isinstance(row_count, bool) or not isinstance(row_count, int) or row_count <= 0:
            raise ValueError('Group row count must be a positive integer.')
        # Python's group boundary only affects packed video writers. Native
        # selection excludes video, so ordinary Parquet/BLOB writes do not roll.

    def write_row(self, row):
        values = row_to_named_values(row, self.table.table_schema.fields)
        names = self._schema.names
        full_row = set(values) == set(self._table_schema.names)
        require_columns(values, self.table.partition_keys, 'write_row')
        from pypaimon.write.row_kind import skip_write_row, with_row_kind
        logical_columns = self._write_cols if full_row else list(values)
        if skip_write_row(self.table, values, row.get_row_kind(), logical_columns):
            return
        require_columns(values, names, 'write_row')
        if self._write_cols is not None:
            routing = list(self.table.partition_keys)
            if ((self.table.bucket_mode() == BucketMode.HASH_FIXED and self.table.options.bucket() > 1)
                    or isinstance(self, NativePostponeFixedBucketTableWrite)):
                routing.extend(self.table.table_schema.bucket_keys)
            routing = list(dict.fromkeys(routing))
            row_kind_field = self.table.options.options.to_map().get('rowkind.field')
            if row_kind_field is not None:
                routing.append(row_kind_field)
            require_columns(values, routing, 'write_row')
            # Keep the selected order, followed by routing/event fields. Values
            # outside this transport remain untouched, including live Blobs.
            names = list(names) + [name for name in self._table_schema.names
                                   if name in routing and name not in names]
        try:
            data = (self._blob_rows.to_batch(self.table, [values], names)
                    if self._blob_rows is not None else
                    self._prepare_native_arrow_data(row_values_to_arrow_table(
                        values, self.table.table_schema.fields, names).to_batches()[0]))
            self._write_native_batch(with_row_kind(self.table, data, row))
        finally:
            # Core externalizes Blob values during write, before buffering PK
            # records. Their live Python objects need not survive the call.
            if self._blob_rows is not None:
                self._blob_rows.release_consumed()
                if not self._defer_blob_input:
                    self._blob_rows.clear()

    def prepare_commit(self, commit_identifier):
        commit_identifier = validate_commit_identifier(commit_identifier)
        if self.stream:
            messages = self._native_writer.prepare_commit(True, commit_identifier)
        else:
            if commit_identifier != BATCH_COMMIT_IDENTIFIER:
                raise ValueError('Batch writer execution requires the batch commit identifier')
            self._batch_committed = True
            try:
                messages = self._native_writer.prepare_commit()
            finally:
                if self._blob_rows is not None and self._defer_blob_input:
                    self._blob_rows.clear()
        return from_native_commit_messages(self.table, messages)

    def close(self):
        if self._native_writer is not None:
            self._native_writer.close()
            self._native_writer = None
        if self._blob_rows is not None:
            self._blob_rows.clear()

    def abort(self):
        # Rust cleans only files still owned by its writer. Prepared commit
        # messages must survive failures whose publication outcome is unknown.
        self.close()


class NativePostponeFixedBucketTableWrite(NativeTableWrite):
    """Fixed-bucket batches retain the standard native ownership contract.

    A distributed coordinator reads bucket metadata through Rust and groups
    inputs with shared arithmetic before workers use the fixed-bucket writer.
    """

    def __init__(self, table, commit_user, static_partition, bucket_plan, native_writer):
        super().__init__(table, commit_user, static_partition, False, native_writer)
        self._bucket_plan = bucket_plan
        # Automatic planning stages input until all partition sizes are known.
        self._defer_blob_input = (bucket_plan is None
                                  and table.options.postpone_default_bucket_num() is None)