# 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 native row-ID updates for batch data-evolution tables."""

import json
from datetime import date, datetime, time

import pyarrow as pa
import pyarrow.compute as pc

from pypaimon.execution.native.table import _native_query_auth_supported
from pypaimon.execution.native.runtime import native_runtime_available
from pypaimon.schema.arrow_schema import arrow_schemas_compatible
from pypaimon.schema.data_types import MapType, PyarrowFieldParser, RowType, VectorType, is_blob_file_field
from pypaimon.common.options.core_options import ChangelogProducer
from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.table.file_store_table import FileStoreTable
from pypaimon.execution.native.write.native_commit import (
    create_native_write_table, from_native_commit_messages,
)
from pypaimon.write.row_input import RowIdUpdateInputMixin, normalize_row_ids
from pypaimon.write.row_utils import _contains_blob_value, value_for_arrow


def _native_update_formats_supported(table):
    return (table.options.file_format() == 'parquet'
            and (not table.options.with_vector_format() or table.options.vector_file_format() == 'parquet'
                 or not any(isinstance(field.type, VectorType) for field in table.table_schema.fields)))


def _native_update_table(table, *, allow_changelog=False):
    """Resolve common native update prerequisites before preparing changes."""
    if (type(table) is not FileStoreTable
            or not table.options.native_write_enabled()
            or not _native_update_formats_supported(table)
            or (not allow_changelog and table.options.changelog_producer() != ChangelogProducer.NONE)
            or not native_runtime_available()
            or not _native_query_auth_supported(table)):
        return None
    native_table = create_native_write_table(table)
    if native_table is not None:
        from pypaimon.execution.native.table import _with_native_read_snapshot
        native_table = _with_native_read_snapshot(table, native_table)
    return native_table


def _native_row_id_table(table):
    """Row-level operations additionally require a data-evolution table."""
    if (type(table) is not FileStoreTable
            or not table.options.data_evolution_enabled()
            or not table.options.row_tracking_enabled()
            or table.is_primary_key_table):
        return None
    return _native_update_table(table)


def create_native_update(table, commit_user, columns):
    """Use the public core updater for direct and grouped row-ID updates."""
    if not _native_update_columns_supported(table, columns):
        return None
    native_table = _native_row_id_table(table)
    if native_table is None:
        return None
    writer = (native_table.new_batch_write_builder()
              .with_commit_user(commit_user)
              .new_update())
    if columns is not None:
        writer.with_update_type(columns)
    return NativeBatchTableUpdate(table, writer)


def create_native_update_by_row_id(table, commit_user, commit_identifier, plan=None):
    """Create a core updater sharing one snapshot across incremental calls."""
    # Columns are selected per call; one core updater pins the snapshot and
    # tracks overlaps across Parquet and dedicated Blob updates.
    if not _native_update_columns_supported(table, None):
        return None
    native_table = _native_row_id_table(
        table._copy_with_snapshot(plan.table._read_snapshot) if plan is not None else table)
    if native_table is None:
        return None
    if commit_identifier == BATCH_COMMIT_IDENTIFIER:
        update = native_table.new_batch_write_builder().with_commit_user(commit_user).new_update()
    else:
        update = native_table.new_stream_write_builder().with_commit_user(commit_user).new_update()
    if plan is not None:
        from pypaimon.execution.native.write.row_id_plan import native_plan
        update.with_plan(native_plan(plan.table, plan.plan))
    writer = (update.new_update_by_row_id() if commit_identifier == BATCH_COMMIT_IDENTIFIER
              else update.new_update_by_row_id(commit_identifier))
    return NativeTableUpdateByRowId(
        table, commit_user, commit_identifier, writer)


def create_native_shard_updator(table, commit_user, projection, columns, shard, shard_count, stream=False):
    """Core owns whole-file sharding, the captured reader and positional updates."""
    if not _native_update_columns_supported(table, columns):
        return None
    native_table = _native_row_id_table(table)
    if native_table is None:
        return None
    builder = (native_table.new_stream_write_builder().with_commit_user(commit_user) if stream else
               native_table.new_batch_write_builder().with_commit_user(commit_user))
    update = builder.new_update()
    read_builder = table.new_read_builder()
    if projection is not None:
        read_builder.with_projection(projection)
    if columns is not None:
        update.with_update_type(columns)
    read = read_builder.new_read()
    writer = update.new_shard_updator(
        shard, shard_count, read_type=json.dumps(RowType(True, read_builder.read_type()).to_dict()))
    return NativeShardTableUpdator(table, read, writer)


