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

from typing import List, Optional

import pyarrow as pa

from pypaimon.common.options.core_options import ChangelogProducer
from pypaimon.schema.arrow_schema import (
    arrow_schemas_compatible, normalize_arrow_strings, normalize_arrow_write_nullability, project_arrow_batch)
from pypaimon.schema.data_types import PyarrowFieldParser
from pypaimon.table.row.blob import BlobConsumer
from pypaimon.write.row_utils import (
    require_columns,
    row_to_named_values,
    validate_write_columns,
)
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.commit_identifier import validate_commit_identifier
from pypaimon.execution.python.write.file_store_write import FileStoreWrite

from pypaimon.write.input_adapter import _pandas_write_batch


class PythonTableWrite:
    """Bound Python writer implementing the common write execution protocol."""

    def __init__(self, public_writer):
        self.public_writer = public_writer

    def initialize(self, table, commit_user, static_partition: Optional[dict] = None,
                   *, restore_snapshot_id=None):
        from pypaimon.table.file_store_table import FileStoreTable

        self.table: FileStoreTable = table
        self.table_pyarrow_schema = PyarrowFieldParser.from_paimon_schema(self.table.table_schema.fields)
        self.commit_user = commit_user
        self.static_partition = static_partition
        self.restore_snapshot_id = restore_snapshot_id
        self.file_store_write = self._create_file_store_write(commit_user)
        # Java overwrite writers ignore previous files even without an explicit
        # restore snapshot. HASH restoration is separately disabled by the row
        # key extractor; sequence restoration must also start from empty state.
        self.file_store_write.restore_snapshot_id = (
            0 if static_partition is not None else restore_snapshot_id)
        if static_partition is not None:
            # An overwrite replaces state, not an input changelog. Java's
            # overwrite commit does not publish changelog manifests; avoid
            # writing unreferenced changelog files in the first place.
            self.file_store_write.changelog_producer = ChangelogProducer.NONE
        self.row_key_extractor = self._create_row_key_extractor(static_partition)

    def _create_file_store_write(self, commit_user):
        return FileStoreWrite(self.table, commit_user)

    def _create_row_key_extractor(self, static_partition):
        if self.restore_snapshot_id is not None:
            from pypaimon.execution.python.write.row_key_extractor import DynamicBucketRowKeyExtractor
            return DynamicBucketRowKeyExtractor(
                self.table.table_schema, table=self.table,
                ignore_existing=static_partition is not None,
                base_snapshot_id=self.restore_snapshot_id)
        from pypaimon.execution.python.write.row_key_extractor import create_row_key_extractor
        return create_row_key_extractor(self.table, ignore_existing=static_partition is not None)

    def write_arrow(self, data, bucket: Optional[int] = None):
        """Write Arrow data, optionally to a bucket assigned upstream."""
        data = self._prepare_arrow_schema(data)
        # Public batch overrides remain observable for Arrow and pandas input.
        # Schema, routing and preparation helpers stay within this executor.
        if isinstance(data, pa.RecordBatch):
            return self.public_writer.write_arrow_batch(data, bucket)
        for batch in data.to_batches():
            self.public_writer.write_arrow_batch(batch, bucket)

    def write_arrow_batch(self, data: pa.RecordBatch, bucket: Optional[int] = None):
        if bucket is not None:
            return self._write_precomputed_bucket(data, bucket)
        data = self._prepare_arrow_data(data)
        if data.num_rows == 0:
            return

        groups = self.row_key_extractor.extract_partition_bucket_groups(data)
        for partition, bucket, row_indices in groups:
            if row_indices is None:
                # Every input row belongs to the same partition/bucket. Passing the
                # original batch through avoids copying large BLOB values through
                # Arrow take before the dedicated BLOB writer consumes them.
                sub_table = data
            else:
                # row_indices is an int64 array of this group's rows in
                # ascending input order (the extractor sorts grouped indices so
                # sequence-number assignment stays latest-wins correct), so the
                # span is just first..last.
                lo = row_indices[0].as_py()
                hi = row_indices[-1].as_py()
                count = len(row_indices)
                if hi - lo + 1 == count:
                    # Distinct row indices spanning exactly `count` values are
                    # contiguous, so share the original Arrow buffers instead of
                    # gathering their rows into newly allocated buffers with take.
                    sub_table = data.slice(lo, count)
                else:
                    sub_table = pa.compute.take(data, row_indices)
            self._write_partition_bucket_batch(partition, bucket, sub_table)

    def roll_before_group_if_needed(self, row_count: int):
        """Roll current files before the next logical write group if needed."""
        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.")
        self.file_store_write.roll_before_group_if_needed(row_count)

    def _write_partition_bucket_batch(self, partition, bucket, data):
        self.file_store_write.write(partition, bucket, data)

    def _write_precomputed_bucket(self, data, bucket):
        """Write one partition/bucket group using normal index maintenance."""
        from pypaimon.table.bucket_mode import BucketMode
        from pypaimon.execution.python.write.row_key_extractor import DynamicBucketRowKeyExtractor

        bucket_mode = self.table.bucket_mode()
        if bucket_mode not in (BucketMode.HASH_FIXED, BucketMode.HASH_DYNAMIC):
            raise ValueError(
                "Precomputed bucket writes are only valid for HASH_FIXED or "
                "HASH_DYNAMIC tables"
            )
        if not isinstance(self.row_key_extractor, DynamicBucketRowKeyExtractor):
            if bucket_mode == BucketMode.HASH_DYNAMIC:
                raise RuntimeError("Dynamic bucket extractor is not configured")
        from pypaimon.index.dynamic_bucket import MAX_DYNAMIC_BUCKETS
        upper_bound = MAX_DYNAMIC_BUCKETS if bucket_mode == BucketMode.HASH_DYNAMIC else self.table.options.bucket()
        if not 0 <= bucket < upper_bound:
            raise ValueError('Bucket id must be between 0 and {}'.format(upper_bound - 1))
        data = self._prepare_arrow_data(data)
        if bucket_mode == BucketMode.HASH_DYNAMIC:
            partition = self.row_key_extractor.notify_precomputed_bucket_batch(data, bucket)
        else:
            partitions = self.row_key_extractor.extract_partitions_batch(data)
            if not partitions:
                return
            partition = tuple(partitions[0])
            if any(tuple(actual) != partition for actual in partitions):
                raise ValueError('A precomputed bucket group contained multiple partitions')
        if partition is None:
            return
        self._write_partition_bucket_batch(partition, bucket, data)

    def write_row(self, row):
        values_by_name = row_to_named_values(row, self.table.table_schema.fields)
        require_columns(values_by_name, self.table.partition_keys, "write_row")
        from pypaimon.write.row_kind import skip_write_row
        write_cols = (self.file_store_write.write_cols
                      if set(values_by_name) == set(self.table.field_names) else list(values_by_name))
        if skip_write_row(self.table, values_by_name, row.get_row_kind(), write_cols):
            return
        partition, bucket = (
            self.row_key_extractor.extract_partition_bucket_row(values_by_name)
        )
        self._write_partition_bucket_row(
            partition, bucket, row, values_by_name
        )

    def _write_partition_bucket_row(
        self, partition, bucket, row, values_by_name
    ):
        self.file_store_write.write_row(
            partition, bucket, row, values_by_name
        )

    def write_pandas(self, dataframe):
        write_cols = self.file_store_write.write_cols
        if write_cols is not None:
            # Column-subset write (append-only ``with_write_type``): build the
            # RecordBatch against the subset schema so the input only needs the
            # written columns, mirroring the ``write_arrow`` path.
            pa_schema = self._write_cols_pyarrow_schema(write_cols)
        else:
            pa_schema = self.table_pyarrow_schema
        record_batch = _pandas_write_batch(
            getattr(self, 'table', None), dataframe, pa_schema, self.table_pyarrow_schema)
        return self.public_writer.write_arrow_batch(record_batch)

    def with_write_type(self, write_cols: List[str]):
        write_cols = validate_write_columns(write_cols)
        for col in write_cols:
            if col not in self.table_pyarrow_schema.names:
                raise ValueError(f"Column {col} is not in table schema.")
        full_names = self.table_pyarrow_schema.names
        # Java preserves the declared append column order in writeCols. PK
        # writes still require the complete table layout in the Python buffer.
        full_pk_type = (self.table.is_primary_key_table and len(write_cols) == len(full_names)
                        and set(write_cols) == set(full_names))
        if write_cols == full_names or full_pk_type:
            write_cols = None
        if self.file_store_write._ever_written:
            if write_cols == self.file_store_write.write_cols:
                return self
            raise RuntimeError('with_write_type must be called before any write operation')
        self.file_store_write.write_cols = write_cols
        return self

    def with_blob_consumer(self, blob_consumer: BlobConsumer):
        if blob_consumer is not None and not callable(blob_consumer):
            raise TypeError('blob_consumer must be callable')
        self.file_store_write.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.file_store_write._ever_written:
            raise RuntimeError(
                "with_blob_uri_reader_factory must be called before any "
                "write operation."
            )
        self.file_store_write.blob_uri_reader_factory = uri_reader_factory
        return self

    def close(self):
        try:
            self.file_store_write.close()
        finally:
            self._release_prepared_indexes()

    def abort(self):
        try:
            self.file_store_write.abort()
        finally:
            abort = getattr(self.row_key_extractor, "abort", None)
            if abort is not None:
                abort()

    def prepare_commit(self, commit_identifier) -> List[CommitMessage]:
        commit_identifier = validate_commit_identifier(commit_identifier)
        commit_messages = self.file_store_write._prepare_commit_messages(commit_identifier)
        prepare_indexes = getattr(self.row_key_extractor, "prepare_commit", None)
        index_changes = prepare_indexes() if prepare_indexes is not None else {}
        messages_by_bucket = {
            (tuple(message.partition), message.bucket): message
            for message in commit_messages
        }
        for (partition, bucket), changes in index_changes.items():
            message = messages_by_bucket.get((partition, bucket))
            if message is None:
                message = CommitMessage(
                    partition=partition,
                    bucket=bucket,
                    new_files=[],
                )
                commit_messages.append(message)
                messages_by_bucket[(partition, bucket)] = message
            message.index_adds.extend(changes.additions)
            message.index_deletes.extend(changes.deletions)
        # Hand off data, changelog and HASH indexes together. Until every
        # preparation succeeds, retain the metadata for retry and abort.
        self.file_store_write._release_prepared_files()
        self._release_prepared_indexes()
        return commit_messages

    def _release_prepared_indexes(self) -> None:
        release = getattr(self.row_key_extractor, "release_prepared", None)
        if release is not None:
            release()

    def _prepare_arrow_data(self, data):
        data = self._prepare_arrow_schema(data)
        # Schema-only writers can normalize input without table options.
        table = getattr(self, 'table', None)
        if table is None:
            return data
        data = self._normalize_arrow_write_schema(data)
        from pypaimon.write.row_kind import filter_write_batch
        selected = (self.file_store_write.write_cols if not table.is_primary_key_table else None)
        write_cols = selected
        if selected is None:
            write_cols = (self.file_store_write.write_cols
                          if data.schema.names == self.table_pyarrow_schema.names else data.schema.names)
        data = filter_write_batch(table, data, write_cols)
        if not data.num_rows:
            return data
        if selected is not None:
            routing = self._write_routing_columns()
            columns = selected + [name for name in data.schema.names
                                  if name in routing and name not in selected]
            data = self._project_write_input(data, columns)
            # Routing fields are not written values. Arrow's whole-table
            # string cast must not reapply their declared top-level NOT NULL.
            schema = pa.schema([field if field.name in selected else field.with_nullable(True)
                                for field in data.schema], metadata=data.schema.metadata)
            data = (pa.RecordBatch.from_arrays(data.columns, schema=schema) if isinstance(data, pa.RecordBatch)
                    else pa.Table.from_arrays(data.columns, schema=schema))
        data = normalize_arrow_strings(data)
        physical = self._project_write_input(data, selected) if selected is not None else data
        target = self.table_pyarrow_schema
        if physical.schema.names != target.names:
            target = self._write_cols_pyarrow_schema(physical.schema.names)
        # Validate only surviving values. Keep the input transport metadata:
        # PyArrow otherwise rejects hidden NULL children of NULL ROW parents.
        validated = normalize_arrow_write_nullability(physical, target)
        self._validate_pyarrow_schema(validated.schema)
        return data

    @staticmethod
    def _project_write_input(data, columns):
        return (project_arrow_batch(data, columns) if isinstance(data, pa.RecordBatch)
                else data.select(columns))

    def _write_routing_columns(self):
        table = self.table
        routing = set(table.partition_keys)
        if table.options.bucket() > 0:
            routing.update(table.table_schema.bucket_keys)
        row_kind = table.options.options.to_map().get('rowkind.field')
        if row_kind is not None:
            routing.add(row_kind)
        return routing

    def _prepare_arrow_schema(self, data):
        # Legacy schema-only callers have no row-kind options to evaluate.
        # A real writer must defer casts until ignored events have been removed.
        if getattr(self, 'table', None) is None:
            data = normalize_arrow_strings(data)
        self._validate_pyarrow_schema(
            data.schema, check_nested_nullability=getattr(self, 'table', None) is None)
        return data

    def _normalize_arrow_write_schema(self, data):
        # Compatible input batches may differ in top-level nullability or
        # metadata. Buffered writes and row-to-Arrow conversion must use the
        # same declared fields, while keeping normalized Arrow array types.
        fields = []
        for field in data.schema:
            declared = self.table_pyarrow_schema.field(field.name)
            fields.append(pa.field(field.name, field.type, nullable=declared.nullable,
                                   metadata=declared.metadata))
        schema = pa.schema(fields, metadata=self.table_pyarrow_schema.metadata)
        from pypaimon.write.row_kind import write_value_schema
        schema = write_value_schema(self.table, schema)
        if data.schema.equals(schema, check_metadata=True):
            return data
        if isinstance(data, pa.RecordBatch):
            return pa.RecordBatch.from_arrays(data.columns, schema=schema)
        return pa.Table.from_arrays(data.columns, schema=schema)

    def _validate_pyarrow_schema(self, data_schema: pa.Schema, *, check_nested_nullability=True):
        if self._is_compatible_pyarrow_schema(
                data_schema, self.table_pyarrow_schema, check_nested_nullability=check_nested_nullability):
            return

        write_cols = self.file_store_write.write_cols
        if write_cols is not None:
            write_cols_schema = self._write_cols_pyarrow_schema(write_cols)
            if self._is_compatible_pyarrow_schema(
                    data_schema, write_cols_schema, check_nested_nullability=check_nested_nullability):
                return
            table = getattr(self, 'table', None)
            if table is not None and not table.is_primary_key_table:
                names = data_schema.names
                selected = set(write_cols)
                allowed = selected | self._write_routing_columns()
                if (len(set(names)) == len(names) and set(names).issubset(allowed)
                        and [name for name in names if name in selected] == write_cols):
                    transport_schema = self._write_cols_pyarrow_schema(names)
                    if self._is_compatible_pyarrow_schema(
                            data_schema, transport_schema, check_nested_nullability=check_nested_nullability):
                        return

        self._raise_inconsistent_schema(data_schema)

    def _is_compatible_pyarrow_schema(
            self, data_schema: pa.Schema, expected_schema: pa.Schema, *, check_nested_nullability=True) -> bool:
        return arrow_schemas_compatible(
            data_schema, expected_schema, check_top_level_nullability=False,
            allow_binary_compatibility=True, check_nested_nullability=check_nested_nullability)

    def _write_cols_pyarrow_schema(self, write_cols: List[str]) -> pa.Schema:
        table_fields = {
            field.name: field for field in self.table_pyarrow_schema
        }
        return pa.schema([table_fields[col] for col in write_cols])

    def _raise_inconsistent_schema(self, data_schema: pa.Schema):
        raise ValueError(f"Input schema isn't consistent with table schema and write cols. "
                         f"Input schema is: {data_schema} "
                         f"Table schema is: {self.table_pyarrow_schema} "
                         f"Write cols is: {self.file_store_write.write_cols}")