"""zg_storage.utils.__init__ module."""

from __future__ import annotations

import base64
import binascii
from typing import List, Union

from web3 import Web3

DEFAULT_CHUNK_SIZE = 256
DEFAULT_SEGMENT_MAX_CHUNKS = 1024
DEFAULT_SEGMENT_SIZE = DEFAULT_CHUNK_SIZE * DEFAULT_SEGMENT_MAX_CHUNKS


def ensure_hex(data: Union[str, bytes]) -> str:
    """Ensure a hex string is 0x-prefixed."""
    if isinstance(data, bytes):
        return "0x" + data.hex()
    if data.startswith("0x"):
        return data
    return "0x" + data


def hex_to_bytes(value: str) -> bytes:
    """Convert a hex string (with or without 0x) to bytes."""
    if value.startswith("0x"):
        value = value[2:]
    return binascii.unhexlify(value)


def decode_bytes(value: Union[str, bytes]) -> bytes:
    """Decode bytes from raw or base64-encoded input."""
    if isinstance(value, bytes):
        return value
    return base64.b64decode(value)


def encode_bytes_base64(data: bytes) -> str:
    return base64.b64encode(data).decode("ascii")


def next_pow2(value: int) -> int:
    """Return the next power-of-two greater than or equal to the value."""
    if value <= 1:
        return 1
    x = value - 1
    x |= x >> 32
    x |= x >> 16
    x |= x >> 8
    x |= x >> 4
    x |= x >> 2
    x |= x >> 1
    return x + 1


def compute_padded_chunks(chunks: int) -> int:
    """Compute chunk count rounded to a shard-friendly padded boundary."""
    next_pow = next_pow2(chunks)
    if next_pow == chunks:
        return next_pow
    min_chunk = next_pow // 16 if next_pow >= 16 else 1
    return ((chunks - 1) // min_chunk + 1) * min_chunk


def build_merkle_root(leaves: List[bytes], hash_leaves: bool = False) -> bytes:
    """Build a Merkle root from leaves, optionally hashing the leaves first."""
    if not leaves:
        return b""
    nodes = [Web3.keccak(l) if hash_leaves else l for l in leaves]
    while len(nodes) > 1:
        next_level: list[bytes] = []
        for i in range(0, len(nodes), 2):
            if i == len(nodes) - 1:
                next_level.append(nodes[i])
            else:
                next_level.append(Web3.keccak(nodes[i] + nodes[i + 1]))
        nodes = next_level
    return nodes[0]


def segment_root(segment_bytes: bytes) -> bytes:
    """Compute the Merkle root for a single segment."""
    leaves: list[bytes] = []
    for i in range(0, len(segment_bytes), DEFAULT_CHUNK_SIZE):
        chunk = segment_bytes[i : i + DEFAULT_CHUNK_SIZE]
        if len(chunk) < DEFAULT_CHUNK_SIZE:
            chunk = chunk + b"\x00" * (DEFAULT_CHUNK_SIZE - len(chunk))
        leaves.append(Web3.keccak(chunk))
    return build_merkle_root(leaves, hash_leaves=False)


def num_splits(total: int, unit: int) -> int:
    return (total - 1) // unit + 1


def compute_data_root(data: bytes, segment_size: int, chunk_size: int) -> bytes:
    """Compute the full data Merkle root from raw bytes."""
    file_size = len(data)
    chunks = num_splits(file_size, chunk_size)
    padded_chunks = compute_padded_chunks(chunks)
    padded_size = padded_chunks * chunk_size
    num_segments = num_splits(padded_size, segment_size)
    segment_roots = []
    for seg_index in range(num_segments):
        offset = seg_index * segment_size
        read_size = min(segment_size, padded_size - offset)
        if offset >= file_size:
            seg_bytes = b"\x00" * read_size
        else:
            chunk = data[offset : offset + read_size]
            if len(chunk) < read_size:
                chunk = chunk + b"\x00" * (read_size - len(chunk))
            seg_bytes = chunk
        segment_roots.append(segment_root(seg_bytes))
    return build_merkle_root(segment_roots, hash_leaves=False)
