"""
SoundBlockMiner — GPU SHA-256 mining orchestrator using sb_miner.wgsl.

Optimizations applied:
  1. Pre-allocated buffers — zero allocation during mining loop
  2. Pre-built bind group — reused across all dispatches
  3. Fused dispatches — multiple dispatches per queue.submit()
  4. Persistent kernel — each thread processes N nonces in shader loop
  5. Shared memory — header + K constants loaded once per workgroup
  6. On-GPU difficulty check — only hits transfer back to CPU

Buffer layout:
  - params (uniform, 16 bytes): base_nonce, nonces_per_thread, difficulty_bits, total_threads
  - header (storage read, 80 bytes): 20 x u32 block header words
  - results (storage rw, 2312 bytes): atomic count + 64 x (nonce + 8 hash words)
  - k_constants (storage read, 256 bytes): 64 x u32 SHA-256 K values

Usage:
    miner = SoundBlockMiner(block_header, difficulty=24)
    result = miner.mine_continuous(duration_seconds=60)
    print(f"Hashrate: {result['hashrate']:.1f} H/s")
"""

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

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

_SHADER_DIR = Path(__file__).parent

# SHA-256 K constants (first 32 bits of cube roots of first 64 primes)
SHA256_K = [
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5,
    0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,
    0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc,
    0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
    0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,
    0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3,
    0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5,
    0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,
    0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
]

# Max results per dispatch (must match shader)
MAX_HITS = 64
RESULTS_SIZE = (1 + MAX_HITS * 9) * 4  # 2312 bytes


def _pack_header_big_endian(header_bytes: bytes) -> bytes:
    """Pack 80-byte header as 20 big-endian u32 words stored as native u32."""
    assert len(header_bytes) == 80, f"Header must be 80 bytes, got {len(header_bytes)}"
    words = []
    for i in range(20):
        word = struct.unpack('>I', header_bytes[i*4:(i+1)*4])[0]
        words.append(word)
    return struct.pack(f'{len(words)}I', *words)


def cpu_double_sha256(header_bytes: bytes, nonce: int) -> str:
    """CPU double-SHA-256 for verification. Returns hex string."""
    data = header_bytes + struct.pack('>I', nonce)
    h1 = hashlib.sha256(data).digest()
    h2 = hashlib.sha256(h1).hexdigest()
    return h2


def count_leading_zero_bits(hex_hash: str) -> int:
    """Count leading zero bits in a hex hash string."""
    bits = bin(int(hex_hash, 16))[2:].zfill(256)
    return len(bits) - len(bits.lstrip('0'))


