#!/usr/bin/env python3
"""
SB Crypto Benchmark — SHA-256 hashing on GPU via Sound Block shader.

Measures:
  1. GPU SHA-256 hashes per second
  2. Comparison with hashlib (CPU) performance
  3. Sound Block nonce packing (Float512: 12 nonces per bin)
  4. Finding a hash below a target difficulty

Usage:
    cd afs-engine
    python3 -m afs.gpu.sb_crypto_benchmark
"""

import numpy as np
import struct
import time
import hashlib
import sys
from pathlib import Path

try:
    import wgpu
    _WGPU_AVAILABLE = True
except ImportError:
    _WGPU_AVAILABLE = False

_SHADER_DIR = Path(__file__).parent


class SBCryptoKernel:
    """GPU SHA-256 kernel using sb_sha256.wgsl."""

    def __init__(self, max_invocations: int = 65536):
        if not _WGPU_AVAILABLE:
            raise ImportError("wgpu package required: pip install wgpu")

        self.max_invocations = max_invocations
        self._device = wgpu.utils.get_default_device()
        self._load_shader()
        self._create_pipeline()

    def _load_shader(self):
        self._shader_src = (_SHADER_DIR / "sb_sha256.wgsl").read_text()

    def _create_pipeline(self):
        d = self._device
        module = d.create_shader_module(code=self._shader_src)

        self._bgl = d.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.uniform}},
            {"binding": 1, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            {"binding": 2, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
        ])

        pl = d.create_pipeline_layout(bind_group_layouts=[self._bgl])
        self._pipeline = d.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "main"})

    def hash_batch(self, header_bytes: bytes, base_nonce: int,
                   num_invocations: int) -> np.ndarray:
        """
        Compute SHA-256 hashes for nonces [base_nonce, base_nonce + num_invocations).

        Args:
            header_bytes: block header (must be padded to multiple of 4 bytes)
            base_nonce: starting nonce value
            num_invocations: number of hashes to compute

        Returns:
            (num_invocations, 8) array of u32 hash words
        """
        d = self._device
        n = min(num_invocations, self.max_invocations)

        # Pad header to multiple of 4 bytes
        header_padded = header_bytes
        while len(header_padded) % 4 != 0:
            header_padded += b'\x00'
        header_words = len(header_padded) // 4

        # Pack header as big-endian u32 words
        header_u32 = []
        for i in range(header_words):
            word = struct.unpack('>I', header_padded[i*4:(i+1)*4])[0]
            header_u32.append(word)
        header_data = struct.pack(f'{header_words}I', *header_u32)

        STORAGE_RW = (wgpu.BufferUsage.STORAGE |
                      wgpu.BufferUsage.COPY_SRC |
                      wgpu.BufferUsage.COPY_DST)
        STORAGE_RO = (wgpu.BufferUsage.STORAGE |
                      wgpu.BufferUsage.COPY_DST)

        # Params: base_nonce, num_invocations, header_words, _pad
        params = struct.pack("IIII", base_nonce, n, header_words, 0)
        buf_params = d.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
        d.queue.write_buffer(buf_params, 0, params)

        # Header buffer
        buf_header = d.create_buffer(
            size=max(len(header_data), 4), usage=STORAGE_RO)
        d.queue.write_buffer(buf_header, 0, header_data)

        # Output buffer: 8 u32 per invocation
        buf_hashes = d.create_buffer(
            size=n * 8 * 4, usage=STORAGE_RW)

        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": buf_params}},
            {"binding": 1, "resource": {"buffer": buf_header}},
            {"binding": 2, "resource": {"buffer": buf_hashes}},
        ])

        workgroups = (n + 255) // 256
        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._pipeline)
        compute_pass.set_bind_group(0, bg)
        compute_pass.dispatch_workgroups(workgroups)
        compute_pass.end()
        d.queue.submit([encoder.finish()])
        d._poll_wait()

        hash_data = d.queue.read_buffer(buf_hashes, size=n * 8 * 4)
        return np.frombuffer(hash_data, dtype=np.uint32).reshape(n, 8).copy()

    @property
    def adapter_info(self) -> dict:
        return dict(self._device.adapter.info)


def hash_to_hex(hash_words: np.ndarray) -> str:
    """Convert 8 u32 hash words to hex string."""
    return ''.join(f'{w:08x}' for w in hash_words)


def cpu_sha256(header_bytes: bytes, nonce: int) -> str:
    """CPU SHA-256 for comparison."""
    data = header_bytes + struct.pack('>I', nonce)
    return hashlib.sha256(data).hexdigest()


def count_leading_zeros(hash_words: np.ndarray) -> int:
    """Count leading zero bits in a hash (for difficulty comparison)."""
    zeros = 0
    for word in hash_words:
        if word == 0:
            zeros += 32
        else:
            # Count leading zeros of this u32
            w = int(word)
            for bit in range(31, -1, -1):
                if w & (1 << bit):
                    return zeros
                zeros += 1
    return zeros


