from __future__ import annotations

import numpy as np
import torch

DEFAULT_HASH_WIDTH = 512


def _validate_width(width: int) -> None:
    if width <= 0:
        raise ValueError("width must be greater than 0")
    if width % 8 != 0:
        raise ValueError("width must be divisible by 8")


def _expected_byte_length(width: int) -> int:
    _validate_width(width)
    return width // 8


def bits_tensor_to_hash_bytes(bits: torch.Tensor, *, width: int = 512) -> bytes:
    _validate_width(width)
    flat = bits.detach().cpu().flatten()
    if flat.numel() != width:
        raise ValueError(
            f"Input tensor must have exactly {width} values, got {flat.numel()}"
        )
    bit_array = (flat.numpy() > 0).astype(np.uint8)
    return np.packbits(bit_array).tobytes()


def hash_bytes_to_bits_array(hash_bytes: bytes, *, width: int = 512) -> np.ndarray:
    expected = _expected_byte_length(width)
    if len(hash_bytes) != expected:
        raise ValueError(
            f"hash_bytes must be exactly {expected} bytes, got {len(hash_bytes)}"
        )
    byte_array = np.frombuffer(hash_bytes, dtype=np.uint8)
    return np.unpackbits(byte_array)[:width].astype(np.uint8, copy=False)


def hash_bytes_to_cam_row(hash_bytes: bytes, *, width: int = 512) -> int:
    expected = _expected_byte_length(width)
    if len(hash_bytes) != expected:
        raise ValueError(
            f"hash_bytes must be exactly {expected} bytes, got {len(hash_bytes)}"
        )
    return int.from_bytes(hash_bytes, byteorder="big", signed=False)


def cam_row_to_hash_bytes(cam_row: int, *, width: int = 512) -> bytes:
    if not (0 <= cam_row < 1 << width):
        raise ValueError(
            f"cam_row {cam_row} is out of range [0, 2**{width})"
        )
    expected = _expected_byte_length(width)
    return int(cam_row).to_bytes(expected, byteorder="big", signed=False)


def bits_tensor_to_cam_row(bits: torch.Tensor, *, width: int = 512) -> int:
    hash_bytes = bits_tensor_to_hash_bytes(bits, width=width)
    return hash_bytes_to_cam_row(hash_bytes, width=width)