# 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 typing import List, Optional, Any

import fastavro
import pyarrow as pa
import pyarrow.dataset as ds
from pyarrow import RecordBatch

from pypaimon.common.file_io import FileIO
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
from pypaimon.schema.data_types import DataField, PyarrowFieldParser


class FormatAvroReader(RecordBatchReader):
    """
    An ArrowBatchReader for reading Avro files using fastavro, filters records based on the
    provided predicate and projection, and converts Avro records to RecordBatch format.
    """

    def __init__(self, file_io: FileIO, file_path: str, read_fields: List[str], full_fields: List[DataField],
                 push_down_predicate: Any, batch_size: int = 1024,
                 nested_name_paths: Optional[List[List[str]]] = None):
        file_path_for_io = file_io.to_filesystem_path(file_path)
        self._file = file_io.filesystem.open_input_file(file_path_for_io)
        self._avro_reader = fastavro.reader(self._file)
        self._batch_size = batch_size
        self._push_down_predicate = push_down_predicate

        self._fields = read_fields
        full_fields_map = {field.name: field for field in full_fields}
        projected_data_fields = [full_fields_map[name] for name in read_fields]
        self._schema = PyarrowFieldParser.from_paimon_schema(projected_data_fields)

        if nested_name_paths is not None and len(nested_name_paths) != len(read_fields):
            raise ValueError(
                "nested_name_paths length {} does not match read_fields length {}".format(
                    len(nested_name_paths), len(read_fields)))
        self._nested_name_paths = nested_name_paths
        self._has_nested = bool(
            nested_name_paths and any(len(p) > 1 for p in nested_name_paths))

    def read_arrow_batch(self) -> Optional[RecordBatch]:
        pydict_data = {name: [] for name in self._fields}
        records_in_batch = 0

        for record in self._avro_reader:
            if self._has_nested:
                for col_name, path in zip(self._fields, self._nested_name_paths):
                    pydict_data[col_name].append(_walk_avro_record(record, path))
            else:
                for col_name in self._fields:
                    pydict_data[col_name].append(record.get(col_name))
            records_in_batch += 1
            if records_in_batch >= self._batch_size:
                break

        if records_in_batch == 0:
            return None
        if self._push_down_predicate is None:
            return pa.RecordBatch.from_pydict(pydict_data, self._schema)
        else:
            pa_batch = pa.Table.from_pydict(pydict_data, self._schema)
            dataset = ds.InMemoryDataset(pa_batch)
            scanner = dataset.scanner(filter=self._push_down_predicate)
            combine_chunks = scanner.to_table().combine_chunks()
            if combine_chunks.num_rows > 0:
                return combine_chunks.to_batches()[0]
            else:
                return None

    def close(self):
        if self._file:
            self._file.close()
            self._file = None


def _walk_avro_record(record, path: List[str]):
    """Walk field names through a nested avro record dict, returning the
    leaf value or ``None`` if any segment is missing.
    """
    current = record
    for name in path:
        if current is None:
            return None
        if isinstance(current, dict):
            current = current.get(name)
            continue
        return None
    return current
