# 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.
# pylint: disable=protected-access,unused-argument,redefined-outer-name

import math
import tempfile
import uuid
from dataclasses import asdict, dataclass
from datetime import (
    date,
    datetime,
    time,
    timedelta,
    timezone,
)
from decimal import Decimal
from typing import (
    Any,
    Dict,
    List,
    Optional,
    Tuple,
    Union,
)

import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from pyiceberg.avro import (
    STRUCT_BOOL,
    STRUCT_DOUBLE,
    STRUCT_FLOAT,
    STRUCT_INT32,
    STRUCT_INT64,
)
from pyiceberg.io.pyarrow import (
    MetricModeTypes,
    MetricsMode,
    PyArrowStatisticsCollector,
    compute_statistics_plan,
    data_file_statistics_from_parquet_metadata,
    match_metrics_mode,
    parquet_path_to_id_mapping,
    schema_to_pyarrow,
)
from pyiceberg.manifest import DataFile
from pyiceberg.schema import Schema, pre_order_visit
from pyiceberg.table.metadata import (
    TableMetadata,
    TableMetadataUtil,
    TableMetadataV1,
    TableMetadataV2,
)
from pyiceberg.types import (
    BooleanType,
    FloatType,
    IntegerType,
    StringType,
)
from pyiceberg.utils.datetime import date_to_days, datetime_to_micros, time_to_micros


@dataclass(frozen=True)
class TestStruct:
    __test__ = False
    x: Optional[int]
    y: Optional[float]


def construct_test_table(
    write_statistics: Union[bool, List[str]] = True,
) -> Tuple[pq.FileMetaData, Union[TableMetadataV1, TableMetadataV2]]:
    table_metadata = {
        "format-version": 2,
        "location": "s3://bucket/test/location",
        "last-column-id": 7,
        "current-schema-id": 0,
        "schemas": [
            {
                "type": "struct",
                "schema-id": 0,
                "fields": [
                    {"id": 1, "name": "strings", "required": False, "type": "string"},
                    {"id": 2, "name": "floats", "required": False, "type": "float"},
                    {
                        "id": 3,
                        "name": "list",
                        "required": False,
                        "type": {"type": "list", "element-id": 6, "element": "long", "element-required": False},
                    },
                    {
                        "id": 4,
                        "name": "maps",
                        "required": False,
                        "type": {
                            "type": "map",
                            "key-id": 7,
                            "key": "long",
                            "value-id": 8,
                            "value": "long",
                            "value-required": False,
                        },
                    },
                    {
                        "id": 5,
                        "name": "structs",
                        "required": False,
                        "type": {
                            "type": "struct",
                            "fields": [
                                {"id": 9, "name": "x", "required": False, "type": "long"},
                                {"id": 10, "name": "y", "required": False, "type": "float", "doc": "comment"},
                            ],
                        },
                    },
                ],
            },
        ],
        "default-spec-id": 0,
        "partition-specs": [{"spec-id": 0, "fields": []}],
        "properties": {},
    }

    table_metadata = TableMetadataUtil.parse_obj(table_metadata)
    arrow_schema = schema_to_pyarrow(table_metadata.schemas[0])

    _strings = ["zzzzzzzzzzzzzzzzzzzz", "rrrrrrrrrrrrrrrrrrrr", None, "aaaaaaaaaaaaaaaaaaaa"]

    _floats = [3.14, math.nan, 1.69, 100]

    _list = [[1, 2, 3], [4, 5, 6], None, [7, 8, 9]]

    _maps: List[Optional[Dict[int, int]]] = [
        {1: 2, 3: 4},
        None,
        {5: 6},
        {},
    ]

    _structs = [
        asdict(TestStruct(1, 0.2)),
        asdict(TestStruct(None, -1.34)),
        None,
        asdict(TestStruct(54, None)),
    ]

    table = pa.Table.from_pydict(
        {
            "strings": _strings,
            "floats": _floats,
            "list": _list,
            "maps": _maps,
            "structs": _structs,
        },
        schema=arrow_schema,
    )
    metadata_collector: List[Any] = []

    with pa.BufferOutputStream() as f:
        with pq.ParquetWriter(
            f, table.schema, metadata_collector=metadata_collector, write_statistics=write_statistics
        ) as writer:
            writer.write_table(table)

    return metadata_collector[0], table_metadata