class NativeShardTableUpdator:
    """Transport Arrow batches and commit messages for a core shard updater."""

    def __init__(self, table, read, writer):
        self._table = table
        self._writer = writer
        self._read = read

    def arrow_reader(self):
        from pypaimon.execution.native.read.table_read import NativeTableRead
        from pypaimon.read.read_support import _ClosableArrowBatchReader
        schema = self._read._output_arrow_schema()
        batches = NativeTableRead(self._read).convert_batches(self._writer.arrow_reader(), schema)
        return _ClosableArrowBatchReader(pa.ipc.RecordBatchReader.from_batches(schema, batches), batches)

    def update_by_arrow_batch(self, data):
        self._writer.update_by_arrow_batch(data)

    def prepare_commit(self):
        return from_native_commit_messages(self._table, self._writer.prepare_commit())

    def close(self):
        self._writer.close()


def _supported_upsert_key_type(data_type):
    return any(check(data_type) for check in (
        pa.types.is_boolean, pa.types.is_integer, pa.types.is_string,
        pa.types.is_large_string, pa.types.is_binary, pa.types.is_large_binary,
        pa.types.is_fixed_size_binary, pa.types.is_date, pa.types.is_decimal,
        pa.types.is_time, pa.types.is_timestamp,
        pa.types.is_floating,
    ))


def _dictionary_value_type(data_type):
    while pa.types.is_dictionary(data_type):
        data_type = data_type.value_type
    return data_type


def _native_upsert_key_type_compatible(source, target):
    """Qualify layouts; exact values and conversion bounds belong to core."""
    dictionary = pa.types.is_dictionary(source)
    source = _dictionary_value_type(source)
    if dictionary:
        # Python's take kernel can select dictionary indices even when it
        # cannot take the underlying view/run-end values directly. These
        # layouts are therefore valid existing matched-upsert inputs.
        while getattr(pa.types, 'is_run_end_encoded', lambda _: False)(source):
            source = _dictionary_value_type(source.value_type)
        if getattr(pa.types, 'is_string_view', lambda _: False)(source):
            source = pa.string()
        elif getattr(pa.types, 'is_binary_view', lambda _: False)(source):
            source = pa.binary()
    if pa.types.is_union(source):
        # A Union can contain unused children of unrelated types. Qualify a
        # possible match here; core validates only the active values exactly.
        return any(_native_upsert_key_type_compatible(field.type, target) for field in source)
    if pa.types.is_null(source):
        return True

    def numeric(data_type):
        return (pa.types.is_integer(data_type) or pa.types.is_floating(data_type)
                or pa.types.is_decimal(data_type))
    if numeric(source) and numeric(target):
        return True
    for family in (pa.types.is_date, pa.types.is_time):
        if family(source) and family(target):
            return True

    def binary(data_type):
        return (pa.types.is_binary(data_type) or pa.types.is_large_binary(data_type)
                or pa.types.is_fixed_size_binary(data_type))
    if binary(source) and binary(target):
        return True
    if pa.types.is_timestamp(source) and pa.types.is_timestamp(target):
        return (source.tz is None) == (target.tz is None)
    return arrow_schemas_compatible(
        pa.schema([pa.field('key', source)]), pa.schema([pa.field('key', target)]),
        check_top_level_nullability=False, allow_binary_compatibility=True)


def _native_upsert_key_compatible(source, target):
    if _native_upsert_key_type_compatible(source.type, target):
        return True

    def nulls(array):
        if pa.types.is_dictionary(array.type):
            # Index validity alone misses a selected NULL dictionary value.
            # Select only its validity bitmap, without decoding key payloads.
            return pc.fill_null(pc.take(nulls(array.dictionary), array.indices), True)
        if pa.types.is_union(array.type):
            # Unions have no validity bitmap. Only the active child's NULL
            # counts; inactive children can contain unrelated non-NULL values.
            # Read the parent buffers at its offset because PyArrow's union
            # type_codes/offsets accessors can omit the offset of a slice.
            buffers = array.buffers()
            codes = pa.Array.from_buffers(pa.int8(), len(array), [None, buffers[1]], offset=array.offset)
            offsets = (pa.Array.from_buffers(pa.int32(), len(array), [None, buffers[2]], offset=array.offset)
                       if array.type.mode == 'dense' else None)
            result = pc.fill_null(pa.nulls(len(array), type=pa.bool_()), False)
            for index, code in enumerate(array.type.type_codes):
                active = pc.equal(codes, code)
                child_nulls = nulls(array.field(index))
                if offsets is not None:
                    child_nulls = pc.take(child_nulls, pc.if_else(active, offsets, None))
                result = pc.if_else(active, child_nulls, result)
            return result
        return array.is_null()

    try:
        chunks = source.chunks if isinstance(source, pa.ChunkedArray) else [source]
        for chunk in chunks:
            if chunk.null_count != len(chunk) and not pc.all(nulls(chunk)).as_py():
                return False
            # Keep Python's existing input contract. In particular, bare view
            # and run-end inputs have no take kernel on supported PyArrow.
            pc.take(chunk, pa.array([], type=pa.int64()))
        return True
    except (pa.ArrowInvalid, pa.ArrowNotImplementedError, TypeError):
        return False


