"""zg_storage.core.data module."""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Optional, Protocol, runtime_checkable


@runtime_checkable
class DataSource(Protocol):
    """DataSource class.

    Purpose:
        SDK component type."""

    def size(self) -> int:
        """Return total readable size in bytes."""
        raise NotImplementedError

    def read_at(self, offset: int, size: int) -> bytes:
        """Read bytes starting at the given offset."""
        raise NotImplementedError

    def close(self) -> None:
        """Release any held resources."""
        raise NotImplementedError


@runtime_checkable
class DataSink(Protocol):
    """DataSink class.

    Purpose:
        SDK component type."""

    def size(self) -> int:
        """Return total size in bytes."""
        raise NotImplementedError

    def write_at(self, offset: int, data: bytes) -> None:
        """Write bytes starting at the given offset."""
        raise NotImplementedError

    def read_at(self, offset: int, size: int) -> bytes:
        """Read bytes starting at the given offset."""
        raise NotImplementedError

    def close(self) -> None:
        """Release any held resources."""
        raise NotImplementedError


@dataclass
class FileDataSource:
    """File-backed random access data source.

    Purpose:
        Provide offset reads for upload segmentation."""

    path: str
    _file: Optional[object] = None
    _size: Optional[int] = None

    def __enter__(self) -> FileDataSource:
        """Open the file and cache size for random-access reads."""
        self._file = open(self.path, "rb")
        self._file.seek(0, 2)
        self._size = self._file.tell()
        self._file.seek(0)
        return self

    def __exit__(self, exc_type, exc, tb) -> None:  # type: ignore[override]
        """Close the file handle on context manager exit."""
        self.close()

    def size(self) -> int:
        """Return the total file size in bytes."""
        if self._size is None:
            raise RuntimeError("file not opened")
        return self._size

    def read_at(self, offset: int, size: int) -> bytes:
        """Read bytes from the file at the given offset."""
        if self._file is None:
            raise RuntimeError("file not opened")
        return os.pread(self._file.fileno(), size, offset)

    def close(self) -> None:
        """Close the underlying file handle."""
        if self._file is not None:
            self._file.close()
            self._file = None


@dataclass
class BytesDataSource:
    """In-memory random access data source.

    Purpose:
        Provide offset reads for upload segmentation."""

    data: bytes | bytearray | memoryview
    offset: int = 0

    def __post_init__(self) -> None:
        """Validate data is non-empty and offset is within bounds."""
        if len(self.data) == 0:
            raise ValueError("data is empty")
        if self.offset < 0 or self.offset > len(self.data):
            raise ValueError("invalid offset")

    def size(self) -> int:
        """Return the readable length from the current offset."""
        return len(self.data) - self.offset

    def read_at(self, offset: int, size: int) -> bytes:
        """Read bytes from the buffer at the given offset."""
        mv = self.data if isinstance(self.data, memoryview) else memoryview(self.data)
        start = self.offset + offset
        end = start + size
        return bytes(mv[start:end])

    def close(self) -> None:
        """No-op close for in-memory data sources."""
        return None


@dataclass
class FileDataSink:
    """File-backed random access data sink.

    Purpose:
        Provide offset writes for downloads."""

    path: str
    _file: Optional[object] = None
    _size: Optional[int] = None

    def __enter__(self) -> FileDataSink:
        """Open the file for random-access writes."""
        self._file = open(self.path, "wb+")
        return self

    def __exit__(self, exc_type, exc, tb) -> None:  # type: ignore[override]
        """Close the file handle on context manager exit."""
        self.close()

    def size(self) -> int:
        """Return the initialized size of the sink."""
        if self._size is None:
            raise RuntimeError("sink size not initialized")
        return self._size

    def init_size(self, size: int) -> None:
        """Initialize the sink size and pre-allocate the file."""
        if self._file is None:
            raise RuntimeError("file not opened")
        self._size = size
        self._file.truncate(size)

    def write_at(self, offset: int, data: bytes) -> None:
        """Write bytes to the file at the given offset."""
        if self._file is None:
            raise RuntimeError("file not opened")
        os.pwrite(self._file.fileno(), data, offset)

    def read_at(self, offset: int, size: int) -> bytes:
        """Read bytes from the file at the given offset."""
        if self._file is None:
            raise RuntimeError("file not opened")
        return os.pread(self._file.fileno(), size, offset)

    def close(self) -> None:
        """Close the underlying file handle."""
        if self._file is not None:
            self._file.close()
            self._file = None


@dataclass
class BytesData:
    """In-memory writable data buffer.

    Purpose:
        Provide offset writes and reads for downloads."""

    data: bytearray

    @classmethod
    def with_size(cls, size: int) -> BytesData:
        """Allocate an in-memory buffer of the given size."""
        return cls(bytearray(size))

    def size(self) -> int:
        """Return the buffer length."""
        return len(self.data)

    def write_at(self, offset: int, data: bytes) -> None:
        """Write bytes to the buffer at the given offset."""
        end = offset + len(data)
        self.data[offset:end] = data

    def read_at(self, offset: int, size: int) -> bytes:
        """Read a range from the buffer starting at the offset."""
        return bytes(self.data[offset : offset + size])

    def close(self) -> None:
        """No-op close for in-memory buffers."""
        return None
