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

"""Pure write schema validation is independent of Native installation."""

import sys
from types import ModuleType
from unittest.mock import Mock, patch

import pyarrow as pa
import pytest

from pypaimon.schema.arrow_schema import normalize_arrow_write_nullability


def _source(array, batch=False):
    schema = pa.schema([pa.field('nested', array.type, metadata={b'field': b'input'})],
                       metadata={b'schema': b'input'})
    return (pa.RecordBatch if batch else pa.Table).from_arrays([array], schema=schema)


def _target(data_type):
    return pa.schema([pa.field('nested', data_type)])


@pytest.mark.parametrize('installed', [False, True])
@pytest.mark.parametrize('batch', [False, True])
def test_visible_child_validation_never_loads_native(installed, batch):
    sdk = ModuleType('pypaimon_rust.data')
    sdk.normalize_array_for_write = Mock(side_effect=AssertionError('Native normalization called'))
    modules = {'pypaimon.execution.native.runtime': None,
               'pypaimon_rust': ModuleType('pypaimon_rust') if installed else None,
               'pypaimon_rust.data': sdk if installed else None}
    target = _target(pa.list_(pa.field('item', pa.int32(), nullable=False)))
    data = _source(pa.array([[1, 2], None, []], type=pa.list_(pa.int32())), batch)
    with patch.dict(sys.modules, modules):
        result = normalize_arrow_write_nullability(data, target)
        assert result.to_pydict() == data.to_pydict()
        assert result.schema.field(0).type == target.field(0).type
        with pytest.raises(ValueError, match='nested.item.*NULL'):
            normalize_arrow_write_nullability(_source(pa.array([[1, None]], type=pa.list_(pa.int32())), batch), target)
    sdk.normalize_array_for_write.assert_not_called()


@pytest.mark.parametrize('kind', ['list', 'large-list', 'fixed-size-list', 'map', 'struct'])
def test_null_parents_and_slices_hide_unselected_null_children(kind):
    child = pa.field('item', pa.int32(), nullable=False)
    if kind == 'struct':
        data = pa.StructArray.from_arrays(
            [pa.array([None, None, 7, None], type=pa.int32())],
            names=['item'], mask=pa.array([False, True, False, False]))
        data_type = pa.struct([child])
    elif kind == 'map':
        offsets = pa.array([0, 1, 2, 3, 4], type=pa.int32())
        data = pa.MapArray.from_arrays(
            offsets, pa.array(['outside', 'hidden', 'visible', 'outside']),
            pa.array([None, None, 7, None], type=pa.int32()),
            mask=pa.array([False, True, False, False]))
        data_type = pa.map_(pa.string(), pa.field('value', pa.int32(), nullable=False))
    else:
        values = pa.array([None, None, 7, None], type=pa.int32())
        mask = pa.array([False, True, False, False])
        if kind == 'fixed-size-list':
            data = pa.FixedSizeListArray.from_arrays(values, 1, mask=mask)
            data_type = pa.list_(child, 1)
        else:
            large = kind == 'large-list'
            cls = pa.LargeListArray if large else pa.ListArray
            offsets = pa.array([0, 1, 2, 3, 4], type=pa.int64() if large else pa.int32())
            data = cls.from_arrays(offsets, values, mask=mask)
            data_type = pa.large_list(child) if large else pa.list_(child)
    source = _source(data.slice(1, 2))
    result = normalize_arrow_write_nullability(source, _target(data_type))
    assert result.to_pydict() == source.to_pydict()
    assert result.schema.field(0).type == data_type
    assert result.schema.metadata == source.schema.metadata
    assert result.schema.field(0).metadata == source.schema.field(0).metadata
    result.validate(full=True)
    with pytest.raises(ValueError, match='NULL'):
        normalize_arrow_write_nullability(_source(data.slice(2, 2)), _target(data_type))