def get_current_schema(
    table_metadata: TableMetadata,
) -> Schema:
    return next(filter(lambda s: s.schema_id == table_metadata.current_schema_id, table_metadata.schemas))


def test_record_count() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())
    assert datafile.record_count == 4


def test_value_counts() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.value_counts) == 7
    assert datafile.value_counts[1] == 4
    assert datafile.value_counts[2] == 4
    assert datafile.value_counts[6] == 10  # 3 lists with 3 items and a None value
    assert datafile.value_counts[7] == 5
    assert datafile.value_counts[8] == 5
    assert datafile.value_counts[9] == 4
    assert datafile.value_counts[10] == 4


def test_column_sizes() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.column_sizes) == 7
    # these values are an artifact of how the write_table encodes the columns
    assert datafile.column_sizes[1] > 0
    assert datafile.column_sizes[2] > 0
    assert datafile.column_sizes[6] > 0
    assert datafile.column_sizes[7] > 0
    assert datafile.column_sizes[8] > 0


def test_null_and_nan_counts() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.null_value_counts) == 7
    assert datafile.null_value_counts[1] == 1
    assert datafile.null_value_counts[2] == 0
    assert datafile.null_value_counts[6] == 1
    assert datafile.null_value_counts[7] == 2
    assert datafile.null_value_counts[8] == 2
    assert datafile.null_value_counts[9] == 2
    assert datafile.null_value_counts[10] == 2

    # #arrow does not include this in the statistics
    # assert len(datafile.nan_value_counts)  == 3
    # assert datafile.nan_value_counts[1]    == 0
    # assert datafile.nan_value_counts[2]    == 1
    # assert datafile.nan_value_counts[3]    == 0


def test_bounds() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.lower_bounds) == 2
    assert datafile.lower_bounds[1].decode() == "aaaaaaaaaaaaaaaa"
    assert datafile.lower_bounds[2] == STRUCT_FLOAT.pack(1.69)

    assert len(datafile.upper_bounds) == 2
    assert datafile.upper_bounds[1].decode() == "zzzzzzzzzzzzzzz{"
    assert datafile.upper_bounds[2] == STRUCT_FLOAT.pack(100)


def test_metrics_mode_parsing() -> None:
    assert match_metrics_mode("none") == MetricsMode(MetricModeTypes.NONE)
    assert match_metrics_mode("nOnE") == MetricsMode(MetricModeTypes.NONE)
    assert match_metrics_mode("counts") == MetricsMode(MetricModeTypes.COUNTS)
    assert match_metrics_mode("Counts") == MetricsMode(MetricModeTypes.COUNTS)
    assert match_metrics_mode("full") == MetricsMode(MetricModeTypes.FULL)
    assert match_metrics_mode("FuLl") == MetricsMode(MetricModeTypes.FULL)
    assert match_metrics_mode(" FuLl") == MetricsMode(MetricModeTypes.FULL)

    assert match_metrics_mode("truncate(16)") == MetricsMode(MetricModeTypes.TRUNCATE, 16)
    assert match_metrics_mode("trUncatE(16)") == MetricsMode(MetricModeTypes.TRUNCATE, 16)
    assert match_metrics_mode("trUncatE(7)") == MetricsMode(MetricModeTypes.TRUNCATE, 7)
    assert match_metrics_mode("trUncatE(07)") == MetricsMode(MetricModeTypes.TRUNCATE, 7)

    with pytest.raises(ValueError) as exc_info:
        match_metrics_mode("trUncatE(-7)")
    assert "Malformed truncate: trUncatE(-7)" in str(exc_info.value)

    with pytest.raises(ValueError) as exc_info:
        match_metrics_mode("trUncatE(0)")
    assert "Truncation length must be larger than 0" in str(exc_info.value)


