"""zg_storage.transfer.uploader module."""

from __future__ import annotations

import logging
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import nullcontext
from dataclasses import dataclass
from enum import Enum
from typing import Optional

from pydantic import BaseModel

from ..contracts.flow_contract import FlowContractClient
from ..contracts.market_contract import MarketContractClient
from ..core.data import DataSource, FileDataSource
from ..core.flow import FlowBuilder, submission_fee, submission_root
from ..node import ZgsNodeClient
from ..utils import (
    DEFAULT_CHUNK_SIZE,
    DEFAULT_SEGMENT_MAX_CHUNKS,
    DEFAULT_SEGMENT_SIZE,
    compute_padded_chunks,
    encode_bytes_base64,
    num_splits,
    segment_root,
)
from .downloader import DownloadNode
from .merkle import MerkleTree

DATA_ALREADY_EXISTS_ERROR = "Invalid params: root; data: already uploaded and finalized"
SEGMENT_ALREADY_EXISTS_ERROR = "segment has already been uploaded or is being uploaded"
TOO_MANY_DATA_ERROR = "too many data writing"
TOO_MANY_DATA_RETRIES = 12


def is_duplicate_error(message: str) -> bool:
    return DATA_ALREADY_EXISTS_ERROR in message or SEGMENT_ALREADY_EXISTS_ERROR in message


def is_too_many_data_error(message: str) -> bool:
    return TOO_MANY_DATA_ERROR in message


class FinalityRequirement(str, Enum):
    """FinalityRequirement class.

    Purpose:
        SDK component type."""

    FILE_FINALIZED = "file_finalized"
    TRANSACTION_PACKED = "transaction_packed"


class UploadOption(BaseModel):
    """Upload configuration options.

    Purpose:
        Control replica selection, task sizing, finality, and RPC settings."""

    submitter: Optional[str] = None
    tags: bytes = b""
    finality_required: FinalityRequirement = FinalityRequirement.TRANSACTION_PACKED
    task_size: int = 10
    expected_replica: int = 1
    skip_tx: bool = True
    fast_mode: bool = True
    fee: Optional[int] = None
    nonce: Optional[int] = None
    max_gas_price: Optional[int] = None
    flow_address: Optional[str] = None
    market_address: Optional[str] = None
    n_retries: int = 0
    step: int = 15
    method: str = "min"
    full_trusted: bool = True
    rpc_timeout: Optional[float] = None
    routines: int = 4
    wait_receipt: bool = True


