# 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
import json

import pytest

from pyiceberg.avro.file import AvroFile
from pyiceberg.avro.reader import (
    AvroStruct,
    BinaryReader,
    BooleanReader,
    DateReader,
    DecimalReader,
    DoubleReader,
    FixedReader,
    FloatReader,
    IntegerReader,
    StringReader,
    TimeReader,
    TimestampReader,
    TimestamptzReader,
    UUIDReader,
    primitive_reader,
)
from pyiceberg.manifest import _convert_pos_to_dict
from pyiceberg.schema import Schema
from pyiceberg.types import (
    BinaryType,
    BooleanType,
    DateType,
    DecimalType,
    DoubleType,
    FixedType,
    FloatType,
    IntegerType,
    ListType,
    LongType,
    MapType,
    NestedField,
    PrimitiveType,
    StringType,
    StructType,
    TimestampType,
    TimestamptzType,
    TimeType,
    UUIDType,
)
from tests.io.test_io import LocalInputFile


def test_read_header(generated_manifest_entry_file: str, iceberg_manifest_entry_schema: Schema):
    with AvroFile(LocalInputFile(generated_manifest_entry_file)) as reader:
        header = reader._read_header()

    assert header.magic == b"Obj\x01"
    assert json.loads(header.meta["avro.schema"]) == {
        "type": "record",
        "name": "manifest_entry",
        "fields": [
            {"field-id": 0, "name": "status", "type": "int"},
            {"field-id": 1, "default": None, "name": "snapshot_id", "type": ["null", "long"]},
            {
                "field-id": 2,
                "name": "data_file",
                "type": {
                    "type": "record",
                    "name": "r2",
                    "fields": [
                        {"field-id": 100, "doc": "Location URI with FS scheme", "name": "file_path", "type": "string"},
                        {
                            "field-id": 101,
                            "doc": "File format name: avro, orc, or parquet",
                            "name": "file_format",
                            "type": "string",
                        },
                        {
                            "field-id": 102,
                            "name": "partition",
                            "type": {
                                "type": "record",
                                "name": "r102",
                                "fields": [{"field-id": 1000, "default": None, "name": "VendorID", "type": ["null", "int"]}],
                            },
                        },
                        {"field-id": 103, "doc": "Number of records in the file", "name": "record_count", "type": "long"},
                        {"field-id": 104, "doc": "Total file size in bytes", "name": "file_size_in_bytes", "type": "long"},
                        {"field-id": 105, "name": "block_size_in_bytes", "type": "long"},
                        {
                            "field-id": 108,
                            "doc": "Map of column id to total size on disk",
                            "default": None,
                            "name": "column_sizes",
                            "type": [
                                "null",
                                {
                                    "logicalType": "map",
                                    "type": "array",
                                    "items": {
                                        "type": "record",
                                        "name": "k117_v118",
                                        "fields": [
                                            {"field-id": 117, "name": "key", "type": "int"},
                                            {"field-id": 118, "name": "value", "type": "long"},
                                        ],
                                    },
                                },
                            ],
                        },
                        {
                            "field-id": 109,
                            "doc": "Map of column id to total count, including null and NaN",
                            "default": None,
                            "name": "value_counts",
                            "type": [
                                "null",
                                {
                                    "logicalType": "map",
                                    "type": "array",
                                    "items": {
                                        "type": "record",
                                        "name": "k119_v120",
                                        "fields": [
                                            {"field-id": 119, "name": "key", "type": "int"},
                                            {"field-id": 120, "name": "value", "type": "long"},
                                        ],
                                    },
                                },
                            ],
                        },
                        {
                            "field-id": 110,
                            "doc": "Map of column id to null value count",
                            "default": None,
                            "name": "null_value_counts",
                            "type": [
                                "null",
                                {
                                    "logicalType": "map",
                                    "type": "array",
                                    "items": {
                                        "type": "record",
                                        "name": "k121_v122",
                                        "fields": [
                                            {"field-id": 121, "name": "key", "type": "int"},
                                            {"field-id": 122, "name": "value", "type": "long"},
                                        ],
                                    },
                                },
                            ],
                        },
                        {
                            "field-id": 137,
                            "doc": "Map of column id to number of NaN values in the column",
                            "default": None,
                            "name": "nan_value_counts",
                            "type": [
                                "null",
                                {
                                    "logicalType": "map",
                                    "type": "array",
                                    "items": {
                                        "type": "record",
                                        "name": "k138_v139",
                                        "fields": [
                                            {"field-id": 138, "name": "key", "type": "int"},
                                            {"field-id": 139, "name": "value", "type": "long"},
                                        ],
                                    },
                                },
                            ],
                        },
                        {
                            "field-id": 125,
                            "doc": "Map of column id to lower bound",
                            "default": None,
                            "name": "lower_bounds",
                            "type": [
                                "null",
                                {
                                    "logicalType": "map",
                                    "type": "array",
                                    "items": {
                                        "type": "record",
                                        "name": "k126_v127",
                                        "fields": [
                                            {"field-id": 126, "name": "key", "type": "int"},
                                            {"field-id": 127, "name": "value", "type": "bytes"},
                                        ],
                                    },
                                },
                            ],
                        },
                        {
                            "field-id": 128,
                            "doc": "Map of column id to upper bound",
                            "default": None,
                            "name": "upper_bounds",
                            "type": [
                                "null",
                                {
                                    "logicalType": "map",
                                    "type": "array",
                                    "items": {
                                        "type": "record",
                                        "name": "k129_v130",
                                        "fields": [
                                            {"field-id": 129, "name": "key", "type": "int"},
                                            {"field-id": 130, "name": "value", "type": "bytes"},
                                        ],
                                    },
                                },
                            ],
                        },
                        {
                            "field-id": 131,
                            "doc": "Encryption key metadata blob",
                            "default": None,
                            "name": "key_metadata",
                            "type": ["null", "bytes"],
                        },
                        {
                            "field-id": 132,
                            "doc": "Splittable offsets",
                            "default": None,
                            "name": "split_offsets",
                            "type": ["null", {"element-id": 133, "type": "array", "items": "long"}],
                        },
                        {
                            "field-id": 140,
                            "doc": "Sort order ID",
                            "default": None,
                            "name": "sort_order_id",
                            "type": ["null", "int"],
                        },
                    ],
                },
            },
        ],
    }

    assert header.get_schema() == iceberg_manifest_entry_schema


