"""zg_storage.types module."""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Union

from pydantic import BaseModel, Field


class ShardConfig(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    shard_id: int = Field(alias="shardId")
    num_shard: int = Field(alias="numShard")

    model_config = {"populate_by_name": True}


class ShardedNode(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    url: str
    config: ShardConfig
    latency: Optional[int] = None
    since: Optional[int] = None


class ShardedNodes(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    trusted: List[ShardedNode] = Field(default_factory=list)
    discovered: List[ShardedNode] = Field(default_factory=list)


class IPLocation(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    city: Optional[str] = None
    region: Optional[str] = None
    country: Optional[str] = None
    location: Optional[str] = Field(default=None, alias="loc")
    timezone: Optional[str] = None

    model_config = {"populate_by_name": True}


class Transaction(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    stream_ids: List[str] = Field(default_factory=list, alias="streamIds")
    data: Optional[Union[str, List[int]]] = None
    data_merkle_root: Optional[str] = Field(default=None, alias="dataMerkleRoot")
    start_entry_index: Optional[int] = Field(default=None, alias="startEntryIndex")
    size: Optional[int] = None
    seq: Optional[int] = None

    model_config = {"populate_by_name": True}


class FileInfo(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    tx: Transaction
    finalized: bool
    is_cached: bool = Field(alias="isCached")
    uploaded_seg_num: int = Field(alias="uploadedSegNum")
    pruned: bool

    model_config = {"populate_by_name": True}


class SegmentWithProof(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    root: str
    data: str
    index: int
    proof: Dict[str, Any]
    file_size: int = Field(alias="fileSize")

    model_config = {"populate_by_name": True}


class Value(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    version: int
    data: str
    size: int


class KeyValue(BaseModel):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    version: int
    key: str
    data: str
    size: int
