"""zg_storage.core.flow module."""

from __future__ import annotations

from contextlib import nullcontext
from typing import List

from web3 import Web3

from ..contracts.types import Submission, SubmissionData, SubmissionNode
from ..core.data import DataSource, FileDataSource
from ..utils import (
    DEFAULT_CHUNK_SIZE,
    DEFAULT_SEGMENT_MAX_CHUNKS,
    build_merkle_root,
    compute_padded_chunks,
    next_pow2,
    segment_root,
)


class FlowBuilder:
    """FlowBuilder class.

    Purpose:
        SDK component type."""

    def __init__(self, data: DataSource | str, tags: bytes) -> None:
        if isinstance(data, str):
            self.data: DataSource = FileDataSource(data)
            self._data_ctx = self.data
        else:
            self.data = data
            self._data_ctx = None
        self.tags = tags

    @staticmethod
    def _num_splits(total: int, unit: int) -> int:
        return (total - 1) // unit + 1

    @staticmethod
    def _next_pow2(value: int) -> int:
        """Return the next power-of-two greater than or equal to the value."""
        return next_pow2(value)

    @staticmethod
    def _compute_padded_chunks(chunks: int) -> int:
        """Compute padded chunk count used for flow submission."""
        return compute_padded_chunks(chunks)

    def _read_at(self, read_size: int, offset: int, padded_size: int, file_size: int) -> bytes:
        """Read data at an offset and pad to the expected size."""
        if offset < 0 or offset >= padded_size:
            raise ValueError("invalid offset")
        max_available = padded_size - offset
        expected = read_size if max_available >= read_size else max_available
        if offset >= file_size:
            return b"\x00" * expected
        data = self.data.read_at(offset, expected)
        if len(data) < expected:
            data = data + b"\x00" * (expected - len(data))
        return data

    def _split_nodes(self, chunks: int) -> List[int]:
        """Split chunk counts into flow submission node sizes."""
        padded_chunks = self._compute_padded_chunks(chunks)
        next_chunk_size = self._next_pow2(chunks)
        nodes: List[int] = []
        while padded_chunks > 0:
            if padded_chunks >= next_chunk_size:
                padded_chunks -= next_chunk_size
                nodes.append(next_chunk_size)
            next_chunk_size //= 2
        return nodes

    @staticmethod
    def _height_for_chunks(chunks: int) -> int:
        return chunks.bit_length() - 1

    def create_submission(self, submitter: str) -> Submission:
        """Build a flow submission for the current data source."""
        ctx_mgr = self._data_ctx if self._data_ctx is not None else nullcontext()
        with ctx_mgr:
            file_size = self.data.size()
            if file_size <= 0:
                raise ValueError("file is empty")

            chunks = self._num_splits(file_size, DEFAULT_CHUNK_SIZE)
            padded_chunks = self._compute_padded_chunks(chunks)
            padded_size = padded_chunks * DEFAULT_CHUNK_SIZE

            nodes: List[SubmissionNode] = []
            offset = 0
            for node_chunks in self._split_nodes(chunks):
                batch_chunks = min(node_chunks, DEFAULT_SEGMENT_MAX_CHUNKS)
                batch_size = batch_chunks * DEFAULT_CHUNK_SIZE
                total_size = node_chunks * DEFAULT_CHUNK_SIZE
                num_tasks = self._num_splits(total_size, batch_size)

                leaves: List[bytes] = []
                for task in range(num_tasks):
                    buf = self._read_at(
                        batch_size,
                        offset + task * batch_size,
                        padded_size,
                        file_size,
                    )
                    leaves.append(segment_root(buf))

                root = build_merkle_root(leaves, hash_leaves=False)
                height = self._height_for_chunks(node_chunks)
                nodes.append(SubmissionNode(root=root, height=height))
                offset += total_size

        submission_data = SubmissionData(length=file_size, tags=self.tags, nodes=nodes)
        return Submission(data=submission_data, submitter=submitter)


def submission_root(submission: Submission) -> bytes:
    """Compute the Merkle root of a submission."""
    nodes = submission.data.nodes
    if not nodes:
        return b""
    root = nodes[-1].root
    for i in range(1, len(nodes)):
        left = nodes[-1 - i].root
        root = Web3.keccak(left + root)
    return root


def submission_fee(submission: Submission, price_per_sector: int) -> int:
    """Compute the total fee based on submission size and sector price."""
    sectors = 0
    for node in submission.data.nodes:
        sectors += 1 << node.height
    return sectors * price_per_sector