def test_read_manifest_entry_file(generated_manifest_entry_file: str):
    with AvroFile(LocalInputFile(generated_manifest_entry_file)) as reader:
        # Consume the generator
        records = list(reader)

    assert len(records) == 2, f"Expected 2 records, got {len(records)}"
    assert records[0] == AvroStruct(
        _data=[
            1,
            8744736658442914487,
            AvroStruct(
                _data=[
                    "/home/iceberg/warehouse/nyc/taxis_partitioned/data/VendorID=null/00000-633-d8a4223e-dc97-45a1-86e1-adaba6e8abd7-00001.parquet",
                    "PARQUET",
                    AvroStruct(_data=[None]),
                    19513,
                    388872,
                    67108864,
                    {
                        1: 53,
                        2: 98153,
                        3: 98693,
                        4: 53,
                        5: 53,
                        6: 53,
                        7: 17425,
                        8: 18528,
                        9: 53,
                        10: 44788,
                        11: 35571,
                        12: 53,
                        13: 1243,
                        14: 2355,
                        15: 12750,
                        16: 4029,
                        17: 110,
                        18: 47194,
                        19: 2948,
                    },
                    {
                        1: 19513,
                        2: 19513,
                        3: 19513,
                        4: 19513,
                        5: 19513,
                        6: 19513,
                        7: 19513,
                        8: 19513,
                        9: 19513,
                        10: 19513,
                        11: 19513,
                        12: 19513,
                        13: 19513,
                        14: 19513,
                        15: 19513,
                        16: 19513,
                        17: 19513,
                        18: 19513,
                        19: 19513,
                    },
                    {
                        1: 19513,
                        2: 0,
                        3: 0,
                        4: 19513,
                        5: 19513,
                        6: 19513,
                        7: 0,
                        8: 0,
                        9: 19513,
                        10: 0,
                        11: 0,
                        12: 19513,
                        13: 0,
                        14: 0,
                        15: 0,
                        16: 0,
                        17: 0,
                        18: 0,
                        19: 0,
                    },
                    {16: 0, 17: 0, 18: 0, 19: 0, 10: 0, 11: 0, 12: 0, 13: 0, 14: 0, 15: 0},
                    {
                        2: b"2020-04-01 00:00",
                        3: b"2020-04-01 00:12",
                        7: b"\x03\x00\x00\x00",
                        8: b"\x01\x00\x00\x00",
                        10: b"\xf6(\\\x8f\xc2\x05S\xc0",
                        11: b"\x00\x00\x00\x00\x00\x00\x00\x00",
                        13: b"\x00\x00\x00\x00\x00\x00\x00\x00",
                        14: b"\x00\x00\x00\x00\x00\x00\xe0\xbf",
                        15: b")\\\x8f\xc2\xf5(\x08\xc0",
                        16: b"\x00\x00\x00\x00\x00\x00\x00\x00",
                        17: b"\x00\x00\x00\x00\x00\x00\x00\x00",
                        18: b"\xf6(\\\x8f\xc2\xc5S\xc0",
                        19: b"\x00\x00\x00\x00\x00\x00\x04\xc0",
                    },
                    {
                        2: b"2020-04-30 23:5:",
                        3: b"2020-05-01 00:41",
                        7: b"\t\x01\x00\x00",
                        8: b"\t\x01\x00\x00",
                        10: b"\xcd\xcc\xcc\xcc\xcc,_@",
                        11: b"\x1f\x85\xebQ\\\xe2\xfe@",
                        13: b"\x00\x00\x00\x00\x00\x00\x12@",
                        14: b"\x00\x00\x00\x00\x00\x00\xe0?",
                        15: b"q=\n\xd7\xa3\xf01@",
                        16: b"\x00\x00\x00\x00\x00`B@",
                        17: b"333333\xd3?",
                        18: b"\x00\x00\x00\x00\x00\x18b@",
                        19: b"\x00\x00\x00\x00\x00\x00\x04@",
                    },
                    None,
                    [4],
                    0,
                ]
            ),
        ]
    )


