"""zg_storage.transfer.downloader module."""

from __future__ import annotations

import logging
import os
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Iterable, Optional

from pydantic import BaseModel

from ..core.data import DataSink, FileDataSink
from ..node import ZgsNodeClient
from ..types import FileInfo, ShardConfig
from ..utils import (
    DEFAULT_CHUNK_SIZE,
    DEFAULT_SEGMENT_MAX_CHUNKS,
    DEFAULT_SEGMENT_SIZE,
    build_merkle_root,
    compute_padded_chunks,
    num_splits,
    segment_root,
)


class DownloadOption(BaseModel):
    """Download configuration options.

    Purpose:
        Control proof usage and parallel routines."""

    with_proof: bool = False
    routines: int = 1


@dataclass
class DownloadNode:
    """DownloadNode class.

    Purpose:
        SDK component type."""

    client: ZgsNodeClient
    shard: ShardConfig


class Downloader:
    """Parallel downloader for stored files.

    Purpose:
        Fetch segments from nodes and reconstruct the file with validation.
    Notes:
        Supports proof-based downloads when requested."""

    def __init__(self, nodes: Optional[list[DownloadNode]] = None) -> None:
        self._nodes = nodes or []
        self._logger = logging.getLogger("zg_storage")
        self._logger.propagate = True

    @staticmethod
    def _segment_range(start_entry_index: int, size: int) -> tuple[int, int, int]:
        """Return start segment, end segment, and total chunk count."""
        num_chunks = num_splits(size, DEFAULT_CHUNK_SIZE)
        start_segment_index = start_entry_index // DEFAULT_SEGMENT_MAX_CHUNKS
        end_segment_index = (start_entry_index + num_chunks - 1) // DEFAULT_SEGMENT_MAX_CHUNKS
        return start_segment_index, end_segment_index, num_chunks

    @staticmethod
    def _select_node(nodes: Iterable[DownloadNode], segment_index: int) -> Optional[DownloadNode]:
        """Pick a node that owns the shard for the given segment index."""
        for node in nodes:
            if node.shard.num_shard < 2:
                return node
            if segment_index % node.shard.num_shard == node.shard.shard_id:
                return node
        return None

    def download(
        self, root: str, target: str | DataSink, option: Optional[DownloadOption] = None
    ) -> None:
        if not self._nodes:
            raise ValueError("storage node not specified")
        opt = option or DownloadOption()
        self._logger.info(
            "Download start root=%s target=%s with_proof=%s",
            root,
            target,
            opt.with_proof,
        )
        if self._logger.isEnabledFor(logging.DEBUG):
            self._logger.debug(
                "Download start root=%s target=%s with_proof=%s",
                root,
                target,
                opt.with_proof,
            )
        info = self._nodes[0].client.get_file_info(root, need_available=True)
        self._download_parallel_with_info(
            root, target, info, with_proof=opt.with_proof, routines=opt.routines
        )
        self._logger.info("Download completed target=%s", target)

    def _download_parallel_with_info(
        self,
        root: str,
        target: str | DataSink,
        info: FileInfo,
        with_proof: bool,
        routines: int,
    ) -> None:
        """Download segments in parallel using known file metadata."""
        sink: DataSink
        if isinstance(target, str):
            if os.path.exists(target):
                raise FileExistsError(f"file already exists: {target}")
            sink = FileDataSink(target)
            ctx = sink
        else:
            sink = target
            ctx = None

        start_seg, end_seg, num_chunks = self._segment_range(
            info.tx.start_entry_index or 0, info.tx.size or 0
        )
        file_size = int(info.tx.size or 0)

        self._logger.info(
            "Parallel download root=%s target=%s size=%s segments=%s-%s chunks=%s routines=%s",
            root,
            target,
            file_size,
            start_seg,
            end_seg,
            num_chunks,
            routines,
        )

        lock = threading.Lock()

        if ctx is not None:
            ctx.__enter__()
        try:
            if isinstance(sink, FileDataSink):
                sink.init_size(file_size)

            def fetch_segment(seg_index: int) -> tuple[int, bytes]:
                """Fetch a single segment and return its index and bytes."""
                local_segment_index = seg_index - start_seg
                start_index = local_segment_index * DEFAULT_SEGMENT_MAX_CHUNKS
                end_index = min(start_index + DEFAULT_SEGMENT_MAX_CHUNKS, num_chunks)

                node = self._select_node(self._nodes, seg_index)
                if node is None:
                    raise RuntimeError(f"no node covers segment {seg_index}")

                if with_proof:
                    data = node.client.download_segment_with_proof_by_tx_seq(
                        info.tx.seq, local_segment_index
                    )
                    if self._logger.isEnabledFor(logging.DEBUG):
                        self._logger.debug(
                            "Downloaded segment %s from %s with proof",
                            seg_index,
                            node.client.url,
                        )
                else:
                    data = node.client.download_segment_by_tx_seq(
                        info.tx.seq, start_index, end_index
                    )
                    if self._logger.isEnabledFor(logging.DEBUG):
                        self._logger.debug(
                            "Downloaded segment %s from %s",
                            seg_index,
                            node.client.url,
                        )

                if data is None:
                    raise RuntimeError(f"segment {seg_index} not found")

                if len(data) % DEFAULT_CHUNK_SIZE != 0:
                    raise RuntimeError("invalid segment length")

                if seg_index == end_seg:
                    last_chunk_size = file_size % DEFAULT_CHUNK_SIZE
                    if last_chunk_size > 0:
                        data = data[: len(data) - (DEFAULT_CHUNK_SIZE - last_chunk_size)]

                return seg_index, data

            futures = []
            with ThreadPoolExecutor(max_workers=routines) as executor:
                for seg_index in range(start_seg, end_seg + 1):
                    futures.append(executor.submit(fetch_segment, seg_index))

                for future in as_completed(futures):
                    seg_index, data = future.result()
                    local_segment_index = seg_index - start_seg
                    offset = local_segment_index * DEFAULT_SEGMENT_SIZE
                    with lock:
                        sink.write_at(offset, data)

            self._validate_download(root, sink, file_size)
            self._logger.info("Validation succeeded root=%s target=%s", root, target)
        finally:
            if ctx is not None:
                ctx.__exit__(None, None, None)

    def _validate_download(self, root: str, source: DataSink, file_size: int) -> None:
        """Validate the downloaded file size and Merkle root."""
        actual_size = source.size()
        if actual_size != file_size:
            raise RuntimeError(f"file size mismatch: expected={file_size} actual={actual_size}")

        chunks = num_splits(file_size, DEFAULT_CHUNK_SIZE)
        padded_chunks = compute_padded_chunks(chunks)
        padded_size = padded_chunks * DEFAULT_CHUNK_SIZE
        num_segments = num_splits(padded_size, DEFAULT_SEGMENT_SIZE)

        segment_roots: list[bytes] = []
        for seg_index in range(num_segments):
            offset = seg_index * DEFAULT_SEGMENT_SIZE
            read_size = min(DEFAULT_SEGMENT_SIZE, padded_size - offset)
            if offset >= file_size:
                segment_bytes = b"\x00" * read_size
            else:
                data = source.read_at(offset, read_size)
                if len(data) < read_size:
                    data = data + b"\x00" * (read_size - len(data))
                segment_bytes = data
            segment_roots.append(segment_root(segment_bytes))

        root_hash = build_merkle_root(segment_roots, hash_leaves=False)
        expected = root.lower()
        actual = "0x" + root_hash.hex()
        if actual.lower() != expected:
            raise RuntimeError(f"merkle root mismatch: expected={expected} actual={actual}")


class NodeDownloaderConfig:
    """Configuration container for `NodeDownloader`.

    Purpose:
        Collect node URLs and RPC timeouts."""

    def __init__(
        self,
        nodes: list[str],
        rpc_timeout: float = 30.0,
    ) -> None:
        self.nodes = nodes
        self.rpc_timeout = rpc_timeout


class NodeDownloader(Downloader):
    """Downloader bound to a fixed node list.

    Purpose:
        Reuse the shared download pipeline with explicit nodes."""

    @classmethod
    def from_config(cls, config: NodeDownloaderConfig) -> "NodeDownloader":
        """Build a node-backed downloader from a config."""
        if not config.nodes:
            raise ValueError("storage node not specified")
        nodes: list[DownloadNode] = []
        for url in config.nodes:
            client = ZgsNodeClient(url, timeout=config.rpc_timeout)
            shard = client.get_shard_config()
            nodes.append(DownloadNode(client=client, shard=shard))
        return cls(nodes=nodes)
