"""
Sound Block BIOS — Basic Input/Output System for the Sound Block VM.

The BIOS is the first thing that runs. Python bootstrap calls BIOS.post(),
BIOS initializes the virtual hardware, establishes the wide memory pipe,
and hands off to the Level 2 OS.

The pipe between BIOS and OS is the ENTIRE GPU VRAM mapped as shared
address space. Not 5KB of bridge — the full memory bus width.

Boot sequence:
  1. POST: test GPU device, verify Fabric primitives
  2. Hardware enumeration: detect VRAM, cores, compute limits
  3. Memory map: allocate VRAM as spectral address space
  4. Region table: dynamic regions sized to hardware capacity
  5. Clock init: configure dispatch cadence and clock domains
  6. Boot: load Level 2 program, dispatch first shader

Usage (the entire Python bootstrap):
    from afs.gpu.sb_bios import SoundBlockBIOS
    bios = SoundBlockBIOS()
    bios.post()       # test hardware
    bios.boot()       # hand off to Level 2 VM — Python is done
"""

import numpy as np
import struct
import time
from typing import Optional


class SoundBlockBIOS:
    """Sound Block BIOS — initializes virtual hardware and wide memory pipe.

    Maps the entire available VRAM as the shared spectral address space.
    Every byte of VRAM is addressable by every layer (CPU, FPGA, ASIC, Colony, GPU).
    """

    # Memory map — dynamic, sized at POST based on available VRAM
    # These are RATIOS of total VRAM, not fixed sizes
    REGION_RATIOS = {
        'cpu_regs':       0.000001,   # 16 words (tiny, fixed)
        'cpu_imem':       0.001,      # instruction memory (~8KB at 8GB)
        'cpu_dmem':       0.01,       # data memory (~80KB at 8GB)
        'fpga_state':     0.05,       # FPGA block states (~400KB at 8GB)
        'fpga_routing':   0.05,       # routing matrices
        'asic_state':     0.001,      # hash state
        'colony_probes':  0.10,       # colony probe amplitudes
        'colony_matrix':  0.05,       # influence matrix
        'colony_results': 0.01,       # strategy outputs
        'bridge_io':      0.05,       # cross-layer I/O
        'gpu_workspace':  0.50,       # GPU compute workspace (largest)
        'cluster_sync':   0.01,       # inter-device synchronization
        'os_heap':        0.17,       # Level 2 OS dynamic allocation
    }

    def __init__(self):
        self.fabric = None
        self.memory_map = {}
        self.total_vram = 0
        self.total_words = 0
        self.hardware = {}
        self.booted = False
        self._post_results = {}

    # ─── POST: Power-On Self-Test ────────────────────────────────────────────

    def post(self) -> dict:
        """Power-On Self-Test. Verify GPU, Fabric primitives, measure capacity.

        This is the FIRST thing the Sound Block VM does.
        If POST fails, the VM cannot boot.
        """
        from .sb_fabric import SoundBlockFabric, gpu_available
        t0 = time.perf_counter()

        results = {
            'gpu_available': False,
            'fabric_ok': False,
            'primitives_tested': 0,
            'vram_bytes': 0,
            'gpu_cores': 0,
            'adapter': '',
            'errors': [],
        }

        # Test 1: GPU available?
        if not gpu_available():
            results['errors'].append('No GPU available')
            self._post_results = results
            return results
        results['gpu_available'] = True

        # Test 2: Create Fabric
        try:
            self.fabric = SoundBlockFabric(max_size=1024)
            results['fabric_ok'] = True
        except Exception as e:
            results['errors'].append(f'Fabric creation failed: {e}')
            self._post_results = results
            return results

        # Test 3: Hardware enumeration
        adapter_info = self.fabric.adapter_info
        results['adapter'] = adapter_info.get('description', 'unknown')

        # Detect VRAM (estimate from max buffer size)
        # wgpu doesn't expose total VRAM directly, so we use max_size
        device = self.fabric._device
        limits = device.limits if hasattr(device, 'limits') else {}

        # Estimate available VRAM from buffer limits
        max_buffer = getattr(limits, 'max_buffer_size', 256 * 1024 * 1024)  # default 256MB
        max_storage = getattr(limits, 'max_storage_buffer_binding_size', 128 * 1024 * 1024)
        self.total_vram = min(int(max_buffer), int(max_storage))
        self.total_words = self.total_vram // 4

        results['vram_bytes'] = self.total_vram
        results['vram_mb'] = self.total_vram / (1024 * 1024)

        self.hardware = {
            'adapter': results['adapter'],
            'vram_bytes': self.total_vram,
            'vram_mb': results['vram_mb'],
        }

        # Test 4: Primitive self-test (quick — one dispatch each)
        primitives_ok = 0
        try:
            # Test map
            data = np.array([1.0, -1.0, 0.5], dtype=np.float32)
            result = self.fabric.map(data, 'abs')
            if np.allclose(result, [1.0, 1.0, 0.5]):
                primitives_ok += 1

            # Test add
            a = np.array([1.0, 2.0], dtype=np.float32)
            b = np.array([3.0, 4.0], dtype=np.float32)
            self.fabric.upload('a', a)
            self.fabric.upload('b', b)
            self.fabric._slot_sizes['a'] = 2
            self.fabric._slot_sizes['b'] = 2
            self.fabric.add_gpu('a', 'b', 'out')
            r = self.fabric.download('out')[:2]
            if np.allclose(r, [4.0, 6.0]):
                primitives_ok += 1

            # Test XOR
            self.fabric.xor_gpu('a', 'b', 'out')
            primitives_ok += 1  # if no error, pass

        except Exception as e:
            results['errors'].append(f'Primitive test failed: {e}')

        results['primitives_tested'] = primitives_ok
        results['post_ms'] = (time.perf_counter() - t0) * 1000

        self._post_results = results
        return results

    # ─── Memory Map ──────────────────────────────────────────────────────────

    def create_memory_map(self, vram_budget: int = None) -> dict:
        """Create the VRAM memory map — the wide pipe between BIOS and OS.

        Maps the entire available VRAM as shared spectral address space.
        Each region sized proportionally to hardware capacity.

        Args:
            vram_budget: override VRAM budget in bytes (default: detected)

        Returns:
            {region_name: {offset, size, size_kb}}
        """
        budget = vram_budget or self.total_vram
        if budget <= 0:
            budget = 8 * 1024 * 1024  # fallback: 8MB

        budget_words = budget // 4

        # Allocate regions proportionally
        offset = 0
        self.memory_map = {}
        for name, ratio in self.REGION_RATIOS.items():
            size = max(16, int(budget_words * ratio))  # minimum 16 words
            # Align to 256 words (1KB) for GPU buffer alignment
            size = ((size + 255) // 256) * 256
            self.memory_map[name] = {
                'offset': offset,
                'size': size,
                'size_kb': size * 4 / 1024,
            }
            offset += size

        self.total_words = offset
        return self.memory_map

    # ─── Boot ────────────────────────────────────────────────────────────────

    def boot(self, program: list = None, os_config: dict = None) -> dict:
        """Boot the Sound Block VM.

        POST must pass first. Creates memory map, initializes all regions,
        optionally loads a Level 2 program.

        After boot(), Python's job is done. The VM runs.

        Returns:
            Boot result with memory map and hardware info
        """
        if not self._post_results.get('fabric_ok'):
            raise RuntimeError('POST failed. Cannot boot.')

        t0 = time.perf_counter()

        # Create memory map (the wide pipe)
        self.create_memory_map()

        # Re-create Fabric with appropriate size for the memory map
        max_size = min(self.total_words, 1_000_000)  # cap at Fabric limit
        from .sb_fabric import SoundBlockFabric
        self.fabric = SoundBlockFabric(max_size=max_size)

        # Initialize bridge with full memory map
        from .sb_spectral_bridge import SpectralBridge
        self.bridge = SpectralBridge(self.fabric)

        # Load Level 2 program if provided
        if program:
            self.bridge.cpu_load(program)

        boot_ms = (time.perf_counter() - t0) * 1000
        self.booted = True

        return {
            'booted': True,
            'memory_map': self.memory_map,
            'total_vram_mb': self.total_vram / (1024 * 1024),
            'total_mapped_words': self.total_words,
            'total_mapped_kb': self.total_words * 4 / 1024,
            'regions': len(self.memory_map),
            'hardware': self.hardware,
            'boot_ms': boot_ms,
            'post': self._post_results,
        }

    # ─── Info ────────────────────────────────────────────────────────────────

    @property
    def is_booted(self):
        return self.booted

    def region_info(self, name: str) -> dict:
        """Get info about a memory region."""
        return self.memory_map.get(name, {})

    def print_memory_map(self):
        """Print the full memory map."""
        lines = [f"Sound Block BIOS Memory Map ({self.total_words * 4 / 1024:.0f} KB total):"]
        for name, info in self.memory_map.items():
            lines.append(f"  {name:>20}: offset={info['offset']:>8}, "
                        f"size={info['size']:>8} words ({info['size_kb']:.1f} KB)")
        return '\n'.join(lines)
