################################################################################
#  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.
################################################################################

import json
from typing import Callable, Dict, List, Optional

import pyarrow as pa

from pypaimon.common.predicate_json_parser import (
    _apply_predicate_transform,
    _collect_all_field_refs_from_transform,
)
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
from pypaimon.read.reader.iface.record_iterator import RecordIterator
from pypaimon.read.reader.iface.record_reader import RecordReader
from pypaimon.table.row.offset_row import OffsetRow


class RecordReaderToBatchAdapter(RecordBatchReader):

    def __init__(self, inner, schema: pa.Schema, chunk_size: int = 65536, include_row_kind: bool = False):
        self._inner = inner
        self._schema = schema
        self._chunk_size = chunk_size
        self._exhausted = False
        self._pending_iterator = None
        self._include_row_kind = include_row_kind
        self.blob_field_indices = getattr(inner, 'blob_field_indices', None)
        self.vector_field_indices = getattr(inner, 'vector_field_indices', None)

    def read_arrow_batch(self) -> Optional[pa.RecordBatch]:
        if self._exhausted:
            return None
        row_tuples = []
        row_kinds = []
        while len(row_tuples) < self._chunk_size:
            if self._pending_iterator is not None:
                row = self._pending_iterator.next()
                while row is not None:
                    row_tuples.append(
                        row.row_tuple[row.offset:row.offset + row.arity])
                    if self._include_row_kind:
                        row_kinds.append(row.get_row_kind().to_string())
                    if len(row_tuples) >= self._chunk_size:
                        return self._flush(row_tuples, row_kinds)
                    row = self._pending_iterator.next()
                self._pending_iterator = None

            row_iterator = self._inner.read_batch()
            if row_iterator is None:
                self._exhausted = True
                break
            self._pending_iterator = row_iterator

        if not row_tuples:
            return None
        return self._flush(row_tuples, row_kinds)

    def _flush(self, row_tuples, row_kinds=None):
        columns_data = list(zip(*row_tuples))
        pydict = {
            name: list(col)
            for name, col in zip(self._schema.names, columns_data)
        }
        batch = pa.RecordBatch.from_pydict(pydict, schema=self._schema)
        if row_kinds:
            row_kind_array = pa.array(row_kinds, type=pa.string())
            row_kind_field = pa.field("_row_kind", pa.string())
            new_schema = pa.schema([row_kind_field] + list(batch.schema))
            columns = [row_kind_array] + [batch.column(i) for i in range(batch.num_columns)]
            batch = pa.RecordBatch.from_arrays(columns, schema=new_schema)
        return batch

    def close(self):
        self._inner.close()


class BatchToRecordReaderAdapter(RecordReader):
    """Convert a batch-level RecordBatchReader back to a row-level RecordReader."""

    def __init__(self, inner: RecordBatchReader):
        self._inner = inner

    def read_batch(self):
        batch = self._inner.read_arrow_batch()
        if batch is None:
            return None
        return _ArrowBatchIterator(batch)

    def close(self):
        self._inner.close()


class _ArrowBatchIterator(RecordIterator):

    def __init__(self, batch: pa.RecordBatch):
        self._batch = batch
        self._idx = 0
        self._has_rk = "_row_kind" in batch.schema.names
        if self._has_rk:
            self._rk_idx = batch.schema.get_field_index("_row_kind")
            self._data_cols = [j for j in range(batch.num_columns) if j != self._rk_idx]
        else:
            self._rk_idx = -1
            self._data_cols = list(range(batch.num_columns))

    def next(self):
        if self._idx >= self._batch.num_rows:
            return None
        row_tuple = tuple(
            self._batch.column(j)[self._idx].as_py()
            for j in self._data_cols
        )
        row = OffsetRow(row_tuple, 0, len(self._data_cols))
        if self._has_rk:
            from pypaimon.table.row.row_kind import RowKind
            kind_str = self._batch.column(self._rk_idx)[self._idx].as_py()
            row.set_row_kind_byte(RowKind.from_string(kind_str).value)
        self._idx += 1
        return row


class AuthFilterReader(RecordBatchReader):

    def __init__(self, inner_reader: RecordBatchReader, filter_fn: Callable[[pa.RecordBatch], pa.Array]):
        self._inner = inner_reader
        self._filter_fn = filter_fn
        self.blob_field_indices = inner_reader.blob_field_indices
        self.vector_field_indices = inner_reader.vector_field_indices

    def read_arrow_batch(self) -> Optional[pa.RecordBatch]:
        batch = self._inner.read_arrow_batch()
        if batch is None:
            return None
        mask = self._filter_fn(batch)
        return batch.filter(mask)

    def close(self):
        self._inner.close()


class AuthMaskingReader(RecordBatchReader):

    def __init__(self, inner_reader: RecordBatchReader, masking_rules: Dict[str, str], read_fields: List):
        self._inner = inner_reader
        self._masking_rules = masking_rules
        self._read_fields = read_fields
        self.blob_field_indices = inner_reader.blob_field_indices
        self.vector_field_indices = inner_reader.vector_field_indices
        read_field_names = {f.name for f in read_fields}
        parsed = {}
        for col, tj in masking_rules.items():
            if col not in read_field_names:
                continue
            if not tj:
                continue
            transform = json.loads(tj)
            if transform is None:
                continue
            parsed[col] = transform
        self._parsed_rules = parsed
        for col_name, transform in self._parsed_rules.items():
            for ref_name in _collect_all_field_refs_from_transform(transform):
                if ref_name not in read_field_names:
                    raise RuntimeError(
                        f"Column masking refers to field '{ref_name}' which is not present "
                        f"in output row type. Available fields: {read_field_names}"
                    )

    def read_arrow_batch(self) -> Optional[pa.RecordBatch]:
        batch = self._inner.read_arrow_batch()
        if batch is None:
            return None
        original_batch = batch
        masked_columns = {}
        for col_name, transform in self._parsed_rules.items():
            if col_name in original_batch.schema.names:
                col_idx = original_batch.schema.get_field_index(col_name)
                target_col_type = original_batch.schema.field(col_idx).type
                masked_columns[col_idx] = self._apply_masking_transform(transform, original_batch, target_col_type)
        for col_idx, masked_array in masked_columns.items():
            original_field = original_batch.schema.field(col_idx)
            batch = batch.set_column(
                col_idx,
                pa.field(original_field.name, masked_array.type, nullable=True),
                masked_array)
        return batch

    def close(self):
        self._inner.close()

    def _apply_masking_transform(
            self,
            transform: dict,
            original_batch: pa.RecordBatch,
            target_col_type: pa.DataType,
    ) -> pa.Array:
        return _apply_predicate_transform(
            transform, original_batch, null_type=target_col_type)


class ColumnProjectReader(RecordBatchReader):

    def __init__(self, inner_reader: RecordBatchReader, columns: List[str]):
        self._inner = inner_reader
        self._columns = columns
        self.blob_field_indices = inner_reader.blob_field_indices
        self.vector_field_indices = inner_reader.vector_field_indices

    def read_arrow_batch(self) -> Optional[pa.RecordBatch]:
        batch = self._inner.read_arrow_batch()
        if batch is None:
            return None
        columns = self._columns
        if "_row_kind" in batch.schema.names and "_row_kind" not in columns:
            columns = ["_row_kind"] + list(columns)
        return batch.select(columns)

    def close(self):
        self._inner.close()
