#!/usr/bin/env python3
"""
isa_trace_wrapper.py - ISA Tracing Wrapper for Helix64 Executor

Wraps helix64_exec to capture ISA traces for receipt generation.

Usage:
    from helix_cdc.receipts.isa_trace_wrapper import TracingExecutor

    executor = TracingExecutor()
    result = executor.execute(bytecode, fuel=1000)

    # Get ISA trace for receipt
    isa_trace = executor.get_isa_trace()
"""

from __future__ import annotations
import sys
import struct
import hashlib
from pathlib import Path
from dataclasses import dataclass, field
from typing import Dict, List, Any, Optional
import numpy as np

# Add scripts to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))

try:
    from helix64_opcodes import OP, OPNAMES
except ImportError:
    # Fallback opcode table
    OP = {
        "NOP": 0x00, "LOAD": 0x01, "STORE": 0x02, "ADD": 0x03,
        "MUL": 0x04, "MAD": 0x05, "DOT": 0x06, "PACK8": 0x30,
        "UNPACK8": 0x31, "CDR_INIT": 0xC0, "CDR_ABSORB": 0xC1,
        "CDR_STEP_BACK": 0xC2, "CDR_SQUEEZE": 0xC3, "CDR_ABSORB_MEM": 0xC4,
        "CDR_STEP_BACK_MEM": 0xC5, "CDR_ABSORB_MEMB": 0xC6,
        "CDR_STEP_BACK_MEMB": 0xC7, "HALT": 0xFF,
    }
    OPNAMES = {v: k for k, v in OP.items()}

from .inference_receipt_v1 import ISAOp, ISATrace


# Cost model (cycles, flops, bytes)
COST_MODEL = {
    OP["NOP"]:     {"cycles": 1,   "flops": 0,    "bytes": 0},
    OP["LOAD"]:    {"cycles": 3,   "flops": 0,    "bytes": 8},
    OP["STORE"]:   {"cycles": 3,   "flops": 0,    "bytes": 8},
    OP["ADD"]:     {"cycles": 2,   "flops": 1,    "bytes": 24},
    OP["MUL"]:     {"cycles": 2,   "flops": 1,    "bytes": 24},
    OP["MAD"]:     {"cycles": 3,   "flops": 2,    "bytes": 32},
    OP["DOT"]:     {"cycles": 10,  "flops": 0,    "bytes": 0},  # Per-element costs added dynamically
    OP["PACK8"]:   {"cycles": 5,   "flops": 0,    "bytes": 16},
    OP["UNPACK8"]: {"cycles": 5,   "flops": 0,    "bytes": 16},
    OP["CDR_INIT"]:        {"cycles": 50,  "flops": 0, "bytes": 200},
    OP["CDR_ABSORB"]:      {"cycles": 100, "flops": 0, "bytes": 64},
    OP["CDR_STEP_BACK"]:   {"cycles": 100, "flops": 0, "bytes": 64},
    OP["CDR_SQUEEZE"]:     {"cycles": 10,  "flops": 0, "bytes": 64},
    OP["CDR_ABSORB_MEM"]:  {"cycles": 100, "flops": 0, "bytes": 64},
    OP["CDR_STEP_BACK_MEM"]: {"cycles": 100, "flops": 0, "bytes": 64},
    OP["CDR_ABSORB_MEMB"]:   {"cycles": 100, "flops": 0, "bytes": 64},
    OP["CDR_STEP_BACK_MEMB"]: {"cycles": 100, "flops": 0, "bytes": 64},
    OP["HALT"]:    {"cycles": 1,   "flops": 0,    "bytes": 0},
}


@dataclass
class TracedOp:
    """A traced operation with full context."""
    pc: int              # Program counter before execution
    opcode: int          # Raw opcode byte
    op_name: str         # Human-readable name
    args: List[int]      # Raw operand bytes
    cycles: int          # Cost in cycles
    flops: int           # Cost in FLOPs
    bytes_accessed: int  # Memory bytes accessed