class SoundBlockMiner:
    """GPU SHA-256 miner using Sound Block architecture.

    Pre-allocates all GPU resources at init. Mining loop only updates
    the params uniform buffer and submits dispatches.
    """

    def __init__(self, block_header: bytes, difficulty: int = 20,
                 nonces_per_thread: int = 32, threads_per_dispatch: int = 262144):
        """
        Args:
            block_header: 80-byte block header (big-endian, standard Bitcoin format)
            difficulty: number of leading zero bits required
            nonces_per_thread: how many nonces each GPU thread processes (inner loop)
            threads_per_dispatch: total threads per dispatch (should be multiple of 256)
        """
        if not _WGPU_AVAILABLE:
            raise ImportError("wgpu package required: pip install wgpu")

        assert len(block_header) == 80, f"Header must be 80 bytes, got {len(block_header)}"

        self.header_bytes = block_header
        self.difficulty = difficulty
        self.nonces_per_thread = nonces_per_thread
        # Round threads to multiple of 256 (workgroup size)
        self.threads_per_dispatch = (threads_per_dispatch + 255) // 256 * 256
        self.nonces_per_dispatch = self.threads_per_dispatch * self.nonces_per_thread

        self._device = wgpu.utils.get_default_device()
        self._load_shader()
        self._create_pipeline()
        self._create_buffers()
        self._upload_static_data()
        self._create_bind_group()

    def _load_shader(self):
        self._shader_src = (_SHADER_DIR / "sb_miner.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: params (uniform)
            {"binding": 0, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.uniform}},
            # binding 1: header (storage read)
            {"binding": 1, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            # binding 2: results (storage rw)
            {"binding": 2, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
            # binding 3: k_constants (storage read)
            {"binding": 3, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
        ])

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

    def _create_buffers(self):
        d = self._device

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

        # Params: 4 x u32 = 16 bytes
        self._buf_params = d.create_buffer(size=16, usage=UNIFORM)

        # Header: 20 x u32 = 80 bytes
        self._buf_header = d.create_buffer(size=80, usage=STORAGE_RO)

        # Results: count + 64 * 9 words = 577 u32s = 2308 bytes, round to 2312
        self._buf_results = d.create_buffer(size=RESULTS_SIZE, usage=STORAGE_RW)

        # K constants: 64 x u32 = 256 bytes
        self._buf_k = d.create_buffer(size=256, usage=STORAGE_RO)

    def _upload_static_data(self):
        """Upload header and K constants (one-time, never changes during mining)."""
        d = self._device

        # Header as big-endian u32 words
        header_data = _pack_header_big_endian(self.header_bytes)
        d.queue.write_buffer(self._buf_header, 0, header_data)

        # K constants
        k_data = struct.pack(f'{len(SHA256_K)}I', *SHA256_K)
        d.queue.write_buffer(self._buf_k, 0, k_data)

    def _create_bind_group(self):
        """Build bind group once — reused for all dispatches."""
        self._bg = self._device.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_params}},
            {"binding": 1, "resource": {"buffer": self._buf_header}},
            {"binding": 2, "resource": {"buffer": self._buf_results}},
            {"binding": 3, "resource": {"buffer": self._buf_k}},
        ])

    def _clear_results(self):
        """Zero out the results buffer (count = 0)."""
        zeros = bytes(RESULTS_SIZE)
        self._device.queue.write_buffer(self._buf_results, 0, zeros)

    def _read_results(self) -> list:
        """Read hits from results buffer. Returns list of (nonce, hash_hex) tuples."""
        data = self._device.queue.read_buffer(self._buf_results, size=RESULTS_SIZE)
        words = struct.unpack(f'{RESULTS_SIZE // 4}I', data)
        count = min(words[0], MAX_HITS)
        hits = []
        for i in range(count):
            base = 1 + i * 9
            nonce = words[base]
            hash_words = words[base+1:base+9]
            hash_hex = ''.join(f'{w:08x}' for w in hash_words)
            hits.append((nonce, hash_hex))
        return hits

    def mine_batch(self, base_nonce: int, num_dispatches: int = 1) -> list:
        """Run mining dispatches and return any hits.

        Fuses multiple dispatches into a single command encoder for
        minimum per-dispatch overhead.

        Args:
            base_nonce: starting nonce
            num_dispatches: number of dispatches to fuse (each covers
                           nonces_per_dispatch nonces)

        Returns:
            List of (nonce, hash_hex) tuples for any hits found
        """
        d = self._device

        # Clear results
        self._clear_results()

        # Build fused command encoder with multiple dispatches
        encoder = d.create_command_encoder()
        workgroups = self.threads_per_dispatch // 256

        for i in range(num_dispatches):
            dispatch_base = base_nonce + i * self.nonces_per_dispatch
            params = struct.pack("IIII",
                                 dispatch_base,
                                 self.nonces_per_thread,
                                 self.difficulty,
                                 self.threads_per_dispatch)
            d.queue.write_buffer(self._buf_params, 0, params)

            compute_pass = encoder.begin_compute_pass()
            compute_pass.set_pipeline(self._pipeline)
            compute_pass.set_bind_group(0, self._bg)
            compute_pass.dispatch_workgroups(workgroups)
            compute_pass.end()

        d.queue.submit([encoder.finish()])
        d._poll_wait()

        return self._read_results()

    def mine_batch_fused(self, base_nonce: int, num_dispatches: int = 10) -> list:
        """Fused mining: multiple dispatches share ONE submit.

        Unlike mine_batch which writes params per-dispatch (last write wins
        for uniform buffers submitted together), this method uses separate
        param buffers per dispatch to truly fuse them.

        Args:
            base_nonce: starting nonce
            num_dispatches: dispatches per submit (each with different nonce range)

        Returns:
            List of (nonce, hash_hex) for any hits
        """
        d = self._device

        # Clear results
        self._clear_results()

        # Create per-dispatch param buffers
        UNIFORM = wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST
        param_buffers = []
        bind_groups = []
        workgroups = self.threads_per_dispatch // 256

        for i in range(num_dispatches):
            dispatch_base = base_nonce + i * self.nonces_per_dispatch
            params = struct.pack("IIII",
                                 dispatch_base,
                                 self.nonces_per_thread,
                                 self.difficulty,
                                 self.threads_per_dispatch)
            buf = d.create_buffer(size=16, usage=UNIFORM)
            d.queue.write_buffer(buf, 0, params)
            param_buffers.append(buf)

            bg = d.create_bind_group(layout=self._bgl, entries=[
                {"binding": 0, "resource": {"buffer": buf}},
                {"binding": 1, "resource": {"buffer": self._buf_header}},
                {"binding": 2, "resource": {"buffer": self._buf_results}},
                {"binding": 3, "resource": {"buffer": self._buf_k}},
            ])
            bind_groups.append(bg)

        # Build single command encoder with all dispatches
        encoder = d.create_command_encoder()
        for i in range(num_dispatches):
            compute_pass = encoder.begin_compute_pass()
            compute_pass.set_pipeline(self._pipeline)
            compute_pass.set_bind_group(0, bind_groups[i])
            compute_pass.dispatch_workgroups(workgroups)
            compute_pass.end()

        d.queue.submit([encoder.finish()])
        d._poll_wait()

        return self._read_results()

    def mine_continuous(self, duration_seconds: int = 60,
                        report_interval: float = 5.0,
                        fused_dispatches: int = 10) -> dict:
        """Run continuous mining for a fixed duration.

        Args:
            duration_seconds: how long to mine
            report_interval: seconds between hashrate reports
            fused_dispatches: dispatches per submit

        Returns:
            dict with total_hashes, elapsed, hashrate, best_difficulty,
                 best_nonce, best_hash, hits
        """
        total_hashes = 0
        current_nonce = 0
        nonces_per_submit = self.nonces_per_dispatch * fused_dispatches
        all_hits = []
        best_difficulty = 0
        best_nonce = 0
        best_hash = ""

        start = time.perf_counter()
        last_report = start

        print(f"Mining started: difficulty={self.difficulty}, "
              f"nonces/dispatch={self.nonces_per_dispatch:,}, "
              f"fused={fused_dispatches}, "
              f"nonces/submit={nonces_per_submit:,}")

        while True:
            now = time.perf_counter()
            elapsed = now - start
            if elapsed >= duration_seconds:
                break

            # Mine a fused batch
            hits = self.mine_batch_fused(current_nonce, fused_dispatches)
            total_hashes += nonces_per_submit
            current_nonce += nonces_per_submit

            # Handle nonce wraparound (u32 max)
            if current_nonce >= 0xFFFFFFFF:
                current_nonce = 0

            # Process hits
            for nonce, hash_hex in hits:
                zeros = count_leading_zero_bits(hash_hex)
                all_hits.append({
                    "nonce": nonce,
                    "hash": hash_hex,
                    "difficulty": zeros,
                    "time": elapsed,
                })
                if zeros > best_difficulty:
                    best_difficulty = zeros
                    best_nonce = nonce
                    best_hash = hash_hex

            # Periodic report
            if now - last_report >= report_interval:
                rate = total_hashes / elapsed
                print(f"  [{elapsed:.1f}s] {total_hashes:>12,} hashes | "
                      f"{rate:>12,.0f} H/s | "
                      f"best: {best_difficulty} bits | "
                      f"nonce: {current_nonce:,}")
                last_report = now

        elapsed = time.perf_counter() - start
        hashrate = total_hashes / elapsed if elapsed > 0 else 0

        return {
            "total_hashes": total_hashes,
            "elapsed_seconds": elapsed,
            "hashrate": hashrate,
            "best_difficulty": best_difficulty,
            "best_nonce": best_nonce,
            "best_hash": best_hash,
            "hits": all_hits,
            "nonces_per_thread": self.nonces_per_thread,
            "threads_per_dispatch": self.threads_per_dispatch,
            "nonces_per_dispatch": self.nonces_per_dispatch,
            "fused_dispatches": fused_dispatches,
        }

    def benchmark(self, warmup_rounds: int = 3, measure_rounds: int = 10,
                  fused_dispatches: int = 10) -> dict:
        """Measure peak hashrate.

        Args:
            warmup_rounds: submit rounds to warm up GPU
            measure_rounds: submit rounds to measure
            fused_dispatches: dispatches per submit

        Returns:
            dict with hashrate, elapsed, config details
        """
        nonces_per_submit = self.nonces_per_dispatch * fused_dispatches

        # Warmup
        for i in range(warmup_rounds):
            self.mine_batch_fused(i * nonces_per_submit, fused_dispatches)

        # Measure
        t0 = time.perf_counter()
        for i in range(measure_rounds):
            base = (warmup_rounds + i) * nonces_per_submit
            self.mine_batch_fused(base, fused_dispatches)
        elapsed = time.perf_counter() - t0

        total = nonces_per_submit * measure_rounds
        hashrate = total / elapsed

        return {
            "hashrate": hashrate,
            "total_hashes": total,
            "elapsed_seconds": elapsed,
            "nonces_per_thread": self.nonces_per_thread,
            "threads_per_dispatch": self.threads_per_dispatch,
            "nonces_per_dispatch": self.nonces_per_dispatch,
            "fused_dispatches": fused_dispatches,
            "nonces_per_submit": nonces_per_submit,
            "adapter_info": dict(self._device.adapter.info),
        }

    def verify_hash(self, nonce: int) -> tuple:
        """CPU-verify a GPU-found hash. Returns (hash_hex, leading_zeros)."""
        h = cpu_double_sha256(self.header_bytes, nonce)
        zeros = count_leading_zero_bits(h)
        return h, zeros

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