"""zg_storage.node.zgs_node module."""

from __future__ import annotations

from ..rpc import RpcClient
from ..types import FileInfo, ShardConfig
from ..utils import decode_bytes


class ZgsNodeClient:
    """Client for storage node RPC endpoints.

    Purpose:
        Fetch shard config, file info, and upload/download segments.
    Notes:
        Methods map directly to `zgs_*` RPC calls."""

    def __init__(self, url: str, timeout: float = 30.0) -> None:
        self._rpc = RpcClient(url, timeout=timeout)

    @property
    def url(self) -> str:
        """Return the node RPC URL."""
        return self._rpc.url

    def get_shard_config(self) -> ShardConfig:
        """Fetch the node shard configuration."""
        result = self._rpc.call("zgs_getShardConfig")
        return ShardConfig.model_validate(result)

    def get_file_info(self, root: str, need_available: bool = True) -> FileInfo:
        """Fetch file info by data root."""
        result = self._rpc.call("zgs_getFileInfo", [root, need_available])
        return FileInfo.model_validate(result)

    def get_file_info_by_tx_seq(self, tx_seq: int) -> FileInfo:
        """Fetch file info by transaction sequence."""
        result = self._rpc.call("zgs_getFileInfoByTxSeq", [tx_seq])
        return FileInfo.model_validate(result)

    def download_segment(self, root: str, start_index: int, end_index: int) -> bytes:
        """Download segment data by root and chunk indices."""
        result = self._rpc.call("zgs_downloadSegment", [root, start_index, end_index])
        if not result:
            return b""
        return decode_bytes(result)

    def download_segment_by_tx_seq(self, tx_seq: int, start_index: int, end_index: int) -> bytes:
        """Download segment data by tx sequence and chunk indices."""
        result = self._rpc.call("zgs_downloadSegmentByTxSeq", [tx_seq, start_index, end_index])
        if not result:
            return b""
        return decode_bytes(result)

    def download_segment_with_proof(self, root: str, index: int) -> bytes:
        """Download segment data with proof by root and segment index."""
        result = self._rpc.call("zgs_downloadSegmentWithProof", [root, index])
        if not result:
            return b""
        data = result["data"] if isinstance(result, dict) else result.data
        return decode_bytes(data)

    def download_segment_with_proof_by_tx_seq(self, tx_seq: int, index: int) -> bytes:
        """Download segment data with proof by tx sequence and segment index."""
        result = self._rpc.call("zgs_downloadSegmentWithProofByTxSeq", [tx_seq, index])
        if not result:
            return b""
        data = result["data"] if isinstance(result, dict) else result.data
        return decode_bytes(data)

    def upload_segments_by_tx_seq(self, segments: list[dict], tx_seq: int) -> int:
        return self._rpc.call("zgs_uploadSegmentsByTxSeq", [segments, tx_seq])
