"""zg_storage.indexer.__init__ module."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Dict, List, Optional

from ..node import ZgsNodeClient
from ..rpc import RpcClient
from ..transfer import Downloader, DownloadNode, NodeUploader, UploadOption
from ..types import IPLocation, ShardedNode, ShardedNodes

if TYPE_CHECKING:
    from ..transfer import DownloadOption


class IndexerClient:
    """Client for indexer RPC endpoints.

    Purpose:
        Select nodes, query locations, and orchestrate uploads/downloads via the indexer.
    Notes:
        Upload uses selected nodes then delegates to the shared upload pipeline."""

    def __init__(
        self,
        url: str,
        timeout: float = 30.0,
        evm_client=None,
        flow_address: Optional[str] = None,
    ) -> None:
        self._rpc = RpcClient(url, timeout=timeout)
        self._logger = logging.getLogger("zg_storage")
        self._logger.propagate = True
        self._evm_client = evm_client
        self._flow_address = flow_address

    def get_sharded_nodes(self) -> ShardedNodes:
        """Fetch all sharded nodes reported by the indexer."""
        result = self._rpc.call("indexer_getShardedNodes")
        return ShardedNodes.model_validate(result)

    def get_selected_nodes(
        self,
        expected_replica: int,
        method: str,
        full_trusted: bool,
        dropped: Optional[List[str]] = None,
    ) -> ShardedNodes:
        """Select nodes based on replica count and selection method."""
        params = [expected_replica, method, full_trusted, dropped or []]
        result = self._rpc.call("indexer_getSelectedNodes", params)
        return ShardedNodes.model_validate(result)

    def get_node_locations(self) -> Dict[str, IPLocation]:
        """Fetch node locations keyed by node URL."""
        result = self._rpc.call("indexer_getNodeLocations")
        return {k: IPLocation.model_validate(v) for k, v in result.items()}

    def get_file_locations(self, root: str) -> List[ShardedNode]:
        """Get node locations storing the given root."""
        self._logger.info("Indexer get_file_locations root=%s", root)
        if self._logger.isEnabledFor(logging.DEBUG):
            self._logger.debug("Indexer RPC call indexer_getFileLocations")
        result = self._rpc.call("indexer_getFileLocations", [root])
        return [ShardedNode.model_validate(item) for item in result]

    def select_nodes(
        self,
        expected_replica: int,
        dropped: Optional[List[str]] = None,
        method: str = "random",
        full_trusted: bool = False,
    ) -> ShardedNodes:
        """Convenience wrapper around `get_selected_nodes`."""
        return self.get_selected_nodes(expected_replica, method, full_trusted, dropped)

    def download(self, root: str, filename: str, option: Optional["DownloadOption"] = None) -> None:
        """Download a file by resolving nodes via the indexer."""
        opt = option
        with_proof = opt.with_proof if opt else False
        self._logger.info(
            "Indexer download root=%s filename=%s with_proof=%s",
            root,
            filename,
            with_proof,
        )
        if self._logger.isEnabledFor(logging.DEBUG):
            self._logger.debug(
                "Indexer download root=%s filename=%s with_proof=%s",
                root,
                filename,
                with_proof,
            )
        downloader = IndexerDownloader(self)
        downloader.download(root=root, filename=filename, option=option)

    def upload(
        self,
        file_path,
        tags: bytes,
        option: Optional["UploadOption"] = None,
        evm_client=None,
        flow_address: Optional[str] = None,
    ) -> tuple[str, str, int]:
        """Upload via indexer-selected nodes using the shared pipeline."""
        opt = option or UploadOption()
        active_evm = evm_client or self._evm_client
        active_flow = flow_address or self._flow_address or opt.flow_address
        if not active_flow:
            raise ValueError("flow_address is required (pass arg or set UploadOption.flow_address)")
        if active_evm is None:
            raise ValueError("evm_client is required (pass arg or set IndexerClient.evm_client)")
        nodes = self.get_selected_nodes(
            opt.expected_replica,
            opt.method,
            opt.full_trusted,
            [],
        )
        download_nodes: list[DownloadNode] = []
        for node in nodes.trusted + nodes.discovered:
            timeout = opt.rpc_timeout if opt.rpc_timeout is not None else 30.0
            client = ZgsNodeClient(node.url, timeout=timeout)
            download_nodes.append(DownloadNode(client=client, shard=node.config))

        uploader = NodeUploader(download_nodes, evm_client=active_evm, flow_address=active_flow)
        return uploader.upload(
            file_path=file_path,
            tags=tags,
            option=opt,
        )


class IndexerDownloader(Downloader):
    """Downloader that resolves nodes via the indexer.

    Purpose:
        Use indexer file locations and then perform a normal download."""

    def __init__(self, indexer: IndexerClient) -> None:
        self._indexer = indexer
        super().__init__([])

    def download(self, root: str, target: str, option: Optional["DownloadOption"] = None) -> None:
        """Resolve nodes via indexer and download to the target file."""
        nodes = self._indexer.get_file_locations(root)
        download_nodes: list[DownloadNode] = []
        for node in nodes:
            client = ZgsNodeClient(node.url)
            download_nodes.append(DownloadNode(client=client, shard=node.config))
        self._nodes = download_nodes
        super().download(root, target, option=option)