if __name__ == "__main__":
    if not _WGPU_AVAILABLE:
        print("ERROR: wgpu not available. Install with: pip install wgpu")
        sys.exit(1)

    print("=" * 70)
    print("SB Crypto Benchmark — GPU SHA-256 via Sound Block Shader")
    print("=" * 70)

    try:
        kernel = SBCryptoKernel(max_invocations=65536)
        info = kernel.adapter_info
        print(f"GPU: {info.get('description', 'unknown')}")
    except Exception as e:
        print(f"GPU init failed: {e}")
        sys.exit(1)

    # Test header (80 bytes of deterministic data)
    header = bytes(range(80))

    # ── 1. GPU SHA-256 throughput ──────────────────────────────────────────

    print("\n--- GPU SHA-256 Throughput ---")

    batch_sizes = [256, 1024, 4096, 16384, 65536]

    for batch_size in batch_sizes:
        # Warmup
        _ = kernel.hash_batch(header, 0, min(batch_size, 256))

        t0 = time.perf_counter()
        hashes = kernel.hash_batch(header, 0, batch_size)
        elapsed = time.perf_counter() - t0

        rate = batch_size / elapsed
        print(f"  Batch {batch_size:>6}: {elapsed*1000:.1f}ms = {rate:,.0f} hashes/sec")

    # ── 2. CPU SHA-256 throughput ──────────────────────────────────────────

    print("\n--- CPU SHA-256 Throughput (hashlib) ---")

    cpu_count = 10000
    t0 = time.perf_counter()
    for i in range(cpu_count):
        _ = hashlib.sha256(header + struct.pack('>I', i)).digest()
    cpu_elapsed = time.perf_counter() - t0
    cpu_rate = cpu_count / cpu_elapsed
    print(f"  {cpu_count} hashes: {cpu_elapsed*1000:.1f}ms = {cpu_rate:,.0f} hashes/sec")

    # ── 3. GPU vs CPU comparison ──────────────────────────────────────────

    print("\n--- GPU vs CPU ---")
    gpu_best_batch = 65536
    t0 = time.perf_counter()
    _ = kernel.hash_batch(header, 0, gpu_best_batch)
    gpu_elapsed = time.perf_counter() - t0
    gpu_rate = gpu_best_batch / gpu_elapsed

    speedup = gpu_rate / cpu_rate if cpu_rate > 0 else 0
    print(f"  CPU: {cpu_rate:,.0f} hashes/sec")
    print(f"  GPU: {gpu_rate:,.0f} hashes/sec")
    print(f"  Speedup: {speedup:.1f}x {'(GPU faster)' if speedup > 1 else '(CPU faster)'}")

    # ── 4. Sound Block nonce packing ─────────────────────────────────────

    print("\n--- Sound Block Nonce Packing ---")
    print("  Float512 = 12 f32s per bin = 12 * 3 bytes = 36 bytes per bin")
    print("  Each nonce = 4 bytes → 9 nonces per Float512 bin")
    print("  (36 / 4 = 9 nonces packed in mantissa bytes)")

    nonces_per_bin = 36 // 4  # 9 nonces per Float512 bin
    bins_for_65k = 65536 // nonces_per_bin + 1
    print(f"  65536 nonces packed into {bins_for_65k} Float512 bins")
    print(f"  Container size: {bins_for_65k * 12 * 4} bytes ({bins_for_65k * 12 * 4 / 1024:.1f} KB)")

    # ── 5. Find hash below target difficulty ─────────────────────────────

    print("\n--- Difficulty Search ---")
    target_zeros = 16  # Find a hash with >= 16 leading zero bits

    total_searched = 0
    found = False
    best_zeros = 0
    best_nonce = 0
    best_hash = None

    search_start = time.perf_counter()
    batch = 65536

    for batch_idx in range(100):  # Max 100 batches = 6.5M nonces
        base = batch_idx * batch
        hashes = kernel.hash_batch(header, base, batch)
        total_searched += batch

        for i in range(batch):
            z = count_leading_zeros(hashes[i])
            if z > best_zeros:
                best_zeros = z
                best_nonce = base + i
                best_hash = hashes[i].copy()
            if z >= target_zeros:
                found = True
                break
        if found:
            break

    search_time = time.perf_counter() - search_start
    search_rate = total_searched / search_time

    print(f"  Target: {target_zeros} leading zero bits")
    print(f"  Searched: {total_searched:,} nonces in {search_time*1000:.1f}ms")
    print(f"  Rate: {search_rate:,.0f} nonces/sec")
    if best_hash is not None:
        print(f"  Best: nonce={best_nonce}, {best_zeros} leading zeros")
        print(f"  Hash: {hash_to_hex(best_hash)}")
    if found:
        print(f"  FOUND hash meeting difficulty {target_zeros}!")
    else:
        print(f"  Not found in {total_searched:,} nonces (best: {best_zeros} zeros)")

    # ── 6. Honest assessment ──────────────────────────────────────────────

    print("\n--- Honest Assessment ---")
    print(f"  GPU compute SHA-256 rate: {gpu_rate:,.0f} H/s")
    print(f"  For reference, a typical GPU miner (ethash/SHA) does 30-100 MH/s")
    print(f"  This is a proof-of-concept, not a competitive miner.")
    print(f"  The value is: nonces organized as Sound Block frequency bins")
    print(f"  can be searched, sorted, and branched using the SB instruction set.")
    print(f"  Mining is not the point — compute substrate is the point.")