def test_read_manifest_file_file(generated_manifest_file_file: str):
    with AvroFile(LocalInputFile(generated_manifest_file_file)) as reader:
        # Consume the generator
        records = list(reader)

    assert len(records) == 1, f"Expected 1 records, got {len(records)}"
    actual = records[0]
    expected = AvroStruct(
        _data=[
            actual.get(0),
            7989,
            0,
            9182715666859759686,
            3,
            0,
            0,
            [AvroStruct(_data=[True, False, b"\x01\x00\x00\x00", b"\x02\x00\x00\x00"])],
            237993,
            0,
            0,
        ]
    )
    assert actual == expected


def test_null_list_convert_pos_to_dict():
    data = _convert_pos_to_dict(
        Schema(
            NestedField(name="field", field_id=1, field_type=ListType(element_id=2, element=StringType(), element_required=False))
        ),
        AvroStruct([None]),
    )
    assert data["field"] is None


def test_null_dict_convert_pos_to_dict():
    data = _convert_pos_to_dict(
        Schema(
            NestedField(
                name="field",
                field_id=1,
                field_type=MapType(key_id=2, key_type=StringType(), value_id=3, value_type=StringType(), value_required=False),
            )
        ),
        AvroStruct([None]),
    )
    assert data["field"] is None


def test_null_struct_convert_pos_to_dict():
    data = _convert_pos_to_dict(
        Schema(
            NestedField(
                name="field",
                field_id=1,
                field_type=StructType(
                    NestedField(2, "required_field", StringType(), True), NestedField(3, "optional_field", IntegerType())
                ),
                required=False,
            )
        ),
        AvroStruct([None]),
    )
    assert data["field"] is None


def test_fixed_reader():
    assert primitive_reader(FixedType(22)) == FixedReader(22)


def test_decimal_reader():
    assert primitive_reader(DecimalType(19, 25)) == DecimalReader(19, 25)


def test_boolean_reader():
    assert primitive_reader(BooleanType()) == BooleanReader()


def test_integer_reader():
    assert primitive_reader(IntegerType()) == IntegerReader()


def test_long_reader():
    assert primitive_reader(LongType()) == IntegerReader()


def test_float_reader():
    assert primitive_reader(FloatType()) == FloatReader()


def test_double_reader():
    assert primitive_reader(DoubleType()) == DoubleReader()


def test_date_reader():
    assert primitive_reader(DateType()) == DateReader()


def test_time_reader():
    assert primitive_reader(TimeType()) == TimeReader()


def test_timestamp_reader():
    assert primitive_reader(TimestampType()) == TimestampReader()


def test_timestamptz_reader():
    assert primitive_reader(TimestamptzType()) == TimestamptzReader()


def test_string_reader():
    assert primitive_reader(StringType()) == StringReader()


def test_binary_reader():
    assert primitive_reader(BinaryType()) == BinaryReader()


def test_unknown_type():
    class UnknownType(PrimitiveType):
        __root__ = "UnknownType"

    with pytest.raises(ValueError) as exc_info:
        primitive_reader(UnknownType())

    assert "Unknown type:" in str(exc_info.value)


def test_uuid_reader() -> None:
    assert primitive_reader(UUIDType()) == UUIDReader()