def test_null_struct_masks_nested_non_nullable_containers():
    lists = pa.ListArray.from_arrays(pa.array([0, 1, 2]), pa.array([None, 7], type=pa.int32()))
    maps = pa.MapArray.from_arrays(
        pa.array([0, 1, 2]), pa.array(['hidden', 'visible']),
        pa.array([None, 8], type=pa.int32()))
    source = _source(pa.StructArray.from_arrays(
        [lists, maps], names=['items', 'attrs'], mask=pa.array([True, False])))
    target = _target(pa.struct([
        pa.field('items', pa.list_(pa.field('item', pa.int32(), nullable=False)), nullable=False),
        pa.field('attrs', pa.map_(pa.string(), pa.field('value', pa.int32(), nullable=False)), nullable=False),
    ]))
    result = normalize_arrow_write_nullability(source, target)
    assert result.to_pydict() == source.to_pydict()
    assert result.schema.field(0).type == target.field(0).type


def test_chunk_boundaries_and_payload_buffers_are_retained():
    data_type = pa.list_(pa.field('item', pa.int32(), nullable=False))
    chunks = [pa.array([[1, 2], [3]], type=pa.list_(pa.int32())),
              pa.array([None, [4]], type=pa.list_(pa.int32()))]
    data = pa.table({'nested': pa.chunked_array(chunks)})
    result = normalize_arrow_write_nullability(data, _target(data_type))
    assert [len(chunk) for chunk in result.column(0).chunks] == [2, 2]
    assert result.to_pydict() == data.to_pydict()
    for actual, original in zip(result.column(0).chunks, chunks):
        assert actual.values.buffers()[1].address == original.values.buffers()[1].address


def test_nullability_restoration_does_not_cast_other_layouts():
    data = pa.table({'nested': pa.array([[1]], type=pa.list_(pa.int32())),
                     'text': pa.array(['value'], type=pa.large_string()),
                     'binary': pa.array([b'x'], type=pa.binary(1))})
    target = pa.schema([('nested', pa.list_(pa.field('item', pa.int32(), nullable=False))),
                        ('text', pa.string()), ('binary', pa.binary())])
    result = normalize_arrow_write_nullability(data, target)
    assert result.schema.field('text').type == pa.large_string()
    assert result.schema.field('binary').type == pa.binary(1)
    incompatible = pa.schema([('nested', pa.list_(pa.field('item', pa.int64(), nullable=False))),
                              ('text', pa.string()), ('binary', pa.binary())])
    assert normalize_arrow_write_nullability(data, incompatible) is data


def test_empty_chunked_input_restores_only_declared_child_fields():
    source = pa.table({'nested': pa.chunked_array([], type=pa.list_(pa.int32()))})
    data_type = pa.list_(pa.field('item', pa.int32(), nullable=False))
    result = normalize_arrow_write_nullability(source, _target(data_type))
    assert result.num_rows == 0 and result.column(0).num_chunks == 0
    assert result.schema.field(0).type == data_type


def test_python_writer_accepts_visible_non_null_children_without_native(tmp_path):
    from pypaimon import CatalogFactory, Schema

    declared = pa.schema([('nested', pa.struct([pa.field('value', pa.int32(), nullable=False)]))])
    catalog = CatalogFactory.create({'warehouse': str(tmp_path / 'warehouse')})
    catalog.create_database('default', True)
    catalog.create_table('default.nested', Schema.from_pyarrow_schema(declared, options={
        'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false',
        'write.native.enabled': 'false', 'commit.native.enabled': 'false',
    }), False)
    table = catalog.get_table('default.nested')
    source = _source(pa.StructArray.from_arrays(
        [pa.array([None, 7], type=pa.int32())], names=['value'], mask=pa.array([True, False])))
    builder = table.new_batch_write_builder()
    writer, commit = builder.new_write(), builder.new_commit()
    try:
        with patch.dict(sys.modules, {'pypaimon.execution.native.runtime': None, 'pypaimon_rust': None}):
            writer.write_arrow(source)
            commit.commit(writer.prepare_commit())
    finally:
        writer.close()
        commit.close()
    reader = table.new_read_builder()
    result = reader.new_read().to_arrow(reader.new_scan().plan().splits())
    assert result.column('nested').to_pylist() == [None, {'value': 7}]