def _native_update_columns_supported(table, columns):
    """Core writes Parquet columns and Java-compatible sparse Blob deltas."""
    if type(table) is not FileStoreTable:
        return False
    names = set(table.field_names if columns is None else columns)
    return not names.intersection(table.options.video_frame_fields())


def _temporal_row_key_supported(value, data_type):
    """Qualify row constructors without changing temporal matching semantics."""
    if value is None:
        return True
    if pa.types.is_date(data_type):
        # A datetime is distinct from a date even at midnight. Arrow's DATE
        # constructor would discard this distinction before key matching.
        return isinstance(value, date) and not isinstance(value, datetime)
    if pa.types.is_timestamp(data_type):
        if not isinstance(value, datetime):
            return False
        offset = value.utcoffset()
        if (offset is not None) != (data_type.tz is not None):
            return False
        nanos = value.microsecond * 1000 + getattr(value, 'nanosecond', 0)
        if offset is not None:
            nanos -= offset.microseconds * 1000
    elif pa.types.is_time(data_type):
        if not isinstance(value, time):
            return False
        if value.utcoffset() is not None:
            return False
        nanos = value.microsecond * 1000
    else:
        return True
    # Compare precision, not datetime objects: equivalent instants expressed
    # with different timezones or DST folds must keep the existing core path.
    resolution = {'s': 1000000000, 'ms': 1000000, 'us': 1000, 'ns': 1}[data_type.unit]
    return nanos % resolution == 0


def _upsert_row_batches(rows, fields, schema, table=None, blob_rows=None, key_names=()):
    """Encode consecutive row shapes without padding absent fields with NULL.

    Keep source order, including duplicates across shapes. Key matching,
    last-write-wins and matched/append column validation belong to Rust core.
    """
    encoded_keys = {}
    for name in key_names:
        key_type = schema.field(name).type
        if any(not _temporal_row_key_supported(row[name], key_type) for row in rows):
            return None
        if not any(check(key_type) for check in (
                pa.types.is_integer, pa.types.is_floating, pa.types.is_decimal,
                pa.types.is_boolean, pa.types.is_string, pa.types.is_large_string,
                pa.types.is_binary, pa.types.is_large_binary, pa.types.is_fixed_size_binary)):
            continue
        encoded = pa.array([row[name] for row in rows], type=key_type)
        # Row constructors can change scalar identity before core sees a key.
        # Qualify all keys before encoding any payload: a later unsafe row
        # must not fall back after an earlier Blob input has been consumed.
        if any(row[name] != encoded[index].as_py() for index, row in enumerate(rows)):
            return None
        encoded_keys[name] = encoded
    batches = []
    start = 0
    while start < len(rows):
        names = set(rows[start])
        if not names <= set(schema.names):
            raise ValueError('upsert row fields must be in the table schema')
        end = start + 1
        while end < len(rows) and set(rows[end]) == names:
            end += 1
        selected = [field for field in fields if field.name in names]
        batch_schema = pa.schema([schema.field(field.name) for field in selected])
        if blob_rows is not None:
            batch = blob_rows.to_batch(table, rows[start:end], batch_schema.names)
        else:
            batch = pa.RecordBatch.from_pydict({
                field.name: (encoded_keys[field.name].slice(start, end - start) if field.name in encoded_keys else
                             [value_for_arrow(row[field.name], field) for row in rows[start:end]])
                for field in selected
            }, schema=batch_schema)
        batches.append(batch)
        start = end
    return batches


