# 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
from typing import Callable

import pytest

from pyiceberg.avro.decoder import BinaryDecoder, StreamingBinaryDecoder
from pyiceberg.avro.decoder_fast import CythonBinaryDecoder
from pyiceberg.avro.file import AvroFile
from pyiceberg.avro.reader import (
    BinaryReader,
    BooleanReader,
    DateReader,
    DecimalReader,
    DoubleReader,
    FixedReader,
    FloatReader,
    IntegerReader,
    StringReader,
    StructReader,
    TimeReader,
    TimestampNanoReader,
    TimestampReader,
    TimestamptzNanoReader,
    TimestamptzReader,
    UnknownReader,
    UUIDReader,
)
from pyiceberg.avro.resolver import construct_reader
from pyiceberg.io.pyarrow import PyArrowFileIO
from pyiceberg.manifest import MANIFEST_ENTRY_SCHEMAS, DataFile, ManifestEntry
from pyiceberg.schema import Schema
from pyiceberg.typedef import Record
from pyiceberg.types import (
    BinaryType,
    BooleanType,
    DateType,
    DecimalType,
    DoubleType,
    FixedType,
    FloatType,
    IntegerType,
    LongType,
    NestedField,
    StringType,
    StructType,
    TimestampNanoType,
    TimestampType,
    TimestamptzNanoType,
    TimestamptzType,
    TimeType,
    UnknownType,
    UUIDType,
)

AVAILABLE_DECODERS = [StreamingBinaryDecoder, CythonBinaryDecoder]


def test_read_header(generated_manifest_entry_file: str, iceberg_manifest_entry_schema: Schema) -> None:
    with AvroFile[ManifestEntry](
        PyArrowFileIO().new_input(generated_manifest_entry_file),
        MANIFEST_ENTRY_SCHEMAS[2],
        {-1: ManifestEntry, 2: DataFile},
    ) as reader:
        header = reader.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": 1001,
                                        "default": None,
                                        "name": "tpep_pickup_datetime",
                                        "type": ["null", {"type": "int", "logicalType": "date"}],
                                    },
                                ],
                            },
                        },
                        {"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_fixed_reader() -> None:
    assert construct_reader(FixedType(22)) == FixedReader(22)


def test_decimal_reader() -> None:
    assert construct_reader(DecimalType(19, 25)) == DecimalReader(19, 25)


def test_boolean_reader() -> None:
    assert construct_reader(BooleanType()) == BooleanReader()


def test_integer_reader() -> None:
    assert construct_reader(IntegerType()) == IntegerReader()


def test_long_reader() -> None:
    assert construct_reader(LongType()) == IntegerReader()


def test_float_reader() -> None:
    assert construct_reader(FloatType()) == FloatReader()


def test_double_reader() -> None:
    assert construct_reader(DoubleType()) == DoubleReader()


def test_date_reader() -> None:
    assert construct_reader(DateType()) == DateReader()


def test_time_reader() -> None:
    assert construct_reader(TimeType()) == TimeReader()


def test_timestamp_reader() -> None:
    assert construct_reader(TimestampType()) == TimestampReader()


def test_timestamp_ns_reader() -> None:
    assert construct_reader(TimestampNanoType()) == TimestampNanoReader()


def test_timestamptz_reader() -> None:
    assert construct_reader(TimestamptzType()) == TimestamptzReader()


def test_timestamptz_ns_reader() -> None:
    assert construct_reader(TimestamptzNanoType()) == TimestamptzNanoReader()


def test_string_reader() -> None:
    assert construct_reader(StringType()) == StringReader()


def test_binary_reader() -> None:
    assert construct_reader(BinaryType()) == BinaryReader()


def test_unknown_type() -> None:
    assert construct_reader(UnknownType()) == UnknownReader()


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


@pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
def test_read_struct(decoder_class: Callable[[bytes], BinaryDecoder]) -> None:
    decoder = decoder_class(b"\x18")
    struct = StructType(NestedField(1, "id", IntegerType(), required=True))
    result = StructReader(((0, IntegerReader()),), Record, struct).read(decoder)
    assert repr(result) == "Record[12]"


@pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
def test_read_struct_lambda(decoder_class: Callable[[bytes], BinaryDecoder]) -> None:
    decoder = decoder_class(b"\x18")

    struct = StructType(NestedField(1, "id", IntegerType(), required=True))
    # You can also pass in an arbitrary function that returns a struct
    result = StructReader(
        ((0, IntegerReader()),),
        Record,
        struct,  # pylint: disable=unnecessary-lambda
    ).read(decoder)
    assert repr(result) == "Record[12]"


@pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
def test_read_not_struct_type(decoder_class: Callable[[bytes], BinaryDecoder]) -> None:
    decoder = decoder_class(b"\x18")

    struct = StructType(NestedField(1, "id", IntegerType(), required=True))
    with pytest.raises(ValueError) as exc_info:
        _ = StructReader(((0, IntegerReader()),), str, struct).read(decoder)  # type: ignore

    assert "Incompatible with StructProtocol: <class 'str'>" in str(exc_info.value)