def test_metrics_mode_none() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    table_metadata.properties["write.metadata.metrics.default"] = "none"
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.value_counts) == 0
    assert len(datafile.null_value_counts) == 0
    assert len(datafile.nan_value_counts) == 0
    assert len(datafile.lower_bounds) == 0
    assert len(datafile.upper_bounds) == 0


def test_metrics_mode_counts() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    table_metadata.properties["write.metadata.metrics.default"] = "counts"
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.value_counts) == 7
    assert len(datafile.null_value_counts) == 7
    assert len(datafile.nan_value_counts) == 0
    assert len(datafile.lower_bounds) == 0
    assert len(datafile.upper_bounds) == 0


def test_metrics_mode_full() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    table_metadata.properties["write.metadata.metrics.default"] = "full"
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.value_counts) == 7
    assert len(datafile.null_value_counts) == 7
    assert len(datafile.nan_value_counts) == 0

    assert len(datafile.lower_bounds) == 2
    assert datafile.lower_bounds[1].decode() == "aaaaaaaaaaaaaaaaaaaa"
    assert datafile.lower_bounds[2] == STRUCT_FLOAT.pack(1.69)

    assert len(datafile.upper_bounds) == 2
    assert datafile.upper_bounds[1].decode() == "zzzzzzzzzzzzzzzzzzzz"
    assert datafile.upper_bounds[2] == STRUCT_FLOAT.pack(100)


def test_metrics_mode_non_default_trunc() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    table_metadata.properties["write.metadata.metrics.default"] = "truncate(2)"
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.value_counts) == 7
    assert len(datafile.null_value_counts) == 7
    assert len(datafile.nan_value_counts) == 0

    assert len(datafile.lower_bounds) == 2
    assert datafile.lower_bounds[1].decode() == "aa"
    assert datafile.lower_bounds[2] == STRUCT_FLOAT.pack(1.69)

    assert len(datafile.upper_bounds) == 2
    assert datafile.upper_bounds[1].decode() == "z{"
    assert datafile.upper_bounds[2] == STRUCT_FLOAT.pack(100)


def test_column_metrics_mode() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    table_metadata.properties["write.metadata.metrics.default"] = "truncate(2)"
    table_metadata.properties["write.metadata.metrics.column.strings"] = "none"
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.value_counts) == 6
    assert len(datafile.null_value_counts) == 6
    assert len(datafile.nan_value_counts) == 0

    assert len(datafile.lower_bounds) == 1
    assert datafile.lower_bounds[2] == STRUCT_FLOAT.pack(1.69)
    assert 1 not in datafile.lower_bounds

    assert len(datafile.upper_bounds) == 1
    assert datafile.upper_bounds[2] == STRUCT_FLOAT.pack(100)
    assert 1 not in datafile.upper_bounds