def create_native_upsert(table, commit_user, data, keys, columns):
    """Prepare one core upsert from Arrow columns or named row values."""
    if table.options.video_frame_fields() or not _native_update_columns_supported(table, columns):
        # An upsert can append unmatched rows; packed video writing is still
        # provided by the Python writer.
        return None
    native_table = _native_row_id_table(table)
    if native_table is None:
        return None
    fields = table.table_schema.fields
    schema = PyarrowFieldParser.from_paimon_schema(fields)
    blob_rows = None
    if any(not _supported_upsert_key_type(schema.field(key).type) for key in set(keys + table.partition_keys)):
        return None
    if not isinstance(data, pa.Table):
        if not columns or not data:
            return None
        if any(_contains_blob_value(value) for row in data for value in row.values()) or any(
                is_blob_file_field(field) and isinstance(field.type, MapType) and field.name in row
                for field in fields for row in data):
            from pypaimon.execution.native.write.native_blob_rows import NativeBlobRows
            blob_rows = NativeBlobRows()
        data = _upsert_row_batches(
            data, fields, schema, table, blob_rows, list(dict.fromkeys(keys + table.partition_keys)))
        if data is None:
            return None
    else:
        if (len(data.column_names) != len(set(data.column_names))
                or not set(data.column_names) <= set(schema.names)):
            return None
        # Keys need a shared representation before matching. Core selects
        # surviving payloads before applying row-ID update or append rules;
        # validating all payload types here would reject valid matched updates.
        key_names = list(dict.fromkeys(keys + table.partition_keys))
        for name in key_names:
            field = data.schema.field(name)
            if pa.types.is_null(field.type):
                # An untyped NULL has the same match identity as a typed NULL.
                # Give core the declared key layout without casting payloads.
                data = data.set_column(
                    data.schema.get_field_index(name), field.with_type(schema.field(name).type),
                    pa.nulls(data.num_rows, type=schema.field(name).type))
        if not all(_native_upsert_key_compatible(data.column(name), schema.field(name).type)
                   for name in key_names):
            return None
    writer = (native_table.new_batch_write_builder()
              .with_commit_user(commit_user)
              .new_update()
              .with_update_type(columns))
    if blob_rows is not None:
        writer.with_blob_uri_reader_factory(blob_rows)
    return NativeTableUpsert(table, writer, keys, data, blob_rows)


def create_native_predicate_update(table, commit_user, predicate, columns=None):
    """Prepare a public core operation before any assignment can run."""
    if not _native_update_columns_supported(table, columns):
        return None
    native_table = _native_row_id_table(table)
    if native_table is None:
        return None
    from pypaimon.execution.native.read.builder import _predicate_to_native
    native_predicate = None if predicate is None else _predicate_to_native(predicate)
    if native_predicate is not None:
        # Check predicate translation while fallback is still safe. Planning,
        # reading and callback execution belong to the core operation below.
        native_table.new_read_builder().with_filter(native_predicate)
    writer = (native_table.new_batch_write_builder()
              .with_commit_user(commit_user).new_update())
    return NativePredicateTableUpdate(table, writer, native_predicate)


def create_native_predicate_delete(table, commit_user, predicate):
    """Delegate matching and partition metadata deletion to the core updater."""
    from pypaimon.execution.native.read.builder import _predicate_to_native
    if type(table) is not FileStoreTable:
        return None

    def canonical_paths(node):
        if node.field_path and node.field_path != [node.field]:
            return False
        return (all(canonical_paths(child) for child in (node.literals or []))
                if node.method in ('and', 'or') else True)

    # Core selects whole-partition metadata deletes for append/PK tables too;
    # only the row-level branch requires row tracking and deletion vectors.
    from pypaimon.write.update_validation import predicate_fields
    fields = set() if predicate is None else predicate_fields(predicate)
    partition_delete = (table.is_primary_key_table and bool(fields)
                        and fields.issubset(table.partition_keys) and canonical_paths(predicate))
    # Partition deletion removes complete file/index metadata. It does not
    # produce row-level changes, even when the PK writer emits changelogs.
    # Core still validates the translated predicate and selects this fast path.
    native_table = _native_update_table(table, allow_changelog=partition_delete)
    if native_table is None:
        return None
    native_predicate = None if predicate is None else _predicate_to_native(predicate)
    if native_predicate is not None:
        # Resolve all fields/operators before selecting the mutating operation.
        native_table.new_read_builder().with_filter(native_predicate)
    writer = (native_table.new_batch_write_builder()
              .with_commit_user(commit_user).new_update())
    return NativePredicateTableUpdate(table, writer, native_predicate)


def create_native_delete(table, commit_user, plan=None):
    """Select Rust's deletion-vector writer for supported batch deletes."""
    if not table.options.deletion_vectors_enabled(False):
        return None
    native_table = _native_row_id_table(
        table._copy_with_snapshot(plan.table._read_snapshot) if plan is not None else table)
    if native_table is None:
        return None
    writer = (native_table.new_batch_write_builder()
              .with_commit_user(commit_user)
              .new_update())
    if plan is not None:
        from pypaimon.execution.native.write.row_id_plan import native_plan
        writer.with_plan(native_plan(plan.table, plan.plan))
    return NativeBatchTableUpdate(table, writer)