class TracingExecutor:
    """
    Helix64 executor with ISA tracing for receipt generation.

    Wraps the real executor and captures a trace of all operations.
    """

    def __init__(self, trace_enabled: bool = True):
        """
        Initialize tracing executor.

        Args:
            trace_enabled: Whether to capture operation traces
        """
        self.trace_enabled = trace_enabled
        self.traced_ops: List[TracedOp] = []
        self.last_result: Optional[Dict[str, Any]] = None

    def execute(
        self,
        bytecode: bytes,
        fuel: int = 1_000_000,
        input_seed: int = 0,
        init_registers: Optional[np.ndarray] = None,
    ) -> Dict[str, Any]:
        """
        Execute bytecode with optional tracing.

        Args:
            bytecode: Helix64 bytecode
            fuel: Instruction limit
            input_seed: Seed for deterministic inputs
            init_registers: Optional initial register state

        Returns:
            Execution result dict
        """
        self.traced_ops = []

        if self.trace_enabled:
            result = self._execute_traced(bytecode, fuel, input_seed, init_registers)
        else:
            # Use real executor without tracing
            try:
                from helix64_exec import execute_real
                result = execute_real(bytecode, fuel, input_seed, init_registers)
            except ImportError:
                result = self._execute_traced(bytecode, fuel, input_seed, init_registers)

        self.last_result = result
        return result

    def _execute_traced(
        self,
        bytecode: bytes,
        fuel: int,
        input_seed: int,
        init_registers: Optional[np.ndarray],
    ) -> Dict[str, Any]:
        """Execute with full tracing."""
        # Initialize state
        R = np.zeros(256, dtype=np.float64) if init_registers is None else init_registers.astype(np.float64).copy()
        MEM = np.zeros(1024, dtype=np.float64)

        pc = 0
        steps = 0
        halt = False
        total_cycles = 0
        total_flops = 0
        total_bytes = 0

        while pc < len(bytecode) and steps < fuel and not halt:
            if pc + 3 >= len(bytecode):
                break

            op = bytecode[pc]
            a = bytecode[pc + 1]
            b = bytecode[pc + 2]
            c = bytecode[pc + 3]

            # Get cost from model
            cost = COST_MODEL.get(op, {"cycles": 1, "flops": 0, "bytes": 0})
            op_cycles = cost["cycles"]
            op_flops = cost["flops"]
            op_bytes = cost["bytes"]

            # Handle DOT special case (variable cost)
            extra_args = []
            if op == OP["DOT"] and pc + 5 < len(bytecode):
                L = struct.unpack("<H", bytecode[pc + 4:pc + 6])[0]
                extra_args = [L]
                op_cycles += 2 * L
                op_flops = L
                op_bytes = L * 16

            # Record trace
            if self.trace_enabled:
                self.traced_ops.append(TracedOp(
                    pc=pc,
                    opcode=op,
                    op_name=OPNAMES.get(op, f"UNK_{op:02X}"),
                    args=[a, b, c] + extra_args,
                    cycles=op_cycles,
                    flops=op_flops,
                    bytes_accessed=op_bytes,
                ))

            # Execute operation (simplified for tracing)
            if op == OP["ADD"]:
                R[a] = R[b] + R[c]
            elif op == OP["MUL"]:
                R[a] = R[b] * R[c]
            elif op == OP["LOAD"]:
                if b < len(MEM):
                    R[a] = MEM[b]
            elif op == OP["STORE"]:
                if a < len(MEM):
                    MEM[a] = R[b]
            elif op == OP["DOT"]:
                if pc + 5 < len(bytecode):
                    L = struct.unpack("<H", bytecode[pc + 4:pc + 6])[0]
                    pc += 2
                    acc = 0.0
                    for i in range(min(L, 256)):
                        acc += R[(b + i) & 0xFF] * R[(c + i) & 0xFF]
                    R[a] = acc
            elif op == OP["HALT"]:
                halt = True

            pc += 4
            steps += 1
            total_cycles += op_cycles
            total_flops += op_flops
            total_bytes += op_bytes

        return {
            "registers_sha256": hashlib.sha256(R.tobytes()).hexdigest(),
            "mem_sha256": hashlib.sha256(MEM.tobytes()).hexdigest(),
            "steps": steps,
            "fuel_exhausted": steps >= fuel and not halt,
            "halted": halt,
            "registers": R.copy(),
            "metrics": {
                "cycles": int(total_cycles),
                "flops": int(total_flops),
                "bytes": int(total_bytes),
            },
        }

    def get_isa_trace(self) -> ISATrace:
        """
        Get ISA trace for receipt generation.

        Returns:
            ISATrace dataclass ready for receipt
        """
        ops = [
            ISAOp(
                op=traced.op_name,
                op_hex=f"0x{traced.opcode:02X}",
                args=traced.args,
                cost_cycles=traced.cycles,
                cost_flops=traced.flops,
                cost_bytes=traced.bytes_accessed,
            )
            for traced in self.traced_ops
        ]

        return ISATrace(ops_executed=ops)

    def get_ops_as_dicts(self) -> List[Dict[str, Any]]:
        """
        Get traced operations as list of dicts (for emitter convenience).

        Returns:
            List of operation dicts
        """
        return [
            {
                "op": traced.op_name,
                "op_hex": f"0x{traced.opcode:02X}",
                "args": traced.args,
                "cost_cycles": traced.cycles,
                "cost_flops": traced.flops,
                "cost_bytes": traced.bytes_accessed,
            }
            for traced in self.traced_ops
        ]

    def get_summary(self) -> Dict[str, int]:
        """Get aggregated trace summary."""
        return {
            "total_ops": len(self.traced_ops),
            "total_cycles": sum(op.cycles for op in self.traced_ops),
            "total_flops": sum(op.flops for op in self.traced_ops),
            "total_bytes": sum(op.bytes_accessed for op in self.traced_ops),
        }


# =============================================================================
# DEMO
# =============================================================================

if __name__ == "__main__":
    print("TracingExecutor Demo")
    print("=" * 60)

    # Create simple bytecode: ADD R0 = R1 + R2; MUL R3 = R0 * R2; HALT
    bytecode = bytes([
        OP["ADD"], 0, 1, 2,   # R0 = R1 + R2
        OP["MUL"], 3, 0, 2,   # R3 = R0 * R2
        OP["HALT"], 0, 0, 0,  # Halt
    ])

    executor = TracingExecutor(trace_enabled=True)

    # Set initial registers
    init_regs = np.zeros(256, dtype=np.float64)
    init_regs[1] = 5.0
    init_regs[2] = 3.0

    result = executor.execute(bytecode, init_registers=init_regs)

    print("\n1. Execution Result:")
    print(f"   Steps: {result['steps']}")
    print(f"   Halted: {result['halted']}")
    print(f"   R0 (5+3): {result['registers'][0]}")
    print(f"   R3 (8*3): {result['registers'][3]}")

    print("\n2. Traced Operations:")
    for op in executor.traced_ops:
        print(f"   PC={op.pc}: {op.op_name} {op.args} (cycles={op.cycles})")

    print("\n3. ISA Trace for Receipt:")
    isa_trace = executor.get_isa_trace()
    print(f"   Total ops: {isa_trace.ops_summary['total_ops']}")
    print(f"   Total cycles: {isa_trace.ops_summary['total_cycles']}")

    print("\n4. Summary:")
    summary = executor.get_summary()
    for k, v in summary.items():
        print(f"   {k}: {v}")