@dataclass
class UploadResult:
    """Pydantic model for RPC data.

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

    tx_hashes: list[str]
    roots: list[str]


class Uploader:
    """Upload pipeline for files or in-memory data.

    Purpose:
        Submit log entry to the flow contract and upload segments to nodes.
    Notes:
        Uses shard-aware task scheduling and optional parallelism."""

    def __init__(self) -> None:
        self._logger = logging.getLogger("zg_storage")
        self._logger.propagate = True

    def splitable_upload(self, *args, **kwargs) -> UploadResult:
        raise NotImplementedError("splitable_upload is not implemented yet")

    def submit_log_entry(
        self,
        file_path: str | DataSource,
        tags: bytes,
        flow_address: Optional[str],
        evm_client,
        submitter: Optional[str] = None,
        fee_wei: Optional[int] = None,
        gas: Optional[int] = None,
        gas_price: Optional[int] = None,
        nonce: Optional[int] = None,
        market_address: Optional[str] = None,
        wait_receipt: bool = False,
    ) -> str:
        """Submit a flow log entry transaction and optionally wait for receipt."""
        if not flow_address:
            raise ValueError("flow_address is required")
        submitter_addr = submitter or evm_client.address()
        builder = FlowBuilder(file_path, tags)
        submission = builder.create_submission(submitter_addr)
        flow = FlowContractClient(evm_client, flow_address)
        if fee_wei is None:
            market_addr = market_address or flow.market_address()
            market = MarketContractClient(evm_client, market_addr)
            price_per_sector = market.price_per_sector()
            fee_wei = submission_fee(submission, price_per_sector)
        self._logger.info(
            "Submit log entry submitter=%s fee_wei=%s gas=%s gas_price=%s nonce=%s",
            submitter_addr,
            fee_wei,
            gas,
            gas_price,
            nonce,
        )
        tx_hash = flow.submit(
            submission, value_wei=fee_wei, gas=gas, gas_price=gas_price, nonce=nonce
        )
        self._logger.info("Succeeded to send transaction to append log entry hash=%s", tx_hash)
        if wait_receipt:
            receipt = evm_client.w3.eth.wait_for_transaction_receipt(tx_hash)
            self._logger.info(
                "Transaction receipt received hash=%s status=%s block=%s",
                tx_hash,
                getattr(receipt, "status", None),
                getattr(receipt, "blockNumber", None),
            )
        return tx_hash

    def wait_for_log_entry(
        self,
        nodes: list[DownloadNode],
        root_hex: str,
        tx_seq: Optional[int] = None,
        finality_required: FinalityRequirement = FinalityRequirement.TRANSACTION_PACKED,
        timeout_secs: int = 600,
        poll_interval: float = 2.0,
    ):
        """Poll storage nodes until the log entry is available at the desired finality."""
        deadline = time.time() + timeout_secs
        last_err: Optional[Exception] = None
        while time.time() < deadline:
            info = None
            all_ok = True
            for node in nodes:
                try:
                    if tx_seq is not None:
                        info = node.client.get_file_info_by_tx_seq(tx_seq)
                    else:
                        info = node.client.get_file_info(root_hex, need_available=True)
                except Exception as exc:  # noqa: BLE001
                    last_err = exc
                    all_ok = False
                    break
                if not info or not info.tx or info.tx.seq is None:
                    all_ok = False
                    break
                if finality_required == FinalityRequirement.FILE_FINALIZED and not info.finalized:
                    all_ok = False
                    break
            if all_ok and info is not None:
                return info
            time.sleep(poll_interval)
        if last_err:
            raise TimeoutError(f"wait_for_log_entry timeout: {last_err}") from last_err
        raise TimeoutError("wait_for_log_entry timeout")

    def upload_slow(
        self,
        file_path: str | DataSource,
        tags: bytes,
        flow_address: str,
        evm_client,
        nodes: list[DownloadNode],
        option: Optional[UploadOption] = None,
    ) -> tuple[str, str, int]:
        """Upload a file or data source through the flow log and segment upload path."""
        if self._logger.isEnabledFor(logging.DEBUG):
            self._logger.debug("Upload slow start file_path=%s", file_path)
        else:
            self._logger.info("Upload slow start file_path=%s", type(file_path).__name__)
        opt = option or UploadOption()
        submitter_addr = evm_client.address()
        builder = FlowBuilder(file_path, tags)
        submission = builder.create_submission(submitter_addr)
        root_hex = "0x" + submission_root(submission).hex()

        info = None
        if opt.skip_tx:
            for node in nodes:
                try:
                    existing = node.client.get_file_info(root_hex, need_available=True)
                except Exception:  # noqa: BLE001
                    existing = None
                if existing and existing.tx and existing.tx.seq is not None:
                    info = existing
                    self._logger.info(
                        "Skip submit tx, existing log entry found root=%s tx_seq=%s node=%s",
                        root_hex,
                        existing.tx.seq,
                        node.client._rpc.url,
                    )
                    break

        tx_hash = ""
        if not opt.skip_tx or info is None:
            self._logger.info("Prepare to submit log entry root=%s", root_hex)
            tx_hash = self.submit_log_entry(
                file_path=file_path,
                tags=tags or opt.tags,
                flow_address=flow_address or opt.flow_address,
                evm_client=evm_client,
                submitter=opt.submitter or submitter_addr,
                fee_wei=opt.fee,
                gas=None,
                gas_price=opt.max_gas_price,
                nonce=opt.nonce,
                market_address=opt.market_address,
                wait_receipt=opt.wait_receipt,
            )

            info = self.wait_for_log_entry(
                nodes,
                root_hex,
                tx_seq=None,
                finality_required=FinalityRequirement.TRANSACTION_PACKED,
            )

        if info is None or info.tx is None or info.tx.seq is None:
            raise RuntimeError("log entry not available on storage nodes")
        # Ensure every node can resolve the transaction by tx_seq before uploading segments.
        _ = self.wait_for_log_entry(
            nodes,
            root_hex,
            tx_seq=info.tx.seq if info and info.tx else None,
            finality_required=FinalityRequirement.TRANSACTION_PACKED,
        )
        if self._logger.isEnabledFor(logging.INFO):
            self._logger.info(
                "Log entry info tx_seq=%s data_root=%s start_entry_index=%s size=%s",
                getattr(info.tx, "seq", None),
                getattr(info.tx, "data_merkle_root", None),
                getattr(info.tx, "start_entry_index", None),
                getattr(info.tx, "size", None),
            )

        self._upload_segments(
            file_path,
            info.tx.seq,
            nodes,
            expected_root_hex=root_hex,
            task_size=opt.task_size,
            start_entry_index=info.tx.start_entry_index,
            routines=opt.routines,
        )

        if opt.finality_required == FinalityRequirement.FILE_FINALIZED:
            _ = self.wait_for_log_entry(
                nodes,
                root_hex,
                tx_seq=info.tx.seq if info and info.tx else None,
                finality_required=FinalityRequirement.FILE_FINALIZED,
            )
        self._logger.info(
            "Upload slow completed root=%s tx=%s tx_seq=%s", root_hex, tx_hash, info.tx.seq
        )
        return tx_hash, root_hex, info.tx.seq

    def _upload_segments(
        self,
        file_path: str | DataSource,
        tx_seq: int,
        nodes: list[DownloadNode],
        expected_root_hex: Optional[str] = None,
        task_size: int = 10,
        start_entry_index: Optional[int] = None,
        routines: int = 1,
    ) -> None:
        """Upload all file segments to storage nodes, optionally in parallel."""
        source: DataSource
        if isinstance(file_path, str):
            source = FileDataSource(file_path)
            ctx = source
        else:
            source = file_path
            ctx = nullcontext()
        with ctx:
            file_size = source.size()
            if file_size <= 0:
                raise ValueError("file is empty")

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

            def read_padded_segment(seg_index: int) -> bytes:
                """Read a segment and pad it to the expected size."""
                offset = seg_index * DEFAULT_SEGMENT_SIZE
                read_size = min(DEFAULT_SEGMENT_SIZE, padded_size - offset)
                if offset >= file_size:
                    return b"\x00" * read_size
                data = source.read_at(offset, read_size)
                if len(data) < read_size:
                    data = data + b"\x00" * (read_size - len(data))
                return data

            def read_upload_segment(seg_index: int) -> bytes:
                """Read a segment and trim to the real data length when needed."""
                start_index = seg_index * DEFAULT_SEGMENT_MAX_CHUNKS
                segment = read_padded_segment(seg_index)
                if start_index + len(segment) // DEFAULT_CHUNK_SIZE >= chunks:
                    expected_len = DEFAULT_CHUNK_SIZE * max(0, chunks - start_index)
                    segment = segment[:expected_len]
                return segment

            segment_roots: list[bytes] = []
            for seg_index in range(num_segments_padded):
                segment_roots.append(segment_root(read_padded_segment(seg_index)))

            tree = MerkleTree(segment_roots)
            root_hex = "0x" + tree.root_hash().hex()
            self._logger.info(
                "Upload segments prepared tx_seq=%s file_size=%s segments_real=%s segments_padded=%s root=%s nodes=%s",
                tx_seq,
                file_size,
                num_segments_real,
                num_segments_padded,
                root_hex,
                [node.client._rpc.url for node in nodes],
            )
            computed_root_hex = "0x" + tree.root_hash().hex()
            if expected_root_hex and expected_root_hex.lower() != computed_root_hex.lower():
                raise RuntimeError(
                    f"root mismatch before upload: expected={expected_root_hex} computed={computed_root_hex}"
                )
            if len(computed_root_hex) != 66:
                raise RuntimeError(f"invalid root length: {computed_root_hex}")
            if not tree.root_hash():
                raise RuntimeError("computed empty root for segments")

            start_segment_index = 0
            if start_entry_index is not None:
                start_segment_index = start_entry_index // DEFAULT_SEGMENT_MAX_CHUNKS
            end_segment_index = start_segment_index + num_segments_real - 1

            def _next_segment_index(start_index: int, shard_id: int, num_shard: int) -> int:
                """Find the first segment index covered by a shard."""
                if num_shard <= 0:
                    return start_index
                remainder = start_index % num_shard
                if remainder <= shard_id:
                    return start_index + (shard_id - remainder)
                return start_index + (num_shard - remainder + shard_id)

            client_tasks: list[list[tuple[int, int, int]]] = []
            for node_idx, node in enumerate(nodes):
                info = None
                try:
                    info = node.client.get_file_info(root_hex, need_available=True)
                except Exception:  # noqa: BLE001
                    info = None
                if info is not None and getattr(info, "finalized", False):
                    continue

                num_shard = int(node.shard.num_shard)
                shard_id = int(node.shard.shard_id)
                if num_shard <= 0:
                    num_shard = 1

                seg_index = _next_segment_index(start_segment_index, shard_id, num_shard)
                tasks: list[tuple[int, int, int]] = []
                while seg_index <= end_segment_index:
                    local_index = seg_index - start_segment_index
                    tasks.append((node_idx, local_index, num_shard))
                    seg_index += num_shard * task_size
                if tasks:
                    client_tasks.append(tasks)

            client_tasks.sort(key=len, reverse=True)
            tasks: list[tuple[int, int, int]] = []
            if client_tasks:
                for task_index in range(len(client_tasks[0])):
                    for i in range(len(client_tasks)):
                        if task_index < len(client_tasks[i]):
                            tasks.append(client_tasks[i][task_index])

            def upload_task(node_idx: int, start_local_index: int, num_shard: int) -> None:
                """Upload a batch of segments for a single node and shard."""
                batch: list[dict] = []
                batch_idx: list[int] = []
                seg_index = start_local_index
                for _ in range(task_size):
                    if seg_index >= num_segments_real:
                        break
                    proof = tree.proof_at(seg_index)
                    seg_bytes = read_upload_segment(seg_index)
                    segment = {
                        "root": "0x" + tree.root_hash().hex(),
                        "data": encode_bytes_base64(seg_bytes),
                        "index": seg_index,
                        "proof": {
                            "lemma": ["0x" + h.hex() for h in proof["lemma"]],
                            "path": proof["path"],
                        },
                        "fileSize": file_size,
                    }
                    batch.append(segment)
                    batch_idx.append(seg_index)
                    seg_index += num_shard

                if not batch:
                    return
                self._logger.debug(
                    "Uploading segments batch node=%s tx_seq=%s indices=%s root=%s",
                    nodes[node_idx].client._rpc.url,
                    tx_seq,
                    batch_idx,
                    batch[0]["root"] if batch else None,
                )
                for attempt in range(1, TOO_MANY_DATA_RETRIES + 1):
                    try:
                        nodes[node_idx].client.upload_segments_by_tx_seq(batch, tx_seq)
                        return
                    except Exception as exc:  # noqa: BLE001
                        msg = str(exc)
                        if is_duplicate_error(msg):
                            self._logger.info(
                                "Segments already uploaded or finalized on node=%s",
                                nodes[node_idx].client._rpc.url,
                            )
                            return
                        if is_too_many_data_error(msg) and attempt < TOO_MANY_DATA_RETRIES:
                            time.sleep(10)
                            continue
                        raise

            if routines <= 1 or len(tasks) <= 1:
                for node_idx, start_local_index, num_shard in tasks:
                    upload_task(node_idx, start_local_index, num_shard)
                return

            max_workers = max(1, routines)
            futures = []
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                for node_idx, start_local_index, num_shard in tasks:
                    futures.append(
                        executor.submit(upload_task, node_idx, start_local_index, num_shard)
                    )
                for future in as_completed(futures):
                    future.result()


class NodeUploaderConfig:
    """Configuration container for `NodeUploader`.

    Purpose:
        Collect node URLs, EVM client, flow address, and RPC timeouts."""

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


class NodeUploader:
    """Uploader bound to a fixed node list.

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

    def __init__(
        self,
        nodes: list[DownloadNode],
        evm_client,
        flow_address: str,
        uploader: Optional[Uploader] = None,
    ) -> None:
        if not nodes:
            raise ValueError("storage node not specified")
        if not flow_address:
            raise ValueError("flow_address is required")
        self._nodes = nodes
        self._evm_client = evm_client
        self._flow_address = flow_address
        self._uploader = uploader or Uploader()

    @classmethod
    def from_config(
        cls,
        config: NodeUploaderConfig,
        uploader: Optional[Uploader] = None,
    ) -> "NodeUploader":
        """Build a node-backed uploader 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,
            evm_client=config.evm_client,
            flow_address=config.flow_address,
            uploader=uploader,
        )

    def upload(
        self,
        file_path: str | DataSource,
        tags: bytes,
        option: Optional[UploadOption] = None,
    ) -> tuple[str, str, int]:
        """Upload via the configured node list and flow contract."""
        return self._uploader.upload_slow(
            file_path=file_path,
            tags=tags,
            flow_address=self._flow_address,
            evm_client=self._evm_client,
            nodes=self._nodes,
            option=option,
        )
