# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright the Vortex contributors

from __future__ import annotations

from typing import final

import numpy as np

import pyarrow as pa

import pytest

from pcodec import ( # pyright: ignore[reportMissingTypeStubs]

ChunkConfig, # pyright: ignore[reportAttributeAccessIssue, reportUnknownVariableType]

)

from pcodec import ( # pyright: ignore[reportMissingTypeStubs]

wrapped as pco, # pyright: ignore[reportAttributeAccessIssue, reportUnknownVariableType]

)

from typing_extensions import override

import vortex as vx

@final

class PCodecArray(vx.PyArray):

@property

@override

def id(self):

return "pcodec.v0"

@override

def __len__(self) -> int:

return self._len

@property

@override

def dtype(self) -> vx.DType:

return self._dtype

def __init__(

self,

length: int,

dtype: vx.DType,

file_header: memoryview,

chunk_header: memoryview,

data: memoryview,

):

(fd, _bytes_read) = pco.FileDecompressor.new(file_header) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]

if dtype == vx.int_(64, nullable=True):

dt = "i64"

else:

raise ValueError(f"Unsupported dtype: {dtype}")

(cd, _bytes_read) = fd.read_chunk_meta(chunk_header, dt) # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]

dst = np.array([0] * length, dtype=np.int64)

cd.read_page_into( # pyright: ignore[reportUnknownMemberType]

data,

page_n=length,

dst=dst,

)

self._len = length

self._dtype = dtype

self._file_header = file_header

self._chunk_header = chunk_header

self._data = data

@classmethod

def encode(cls, array: pa.Array[pa.Scalar[pa.DataType]], config: ChunkConfig | None = None) -> PCodecArray: # pyright: ignore[reportUnknownParameterType]

assert array.null_count == 0, "Cannot compress arrays with nulls"

config = config or ChunkConfig() # pyright: ignore[reportUnknownVariableType]

fc = pco.FileCompressor() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]

file_header = fc.write_header() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]

cc = fc.chunk_compressor(array.to_numpy(), config) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]

chunk_header = cc.write_chunk_meta() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]

data = b""

for i, _n in enumerate(cc.n_per_page()): # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportUnknownVariableType]

data += cc.write_page(i) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]

return PCodecArray(

len(array),

vx.DType.from_arrow(array.type),

file_header, # pyright: ignore[reportUnknownArgumentType]

chunk_header, # pyright: ignore[reportUnknownArgumentType]

memoryview(data), # pyright: ignore[reportUnknownArgumentType]

)

@override

@classmethod

def decode(cls, parts: vx.SerializedArray, ctx: vx.ArrayContext, dtype: vx.DType, len: int) -> vx.Array:

"""Decode the serialized array parts into an array."""

assert pco

raise NotImplementedError

@pytest.mark.skip(reason="Not implemented yet")

def test_pcodec():

_ = PCodecArray.encode(pa.array([0, 1, 2, 3, 4])) # pyright: ignore[reportUnknownMemberType]

vx.registry.register(PCodecArray)