# 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 struct
from decimal import Decimal
from typing import Any, List, Optional, Tuple

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

from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor
from pypaimon.common.file_io import FileIO
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
from pypaimon.schema.data_types import (
    ArrayType, DataField, MapType, MultisetType, PyarrowFieldParser, RowType, VectorType, AtomicType
)
from pypaimon.table.special_fields import SpecialFields

FOOTER_SIZE = 32
MAGIC = 0x524F5753  # "ROWS"
VERSION = 1


def _special_field(name: str):
    if name == SpecialFields.ROW_ID.name:
        return SpecialFields.ROW_ID
    if name == SpecialFields.SEQUENCE_NUMBER.name:
        return SpecialFields.SEQUENCE_NUMBER
    if name == SpecialFields.VALUE_KIND.name:
        return SpecialFields.VALUE_KIND
    return None


class FormatRowReader(RecordBatchReader):

    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, row_indices: Optional[List[int]] = None):
        self._file_io = file_io
        self._file_path = file_path
        self._push_down_predicate = push_down_predicate
        self._batch_size = batch_size

        self._file_size = file_io.get_file_size(file_path)

        full_fields_map = {field.name: field for field in full_fields}
        self._read_field_names = list(read_fields)
        self._projected_fields = [full_fields_map[name] for name in read_fields
                                  if name in full_fields_map]
        self._missing_fields = [name for name in read_fields
                                if name not in full_fields_map]
        self._all_fields = full_fields
        self._physical_schema = PyarrowFieldParser.from_paimon_schema(self._projected_fields)

        self._block_compressed_sizes: List[int] = []
        self._block_uncompressed_sizes: List[int] = []
        self._block_row_starts: List[int] = []
        self._block_offsets: List[int] = []
        self._total_row_count = 0
        self._block_count = 0
        self._current_block_idx = 0

        self._row_indices = row_indices
        self._row_indices_pos = 0

        self._read_metadata()

        if self._row_indices is not None:
            self._blocks_to_read = self._compute_blocks_for_indices()
            self._blocks_to_read_pos = 0
        else:
            self._blocks_to_read = None

    def _compute_blocks_for_indices(self) -> List[tuple]:
        """Group row_indices by block. Returns list of (block_idx, [local_row_offsets])."""
        import bisect
        result = []
        row_starts = self._block_row_starts
        for idx in self._row_indices:
            block_idx = bisect.bisect_right(row_starts, idx) - 1
            local_row = idx - row_starts[block_idx]
            if result and result[-1][0] == block_idx:
                result[-1][1].append(local_row)
            else:
                result.append((block_idx, [local_row]))
        return result

    def read_arrow_batch(self) -> Optional[RecordBatch]:
        if self._row_indices is not None:
            return self._read_batch_indexed()
        return self._read_batch_sequential()

    def _read_batch_sequential(self) -> Optional[RecordBatch]:
        if self._current_block_idx >= self._block_count:
            return None

        block_data = self._read_and_decompress_block(self._current_block_idx)
        self._current_block_idx += 1

        columns, row_count = self._decode_block(block_data)

        if row_count == 0:
            return None

        table = self._build_table(columns, row_count)

        if self._push_down_predicate is not None:
            dataset = ds.InMemoryDataset(table)
            scanner = dataset.scanner(filter=self._push_down_predicate)
            table = scanner.to_table().combine_chunks()

        if table.num_rows == 0:
            return self._read_batch_sequential()

        return table.to_batches()[0]

    def _read_batch_indexed(self) -> Optional[RecordBatch]:
        if self._blocks_to_read_pos >= len(self._blocks_to_read):
            return None

        block_idx, local_rows = self._blocks_to_read[self._blocks_to_read_pos]
        self._blocks_to_read_pos += 1

        block_data = self._read_and_decompress_block(block_idx)
        columns, row_count = self._decode_block(block_data, row_filter=local_rows)

        if row_count == 0:
            return self._read_batch_indexed()

        table = self._build_table(columns, row_count)

        if self._push_down_predicate is not None:
            dataset = ds.InMemoryDataset(table)
            scanner = dataset.scanner(filter=self._push_down_predicate)
            table = scanner.to_table().combine_chunks()

        if table.num_rows == 0:
            return self._read_batch_indexed()

        return table.to_batches()[0]

    def close(self):
        pass

    def _read_metadata(self):
        with self._file_io.new_input_stream(self._file_path) as f:
            f.seek(self._file_size - FOOTER_SIZE)
            footer_bytes = f.read(FOOTER_SIZE)

        if len(footer_bytes) != FOOTER_SIZE:
            raise IOError("Invalid row file: cannot read footer")

        magic = struct.unpack_from('<I', footer_bytes, 28)[0]
        if magic != MAGIC:
            raise IOError(f"Invalid row file magic: expected 0x{MAGIC:08X}, got 0x{magic:08X}")

        version = footer_bytes[24]
        if version != VERSION:
            raise IOError(f"Unsupported row file version: {version}")

        self._total_row_count = struct.unpack_from('<q', footer_bytes, 0)[0]
        self._block_count = struct.unpack_from('<i', footer_bytes, 8)[0]
        index_offset = struct.unpack_from('<q', footer_bytes, 12)[0]
        index_length = struct.unpack_from('<i', footer_bytes, 20)[0]

        with self._file_io.new_input_stream(self._file_path) as f:
            f.seek(index_offset)
            index_bytes = f.read(index_length)

        if len(index_bytes) != index_length:
            raise IOError("Invalid row file: cannot read block index")

        self._parse_block_index(index_bytes)

    def _parse_block_index(self, index_data: bytes):
        pos = 0

        len1, consumed = _decode_var_int(index_data, pos)
        pos += consumed
        self._block_compressed_sizes = DeltaVarintCompressor.decompress(index_data[pos:pos + len1])
        pos += len1

        len2, consumed = _decode_var_int(index_data, pos)
        pos += consumed
        self._block_uncompressed_sizes = DeltaVarintCompressor.decompress(index_data[pos:pos + len2])
        pos += len2

        len3, consumed = _decode_var_int(index_data, pos)
        pos += consumed
        self._block_row_starts = DeltaVarintCompressor.decompress(index_data[pos:pos + len3])

        offset = 0
        self._block_offsets = []
        for size in self._block_compressed_sizes:
            self._block_offsets.append(offset)
            offset += size

    def _read_and_decompress_block(self, block_idx: int) -> bytes:
        import zstandard as zstd

        offset = self._block_offsets[block_idx]
        compressed_size = self._block_compressed_sizes[block_idx]

        with self._file_io.new_input_stream(self._file_path) as f:
            f.seek(offset)
            compressed_data = f.read(compressed_size)

        decompressor = zstd.ZstdDecompressor()
        uncompressed_size = self._block_uncompressed_sizes[block_idx]
        return decompressor.decompress(compressed_data, max_output_size=uncompressed_size)

    def _build_table(self, columns: List[List], row_count: int) -> pa.Table:
        pydict = {
            field.name: columns[i]
            for i, field in enumerate(self._projected_fields)
        }
        if not self._missing_fields:
            return pa.Table.from_pydict(pydict, self._physical_schema)

        arrays = []
        fields = []
        projected_by_name = {field.name: field for field in self._projected_fields}
        for name in self._read_field_names:
            if name in projected_by_name:
                field = projected_by_name[name]
                pa_type = PyarrowFieldParser.from_paimon_type(field.type)
                arrays.append(pa.array(pydict[name], type=pa_type))
                fields.append(pa.field(name, pa_type, nullable=field.type.nullable))
            else:
                special_field = _special_field(name)
                pa_type = (
                    PyarrowFieldParser.from_paimon_type(special_field.type)
                    if special_field is not None else pa.null()
                )
                arrays.append(pa.nulls(row_count, type=pa_type))
                fields.append(pa.field(name, pa_type, nullable=not SpecialFields.is_system_field(name)))
        return pa.Table.from_arrays(arrays, schema=pa.schema(fields))

    def _decode_block(self, block_data: bytes,
                      row_filter: Optional[List[int]] = None) -> Tuple[List[List], int]:
        data_len = len(block_data)
        row_count = struct.unpack_from('<i', block_data, data_len - 4)[0]
        offset_array_start = data_len - 4 - row_count * 4

        projected_indices = []
        for pf in self._projected_fields:
            for i, af in enumerate(self._all_fields):
                if af.name == pf.name:
                    projected_indices.append(i)
                    break

        columns = [[] for _ in self._projected_fields]
        arity = len(self._all_fields)
        header_size = (arity + 7) // 8

        rows_to_read = row_filter if row_filter is not None else range(row_count)
        rows_to_read_count = len(rows_to_read)
        proj_set = set(projected_indices)
        proj_col_map = {field_idx: col_idx for col_idx, field_idx in enumerate(projected_indices)}

        for row_idx in rows_to_read:
            row_offset = struct.unpack_from('<i', block_data, offset_array_start + row_idx * 4)[0]
            decoder = _RowDecoder(block_data, row_offset)

            null_bitmap = block_data[decoder.pos:decoder.pos + header_size]
            decoder.pos += header_size

            for field_idx in range(arity):
                is_null = (null_bitmap[field_idx // 8] & (1 << (field_idx % 8))) != 0
                if is_null:
                    if field_idx in proj_set:
                        columns[proj_col_map[field_idx]].append(None)
                else:
                    value = _read_field(decoder, self._all_fields[field_idx].type)
                    if field_idx in proj_set:
                        columns[proj_col_map[field_idx]].append(value)

        return columns, rows_to_read_count


class _RowDecoder:

    __slots__ = ('data', 'pos')

    def __init__(self, data: bytes, pos: int = 0):
        self.data = data
        self.pos = pos

    def read_boolean(self) -> bool:
        v = self.data[self.pos] != 0
        self.pos += 1
        return v

    def read_byte(self) -> int:
        v = struct.unpack_from('<b', self.data, self.pos)[0]
        self.pos += 1
        return v

    def read_short(self) -> int:
        v = struct.unpack_from('<h', self.data, self.pos)[0]
        self.pos += 2
        return v

    def read_int(self) -> int:
        v = struct.unpack_from('<i', self.data, self.pos)[0]
        self.pos += 4
        return v

    def read_long(self) -> int:
        v = struct.unpack_from('<q', self.data, self.pos)[0]
        self.pos += 8
        return v

    def read_float(self) -> float:
        v = struct.unpack_from('<f', self.data, self.pos)[0]
        self.pos += 4
        return v

    def read_double(self) -> float:
        v = struct.unpack_from('<d', self.data, self.pos)[0]
        self.pos += 8
        return v

    def read_var_int(self) -> int:
        result = 0
        shift = 0
        while True:
            b = self.data[self.pos]
            self.pos += 1
            result |= (b & 0x7F) << shift
            if (b & 0x80) == 0:
                return result
            shift += 7

    def read_string(self) -> str:
        length = self.read_var_int()
        s = self.data[self.pos:self.pos + length].decode('utf-8')
        self.pos += length
        return s

    def read_bytes(self) -> bytes:
        length = self.read_var_int()
        b = self.data[self.pos:self.pos + length]
        self.pos += length
        return bytes(b)


def _decode_var_int(data: bytes, offset: int) -> tuple:
    result = 0
    shift = 0
    pos = offset
    while True:
        b = data[pos]
        pos += 1
        result |= (b & 0x7F) << shift
        if (b & 0x80) == 0:
            return result, pos - offset
        shift += 7


def _read_field(decoder: _RowDecoder, data_type) -> Any:
    if isinstance(data_type, AtomicType):
        type_name = data_type.type.upper()
        if type_name == 'BOOLEAN':
            return decoder.read_boolean()
        elif type_name == 'TINYINT':
            return decoder.read_byte()
        elif type_name == 'SMALLINT':
            return decoder.read_short()
        elif type_name in ('INT', 'INTEGER', 'DATE', 'TIME'):
            return decoder.read_int()
        elif type_name.startswith('TIME') and not type_name.startswith('TIMESTAMP'):
            return decoder.read_int()
        elif type_name == 'BIGINT':
            return decoder.read_long()
        elif type_name == 'FLOAT':
            return decoder.read_float()
        elif type_name == 'DOUBLE':
            return decoder.read_double()
        elif type_name == 'STRING' or type_name.startswith('CHAR') or type_name.startswith('VARCHAR'):
            return decoder.read_string()
        elif type_name == 'BYTES' or type_name.startswith('BINARY') or type_name.startswith('VARBINARY'):
            return decoder.read_bytes()
        elif type_name == 'BLOB':
            return decoder.read_bytes()
        elif type_name.startswith('DECIMAL'):
            precision, scale = _parse_decimal_params(type_name)
            if precision <= 18:
                unscaled = decoder.read_long()
                return Decimal(unscaled) / Decimal(10 ** scale)
            else:
                raw = decoder.read_bytes()
                unscaled = int.from_bytes(raw, byteorder='big', signed=True)
                return Decimal(unscaled) / Decimal(10 ** scale)
        elif type_name.startswith('TIMESTAMP'):
            precision = _parse_timestamp_precision(type_name)
            millis = decoder.read_long()
            if precision <= 3:
                return millis
            else:
                nano_of_milli = decoder.read_var_int()
                micros = millis * 1000 + nano_of_milli // 1000
                return micros
        elif type_name == 'VARIANT':
            value_bytes = decoder.read_bytes()
            metadata_bytes = decoder.read_bytes()
            return {'value': value_bytes, 'metadata': metadata_bytes}
        else:
            raise ValueError(f"Unsupported atomic type: {type_name}")

    elif isinstance(data_type, ArrayType):
        return _read_array(decoder, data_type.element)

    elif isinstance(data_type, VectorType):
        return _read_vector(decoder, data_type.element)

    elif isinstance(data_type, MapType):
        keys = _read_array_elements(decoder, data_type.key)
        values = _read_array_elements(decoder, data_type.value)
        return list(zip(keys, values))

    elif isinstance(data_type, MultisetType):
        keys = _read_array_elements(decoder, data_type.element)
        counts = _read_array_elements(decoder, AtomicType("INT"))
        return list(zip(keys, counts))

    elif isinstance(data_type, RowType):
        return _read_nested_row(decoder, data_type)

    else:
        raise ValueError(f"Unsupported data type: {data_type}")


def _read_array(decoder: _RowDecoder, element_type) -> list:
    return _read_array_elements(decoder, element_type)


def _read_array_elements(decoder: _RowDecoder, element_type) -> list:
    size = decoder.read_var_int()
    null_bitmap_bytes = (size + 7) // 8
    null_bitmap = decoder.data[decoder.pos:decoder.pos + null_bitmap_bytes]
    decoder.pos += null_bitmap_bytes

    elements = []
    for i in range(size):
        is_null = (null_bitmap[i // 8] & (1 << (i % 8))) != 0
        if is_null:
            elements.append(None)
        else:
            elements.append(_read_field(decoder, element_type))
    return elements


def _read_vector(decoder: _RowDecoder, element_type) -> list:
    size = decoder.read_var_int()
    elements = []
    for _ in range(size):
        elements.append(_read_field(decoder, element_type))
    return elements


def _read_nested_row(decoder: _RowDecoder, row_type: RowType) -> dict:
    fields = row_type.fields
    arity = len(fields)
    header_size = (arity + 7) // 8
    null_bitmap = decoder.data[decoder.pos:decoder.pos + header_size]
    decoder.pos += header_size

    result = {}
    for i, field in enumerate(fields):
        is_null = (null_bitmap[i // 8] & (1 << (i % 8))) != 0
        if is_null:
            result[field.name] = None
        else:
            result[field.name] = _read_field(decoder, field.type)
    return result


def _parse_decimal_params(type_name: str) -> tuple:
    import re
    match = re.fullmatch(r'DECIMAL\((\d+),\s*(\d+)\)', type_name)
    if match:
        return int(match.group(1)), int(match.group(2))
    match = re.fullmatch(r'DECIMAL\((\d+)\)', type_name)
    if match:
        return int(match.group(1)), 0
    return 10, 0


def _parse_timestamp_precision(type_name: str) -> int:
    import re
    match = re.search(r'\((\d+)\)', type_name)
    if match:
        return int(match.group(1))
    return 6
