# 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
import zlib
from typing import BinaryIO, List, Optional

from pypaimon.table.row.blob import Blob, BlobData, BlobDescriptor, BlobConsumer
from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor


class BlobFormatWriter:
    VERSION = 1
    MAGIC_NUMBER = 1481511375
    NULL_LENGTH = -1
    PLACE_HOLDER_LENGTH = -2
    BUFFER_SIZE = 4096
    METADATA_SIZE = 12  # 8-byte length + 4-byte CRC

    def __init__(self, output_stream: BinaryIO,
                 blob_consumer: Optional[BlobConsumer] = None,
                 file_path: Optional[str] = None):
        self.output_stream = output_stream
        self._blob_consumer = blob_consumer
        self._file_path = file_path
        self.lengths: List[int] = []
        self.position = 0

    def add_element(self, row) -> None:
        if not hasattr(row, 'values') or len(row.values) != 1:
            raise ValueError("BlobFormatWriter only supports one field")

        blob_field_name = row.fields[0].name
        blob_value = row.values[0]
        if blob_value is None:
            self.lengths.append(self.NULL_LENGTH)
            if self._blob_consumer is not None:
                self._blob_consumer(blob_field_name, None)
            return

        if not isinstance(blob_value, Blob):
            raise ValueError("Field must be Blob/BlobData instance")

        if blob_value is Blob.PLACE_HOLDER:
            self.lengths.append(self.PLACE_HOLDER_LENGTH)
            return

        previous_pos = self.position
        crc32 = 0  # Initialize CRC32

        # Write magic number
        magic_bytes = struct.pack('<I', self.MAGIC_NUMBER)  # Little endian
        crc32 = self._write_with_crc(magic_bytes, crc32)

        blob_pos = self.position

        # Write blob data
        if isinstance(blob_value, BlobData):
            data = blob_value.to_data()
            crc32 = self._write_with_crc(data, crc32)
        else:
            # Stream from BlobRef/Blob
            stream = blob_value.new_input_stream()
            try:
                chunk = stream.read(self.BUFFER_SIZE)
                while chunk:
                    crc32 = self._write_with_crc(chunk, crc32)
                    chunk = stream.read(self.BUFFER_SIZE)
            finally:
                stream.close()

        blob_length = self.position - blob_pos

        # Calculate total length including magic + data + metadata (length + CRC)
        bin_length = self.position - previous_pos + self.METADATA_SIZE
        self.lengths.append(bin_length)

        # Write length (8 bytes, little endian)
        length_bytes = struct.pack('<Q', bin_length)
        self.output_stream.write(length_bytes)
        self.position += 8

        # Write CRC32 (4 bytes, little endian)
        crc_bytes = struct.pack('<I', crc32 & 0xffffffff)
        self.output_stream.write(crc_bytes)
        self.position += 4

        if self._blob_consumer is not None:
            descriptor = BlobDescriptor(self._file_path, blob_pos, blob_length)
            flush = self._blob_consumer(blob_field_name, descriptor)
            if flush:
                self.output_stream.flush()

    def _write_with_crc(self, data: bytes, crc32: int) -> int:
        crc32 = zlib.crc32(data, crc32)
        self.output_stream.write(data)
        self.position += len(data)
        return crc32

    def write_value(self, col_data, fields, uri_reader_factory=None) -> None:
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.table.row.row_kind import RowKind

        is_blob = hasattr(fields[0].type, 'type') and fields[0].type.type == "BLOB"

        if col_data is None:
            if not is_blob:
                raise RuntimeError("Null values are only supported for BLOB type fields")
            self.lengths.append(self.NULL_LENGTH)
            return

        if is_blob:
            if hasattr(col_data, 'as_py'):
                col_data = col_data.as_py()
            if isinstance(col_data, str):
                col_data = col_data.encode('utf-8')
            if isinstance(col_data, bytearray):
                col_data = bytes(col_data)

            if isinstance(col_data, bytes):
                if BlobDescriptor.is_blob_descriptor(col_data):
                    descriptor = BlobDescriptor.deserialize(col_data)
                    uri_reader = uri_reader_factory.create(descriptor.uri)
                    blob_value = Blob.from_descriptor(uri_reader, descriptor)
                else:
                    blob_value = BlobData(col_data)
            else:
                raise RuntimeError(
                    "Blob field value must be bytes/blob or serialized BlobDescriptor bytes."
                )
            row_values = [blob_value]
        else:
            row_values = [col_data]

        row = GenericRow(row_values, fields, RowKind.INSERT)
        self.add_element(row)

    def reach_target_size(self, target_size: int) -> bool:
        return self.position >= target_size

    def close(self) -> None:
        index_bytes = DeltaVarintCompressor.compress(self.lengths)
        self.output_stream.write(index_bytes)

        # Write header (index length + version)
        header = struct.pack('<I', len(index_bytes))  # Index length (4 bytes, little endian)
        header += struct.pack('<B', self.VERSION)  # Version (1 byte)
        self.output_stream.write(header)

        # Flush and close
        if hasattr(self.output_stream, 'flush'):
            self.output_stream.flush()
        if hasattr(self.output_stream, 'close'):
            self.output_stream.close()