def construct_test_table_primitive_types() -> Tuple[pq.FileMetaData, Union[TableMetadataV1, TableMetadataV2]]:
    table_metadata = {
        "format-version": 2,
        "location": "s3://bucket/test/location",
        "last-column-id": 7,
        "current-schema-id": 0,
        "schemas": [
            {
                "type": "struct",
                "schema-id": 0,
                "fields": [
                    {"id": 1, "name": "booleans", "required": False, "type": "boolean"},
                    {"id": 2, "name": "ints", "required": False, "type": "int"},
                    {"id": 3, "name": "longs", "required": False, "type": "long"},
                    {"id": 4, "name": "floats", "required": False, "type": "float"},
                    {"id": 5, "name": "doubles", "required": False, "type": "double"},
                    {"id": 6, "name": "dates", "required": False, "type": "date"},
                    {"id": 7, "name": "times", "required": False, "type": "time"},
                    {"id": 8, "name": "timestamps", "required": False, "type": "timestamp"},
                    {"id": 9, "name": "timestamptzs", "required": False, "type": "timestamptz"},
                    {"id": 10, "name": "strings", "required": False, "type": "string"},
                    {"id": 11, "name": "uuids", "required": False, "type": "uuid"},
                    {"id": 12, "name": "binaries", "required": False, "type": "binary"},
                    {"id": 13, "name": "decimal8", "required": False, "type": "decimal(5, 2)"},
                    {"id": 14, "name": "decimal16", "required": False, "type": "decimal(16, 6)"},
                    {"id": 15, "name": "decimal32", "required": False, "type": "decimal(19, 6)"},
                ],
            },
        ],
        "default-spec-id": 0,
        "partition-specs": [{"spec-id": 0, "fields": []}],
        "properties": {},
    }

    table_metadata = TableMetadataUtil.parse_obj(table_metadata)
    arrow_schema = schema_to_pyarrow(table_metadata.schemas[0])
    tz = timezone(timedelta(seconds=19800))

    booleans = [True, False]
    ints = [23, 89]
    longs = [54, 2]
    floats = [454.1223, 24342.29]
    doubles = [8542.12, -43.9]
    dates = [date(2022, 1, 2), date(2023, 2, 4)]
    times = [time(17, 30, 34), time(13, 21, 4)]
    timestamps = [datetime(2022, 1, 2, 17, 30, 34, 399), datetime(2023, 2, 4, 13, 21, 4, 354)]
    timestamptzs = [datetime(2022, 1, 2, 17, 30, 34, 399, tz), datetime(2023, 2, 4, 13, 21, 4, 354, tz)]
    strings = ["hello", "world"]
    uuids = [uuid.uuid3(uuid.NAMESPACE_DNS, "foo").bytes, uuid.uuid3(uuid.NAMESPACE_DNS, "bar").bytes]
    binaries = [b"hello", b"world"]
    decimal8 = pa.array([Decimal("123.45"), Decimal("678.91")], pa.decimal128(8, 2))
    decimal16 = pa.array([Decimal("12345679.123456"), Decimal("67891234.678912")], pa.decimal128(16, 6))
    decimal32 = pa.array([Decimal("1234567890123.123456"), Decimal("9876543210703.654321")], pa.decimal128(19, 6))

    table = pa.Table.from_pydict(
        {
            "booleans": booleans,
            "ints": ints,
            "longs": longs,
            "floats": floats,
            "doubles": doubles,
            "dates": dates,
            "times": times,
            "timestamps": timestamps,
            "timestamptzs": timestamptzs,
            "strings": strings,
            "uuids": uuids,
            "binaries": binaries,
            "decimal8": decimal8,
            "decimal16": decimal16,
            "decimal32": decimal32,
        },
        schema=arrow_schema,
    )

    metadata_collector: List[Any] = []

    with pa.BufferOutputStream() as f:
        with pq.ParquetWriter(f, table.schema, metadata_collector=metadata_collector, store_decimal_as_integer=True) as writer:
            writer.write_table(table)

    return metadata_collector[0], table_metadata


