# 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 collections.abc import Generator
from tempfile import TemporaryDirectory
from typing import (
    Any,
)

import pytest
from pydantic_core import to_json
from pyiceberg.partitioning import PartitionField, PartitionSpec
from pyiceberg.schema import Schema
from pyiceberg.transforms import DayTransform, IdentityTransform
from pyiceberg.types import (
    IntegerType,
    NestedField,
    TimestampType,
)
from pyiceberg.utils.config import Config


@pytest.fixture(scope="session", autouse=True)
def isolate_pyiceberg_config() -> Generator[None, None, None]:
    monkeypatch = pytest.MonkeyPatch()
    with TemporaryDirectory() as empty_home_dir:
        monkeypatch.setenv("HOME", empty_home_dir)
        monkeypatch.setenv("PYICEBERG_HOME", empty_home_dir)

        from pyiceberg import catalog

        monkeypatch.setattr(catalog, "_ENV_CONFIG", Config())
        yield
    monkeypatch.undo()


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


manifest_entry_records = [
    {
        "status": 1,
        "snapshot_id": 8744736658442914487,
        "data_file": {
            "file_path": "/home/iceberg/warehouse/nyc/taxis_partitioned/data/VendorID=null/00000-633-d8a4223e-dc97-45a1-86e1-adaba6e8abd7-00001.parquet",
            "file_format": "PARQUET",
            "partition": {"VendorID": 1, "tpep_pickup_datetime": 1925},
            "record_count": 19513,
            "file_size_in_bytes": 388872,
            "block_size_in_bytes": 67108864,
            "column_sizes": [
                {"key": 1, "value": 53},
                {"key": 2, "value": 98153},
            ],
            "value_counts": [
                {"key": 1, "value": 19513},
                {"key": 2, "value": 19513},
            ],
            "null_value_counts": [
                {"key": 1, "value": 19513},
                {"key": 2, "value": 0},
            ],
            "nan_value_counts": [],
            "lower_bounds": [
                {"key": 1, "value": b"\x00\x00\x00\x00"},
                {"key": 2, "value": b"\x00\x00\x00\x00\x00\x00\x00\x00"},
            ],
            "upper_bounds": [
                {"key": 1, "value": b"\x00\x00\xe0\xbf"},
                {"key": 2, "value": b"\x00\x00\x00\x00\x00\x00\xe0\xbf"},
            ],
            "key_metadata": None,
            "split_offsets": [4],
            "sort_order_id": 0,
        },
    },
    {
        "status": 1,
        "snapshot_id": 8744736658442914487,
        "data_file": {
            "file_path": "/home/iceberg/warehouse/nyc/taxis_partitioned/data/VendorID=1/00000-633-d8a4223e-dc97-45a1-86e1-adaba6e8abd7-00002.parquet",
            "file_format": "PARQUET",
            "partition": {"VendorID": 1, "tpep_pickup_datetime": None},
            "record_count": 95050,
            "file_size_in_bytes": 1265950,
            "block_size_in_bytes": 67108864,
            "column_sizes": [
                {"key": 1, "value": 318},
                {"key": 2, "value": 329806},
            ],
            "value_counts": [
                {"key": 1, "value": 95050},
                {"key": 2, "value": 95050},
            ],
            "null_value_counts": [
                {"key": 1, "value": 0},
                {"key": 2, "value": 0},
            ],
            "nan_value_counts": [],
            "lower_bounds": [
                {"key": 1, "value": b"\x00\x00\x00\x00"},
                {"key": 2, "value": b"\x00\x00\x00\x00\x00\x00\x00\x00"},
            ],
            "upper_bounds": [
                {"key": 1, "value": b"\x00\x00\xe0\xbf"},
                {"key": 2, "value": b"\x00\x00\x00\x00\x00\x00\xe0\xbf"},
            ],
            "key_metadata": None,
            "split_offsets": [4],
            "sort_order_id": 0,
        },
    },
]


@pytest.fixture(scope="session")
def test_schema() -> Schema:
    return Schema(
        NestedField(1, "VendorID", IntegerType(), False),
        NestedField(2, "tpep_pickup_datetime", TimestampType(), False),
    )


@pytest.fixture(scope="session")
def test_partition_spec() -> Schema:
    return PartitionSpec(
        PartitionField(1, 1000, IdentityTransform(), "VendorID"),
        PartitionField(2, 1001, DayTransform(), "tpep_pickup_day"),
    )


@pytest.fixture(scope="session")
def generated_manifest_entry_file(
    avro_schema_manifest_entry: dict[str, Any],
    test_schema: Schema,
    test_partition_spec: PartitionSpec,
) -> Generator[str, None, None]:
    from fastavro import parse_schema, writer

    parsed_schema = parse_schema(avro_schema_manifest_entry)

    with TemporaryDirectory() as tmpdir:
        tmp_avro_file = tmpdir + "/manifest.avro"
        with open(tmp_avro_file, "wb") as out:
            writer(
                out,
                parsed_schema,
                manifest_entry_records,
                metadata={
                    "schema": test_schema.model_dump_json(),
                    "partition-spec": to_json(test_partition_spec.fields).decode(
                        "utf-8"
                    ),
                },
            )
        yield tmp_avro_file