def create_native_row_id_delete(table, plan=None):
    """Adapt the direct deletion API to the existing core table updater."""
    operation = create_native_delete(table, table.new_batch_write_builder().commit_user, plan)
    return None if operation is None else NativeTableDeleteByRowId(operation)


class NativeTableDeleteByRowId:

    def __init__(self, operation):
        self._operation = operation

    def delete(self, row_ids):
        return self._operation.delete_by_row_id(row_ids)


def _raise_native_row_id_error(error):
    detail = str(error)
    if 'duplicate UPDATE operations' in detail:
        raise ValueError('duplicate _ROW_ID: ' + detail) from error
    if 'No file found for _ROW_ID' in detail:
        raise ValueError(
            detail + ' does not belong to any valid range') from error
    raise error


class NativeBatchTableUpdate:
    """Wrap public core operations and decode their commit messages."""

    def __init__(self, table, writer):
        self.table = table
        self.writer = writer

    def update_by_arrow_with_row_id(self, data: pa.Table):
        try:
            messages = self.writer.update_by_arrow_with_row_id(data)
        except ValueError as error:
            _raise_native_row_id_error(error)
        return from_native_commit_messages(self.table, messages)

    def update_by_arrow_batches_with_row_id(self, tables):
        try:
            messages = self.writer.update_by_arrow_batches_with_row_id(tables)
        except ValueError as error:
            _raise_native_row_id_error(error)
        return from_native_commit_messages(self.table, messages)

    def delete_by_row_id(self, row_ids):
        ids = normalize_row_ids(row_ids)
        try:
            messages = self.writer.delete_by_row_id(ids)
        except ValueError as error:
            _raise_native_row_id_error(error)
        return from_native_commit_messages(self.table, messages)


class NativeTableUpdateByRowId(RowIdUpdateInputMixin):
    """Adapt Python row inputs around the core incremental Arrow updater."""

    def __init__(self, table, commit_user, commit_identifier, writer):
        self.table = table
        self.commit_user = commit_user
        self.commit_identifier = commit_identifier
        self.writer = writer

    @property
    def commit_messages(self):
        return from_native_commit_messages(self.table, self.writer.commit_messages)

    def update_columns(self, data, column_names):
        try:
            messages = self.writer.update_columns(data, column_names)
        except ValueError as error:
            _raise_native_row_id_error(error)
        return from_native_commit_messages(self.table, messages)

    def _write_row_columns(self, data, column_names, blob_object_columns):
        # Row input normalization happens in the public Python API; physical
        # field layouts and Blob references are handled by the core updater.
        if not blob_object_columns:
            return self.update_columns(data, column_names)
        from pypaimon.execution.native.write.native_blob_rows import NativeBlobRows
        blobs = NativeBlobRows()
        for name, values in blob_object_columns.items():
            field = self.table.field_dict[name]
            arrow_field = PyarrowFieldParser.from_paimon_field(field)
            if isinstance(field.type, MapType):
                # Arrow Map disallows NULL keys; Java's Blob map records allow
                # one. Core accepts row entries as a list of nullable-key structs.
                entries = pa.struct([pa.field('key', arrow_field.type.key_type),
                                     pa.field('value', arrow_field.type.item_type)])
                arrow_field = pa.field(name, pa.list_(entries), nullable=arrow_field.nullable)
            data = data.append_column(arrow_field, pa.array(
                [blobs.encode(value, field.type) for value in values], type=arrow_field.type))
        self.writer.with_blob_uri_reader_factory(blobs)
        try:
            return self.update_columns(data, column_names)
        finally:
            self.writer.with_blob_uri_reader_factory(None)


class NativeTableUpsert:
    """Submit Arrow tables or row-shape batches to the core Rust upsert writer."""

    def __init__(self, table, writer, keys, data, blob_rows=None):
        self.table = table
        self.writer = writer
        self.keys = keys
        self.data = data
        self.blob_rows = blob_rows

    def upsert(self):
        try:
            return from_native_commit_messages(
                self.table,
                self.writer.upsert_by_arrow_with_key(self.data, self.keys))
        finally:
            if self.blob_rows is not None:
                self.blob_rows.clear()


class NativePredicateTableUpdate:
    """Convert operation inputs and commit messages around the core updater."""

    def __init__(self, table, writer, predicate):
        self.table = table
        self.writer = writer
        self.predicate = predicate

    def update(self, assignments, read_columns):
        return from_native_commit_messages(
            self.table,
            self.writer.update_by_predicate(
                self.predicate, dict(assignments), list(read_columns or ())))

    def delete(self):
        return from_native_commit_messages(
            self.table, self.writer.delete_by_predicate(self.predicate))