def test_metrics_primitive_types() -> None:
    metadata, table_metadata = construct_test_table_primitive_types()

    schema = get_current_schema(table_metadata)
    table_metadata.properties["write.metadata.metrics.default"] = "truncate(2)"
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.value_counts) == 15
    assert len(datafile.null_value_counts) == 15
    assert len(datafile.nan_value_counts) == 0

    tz = timezone(timedelta(seconds=19800))

    assert len(datafile.lower_bounds) == 15
    assert datafile.lower_bounds[1] == STRUCT_BOOL.pack(False)
    assert datafile.lower_bounds[2] == STRUCT_INT32.pack(23)
    assert datafile.lower_bounds[3] == STRUCT_INT64.pack(2)
    assert datafile.lower_bounds[4] == STRUCT_FLOAT.pack(454.1223)
    assert datafile.lower_bounds[5] == STRUCT_DOUBLE.pack(-43.9)
    assert datafile.lower_bounds[6] == STRUCT_INT32.pack(date_to_days(date(2022, 1, 2)))
    assert datafile.lower_bounds[7] == STRUCT_INT64.pack(time_to_micros(time(13, 21, 4)))
    assert datafile.lower_bounds[8] == STRUCT_INT64.pack(datetime_to_micros(datetime(2022, 1, 2, 17, 30, 34, 399)))
    assert datafile.lower_bounds[9] == STRUCT_INT64.pack(datetime_to_micros(datetime(2022, 1, 2, 17, 30, 34, 399, tz)))
    assert datafile.lower_bounds[10] == b"he"
    assert datafile.lower_bounds[11] == uuid.uuid3(uuid.NAMESPACE_DNS, "foo").bytes
    assert datafile.lower_bounds[12] == b"he"
    assert datafile.lower_bounds[13][::-1].ljust(4, b"\x00") == STRUCT_INT32.pack(12345)
    assert datafile.lower_bounds[14][::-1].ljust(8, b"\x00") == STRUCT_INT64.pack(12345679123456)
    assert str(int.from_bytes(datafile.lower_bounds[15], byteorder="big", signed=True)).encode("utf-8") == b"1234567890123123456"

    assert len(datafile.upper_bounds) == 15
    assert datafile.upper_bounds[1] == STRUCT_BOOL.pack(True)
    assert datafile.upper_bounds[2] == STRUCT_INT32.pack(89)
    assert datafile.upper_bounds[3] == STRUCT_INT64.pack(54)
    assert datafile.upper_bounds[4] == STRUCT_FLOAT.pack(24342.29)
    assert datafile.upper_bounds[5] == STRUCT_DOUBLE.pack(8542.12)
    assert datafile.upper_bounds[6] == STRUCT_INT32.pack(date_to_days(date(2023, 2, 4)))
    assert datafile.upper_bounds[7] == STRUCT_INT64.pack(time_to_micros(time(17, 30, 34)))
    assert datafile.upper_bounds[8] == STRUCT_INT64.pack(datetime_to_micros(datetime(2023, 2, 4, 13, 21, 4, 354)))
    assert datafile.upper_bounds[9] == STRUCT_INT64.pack(datetime_to_micros(datetime(2023, 2, 4, 13, 21, 4, 354, tz)))
    assert datafile.upper_bounds[10] == b"wp"
    assert datafile.upper_bounds[11] == uuid.uuid3(uuid.NAMESPACE_DNS, "bar").bytes
    assert datafile.upper_bounds[12] == b"wp"
    assert datafile.upper_bounds[13][::-1].ljust(4, b"\x00") == STRUCT_INT32.pack(67891)
    assert datafile.upper_bounds[14][::-1].ljust(8, b"\x00") == STRUCT_INT64.pack(67891234678912)
    assert str(int.from_bytes(datafile.upper_bounds[15], byteorder="big", signed=True)).encode("utf-8") == b"9876543210703654321"


def construct_test_table_invalid_upper_bound() -> Tuple[pq.FileMetaData, Union[TableMetadataV1, TableMetadataV2]]:
    table_metadata = {
        "format-version": 2,
        "location": "s3://bucket/test/location",
        "last-column-id": 7,
        "current-schema-id": 0,
        "schemas": [
            {
                "type": "struct",
                "schema-id": 0,
                "fields": [
                    {"id": 1, "name": "valid_upper_binary", "required": False, "type": "binary"},
                    {"id": 2, "name": "invalid_upper_binary", "required": False, "type": "binary"},
                    {"id": 3, "name": "valid_upper_string", "required": False, "type": "string"},
                    {"id": 4, "name": "invalid_upper_string", "required": False, "type": "string"},
                ],
            },
        ],
        "default-spec-id": 0,
        "partition-specs": [{"spec-id": 0, "fields": []}],
        "properties": {},
    }

    table_metadata = TableMetadataUtil.parse_obj(table_metadata)
    arrow_schema = schema_to_pyarrow(table_metadata.schemas[0])

    valid_binaries = [b"\x00\x00\x00", b"\xff\xfe\x00"]
    invalid_binaries = [b"\x00\x00\x00", b"\xff\xff\x00"]

    valid_strings = ["\x00\x00\x00", "".join([chr(0x10FFFF), chr(0x10FFFE), chr(0x0)])]
    invalid_strings = ["\x00\x00\x00", "".join([chr(0x10FFFF), chr(0x10FFFF), chr(0x0)])]

    table = pa.Table.from_pydict(
        {
            "valid_upper_binary": valid_binaries,
            "invalid_upper_binary": invalid_binaries,
            "valid_upper_string": valid_strings,
            "invalid_upper_string": invalid_strings,
        },
        schema=arrow_schema,
    )

    metadata_collector: List[Any] = []

    with pa.BufferOutputStream() as f:
        with pq.ParquetWriter(f, table.schema, metadata_collector=metadata_collector) as writer:
            writer.write_table(table)

    return metadata_collector[0], table_metadata


def test_metrics_invalid_upper_bound() -> None:
    metadata, table_metadata = construct_test_table_invalid_upper_bound()

    schema = get_current_schema(table_metadata)
    table_metadata.properties["write.metadata.metrics.default"] = "truncate(2)"
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert len(datafile.value_counts) == 4
    assert len(datafile.null_value_counts) == 4
    assert len(datafile.nan_value_counts) == 0

    assert len(datafile.lower_bounds) == 4
    assert datafile.lower_bounds[1] == b"\x00\x00"
    assert datafile.lower_bounds[2] == b"\x00\x00"
    assert datafile.lower_bounds[3] == b"\x00\x00"
    assert datafile.lower_bounds[4] == b"\x00\x00"

    assert len(datafile.upper_bounds) == 2
    assert datafile.upper_bounds[1] == b"\xff\xff"
    assert datafile.upper_bounds[3] == "".join([chr(0x10FFFF), chr(0x10FFFF)]).encode()


def test_offsets() -> None:
    metadata, table_metadata = construct_test_table()

    schema = get_current_schema(table_metadata)
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )
    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    assert datafile.split_offsets is not None
    assert len(datafile.split_offsets) == 1
    assert datafile.split_offsets[0] == 4


def test_write_and_read_stats_schema(table_schema_nested: Schema) -> None:
    tbl = pa.Table.from_pydict(
        {
            "foo": ["a", "b"],
            "bar": [1, 2],
            "baz": [False, True],
            "qux": [["a", "b"], ["c", "d"]],
            "quux": [[("a", (("aa", 1), ("ab", 2)))], [("b", (("ba", 3), ("bb", 4)))]],
            "location": [[(52.377956, 4.897070), (4.897070, -122.431297)], [(43.618881, -116.215019), (41.881832, -87.623177)]],
            "person": [("Fokko", 33), ("Max", 42)],  # Possible data quality issue
        },
        schema=schema_to_pyarrow(table_schema_nested),
    )
    stats_columns = pre_order_visit(table_schema_nested, PyArrowStatisticsCollector(table_schema_nested, {}))

    visited_paths = []

    def file_visitor(written_file: Any) -> None:
        visited_paths.append(written_file)

    with tempfile.TemporaryDirectory() as tmpdir:
        pq.write_to_dataset(tbl, tmpdir, file_visitor=file_visitor)

    assert visited_paths[0].metadata.num_columns == len(stats_columns)


def test_stats_types(table_schema_nested: Schema) -> None:
    stats_columns = pre_order_visit(table_schema_nested, PyArrowStatisticsCollector(table_schema_nested, {}))

    # the field-ids should be sorted
    assert all(stats_columns[i].field_id <= stats_columns[i + 1].field_id for i in range(len(stats_columns) - 1))
    assert [col.iceberg_type for col in stats_columns] == [
        StringType(),
        IntegerType(),
        BooleanType(),
        StringType(),
        StringType(),
        StringType(),
        IntegerType(),
        FloatType(),
        FloatType(),
        StringType(),
        IntegerType(),
    ]


def test_read_missing_statistics() -> None:
    # write statistics for only for "strings" column
    metadata, table_metadata = construct_test_table(write_statistics=["strings"])

    # expect only "strings" column to have statistics in metadata
    # and all other columns to have no statistics
    for r in range(metadata.num_row_groups):
        for pos in range(metadata.num_columns):
            if metadata.row_group(r).column(pos).path_in_schema == "strings":
                assert metadata.row_group(r).column(pos).is_stats_set is True
                assert metadata.row_group(r).column(pos).statistics is not None
            else:
                assert metadata.row_group(r).column(pos).is_stats_set is False
                assert metadata.row_group(r).column(pos).statistics is None

    schema = get_current_schema(table_metadata)
    statistics = data_file_statistics_from_parquet_metadata(
        parquet_metadata=metadata,
        stats_columns=compute_statistics_plan(schema, table_metadata.properties),
        parquet_column_mapping=parquet_path_to_id_mapping(schema),
    )

    datafile = DataFile.from_args(**statistics.to_serialized_dict())

    # expect only "strings" column values to be reflected in the
    # upper_bound, lower_bound and null_value_counts props of datafile
    string_col_idx = 1
    assert len(datafile.lower_bounds) == 1
    assert datafile.lower_bounds[string_col_idx].decode() == "aaaaaaaaaaaaaaaa"
    assert len(datafile.upper_bounds) == 1
    assert datafile.upper_bounds[string_col_idx].decode() == "zzzzzzzzzzzzzzz{"
    assert len(datafile.null_value_counts) == 1
    assert datafile.null_value_counts[string_col_idx] == 1


# This is commented out for now because write_to_dataset drops the partition
# columns making it harder to calculate the mapping from the column index to
# datatype id
#
# def test_dataset() -> pa.Buffer:

#     table_metadata = {
#         "format-version": 2,
#         "location": "s3://bucket/test/location",
#         "last-column-id": 7,
#         "current-schema-id": 0,
#         "schemas": [
#             {
#                 "type": "struct",
#                 "schema-id": 0,
#                 "fields": [
#                     {"id": 1, "name": "ints", "required": False, "type": "long"},
#                     {"id": 2, "name": "even", "required": False, "type": "boolean"},
#                 ],
#             },
#         ],
#         "default-spec-id": 0,
#         "partition-specs": [{"spec-id": 0, "fields": []}],
#         "properties": {},
#     }

#     table_metadata = TableMetadataUtil.parse_obj(table_metadata)
#     schema = schema_to_pyarrow(table_metadata.schemas[0])

#     _ints = [0, 2, 4, 8, 1, 3, 5, 7]
#     parity = [True, True, True, True, False, False, False, False]

#     table = pa.Table.from_pydict({"ints": _ints, "even": parity}, schema=schema)

#     visited_paths = []

#     def file_visitor(written_file: Any) -> None:
#         visited_paths.append(written_file)

#     with TemporaryDirectory() as tmpdir:
#         pq.write_to_dataset(table, tmpdir, partition_cols=["even"], file_visitor=file_visitor)

#     even = None
#     odd = None

#     assert len(visited_paths) == 2

#     for written_file in visited_paths:
#         df = DataFile()

#         fill_parquet_file_metadata(df, written_file.metadata, written_file.size, table_metadata)

#         if "even=true" in written_file.path:
#             even = df

#         if "even=false" in written_file.path:
#             odd = df

#     assert even is not None
#     assert odd is not None

#     assert len(even.value_counts) == 1
#     assert even.value_counts[1] == 4
#     assert len(even.lower_bounds) == 1
#     assert even.lower_bounds[1] == STRUCT_INT64.pack(0)
#     assert len(even.upper_bounds) == 1
#     assert even.upper_bounds[1] == STRUCT_INT64.pack(8)

#     assert len(odd.value_counts) == 1
#     assert odd.value_counts[1] == 4
#     assert len(odd.lower_bounds) == 1
#     assert odd.lower_bounds[1] == STRUCT_INT64.pack(1)
#     assert len(odd.upper_bounds) == 1
#     assert odd.upper_bounds[1] == STRUCT_INT64.pack(7)