import struct
from pathlib import Path
import sys

# Make repo-local imports work when running pytest from project root
sys.path.insert(0, str(Path(__file__).parent.parent))

from kernels.dot_product_code import dot_product_code
from sdk.code_iter import FILL8, COPY, XOR8, ANCHOR, END, CodeIter

# Simple virtual accessors used in KATs
def v_ones(anchor_id: int, anchor_gain: float, virt_off: int, length: int) -> float:
    """Return the dot contribution for a span when every element = 1.0."""
    return float(length)

def v_gain(anchor_id: int, anchor_gain: float, virt_off: int, length: int) -> float:
    """Return the dot contribution for a span when every element = anchor_gain."""
    return float(length) * float(anchor_gain)

def build_bytes(segments):
    """
    segments: list of tuples describing ops.
      Supported forms:
        ("FILL8", len, val)
        ("COPY",  src_off, len)
        ("XOR8",  len, key)
        ("ANCHOR", anchor_id:u32, gain:f32)
        ("END",)
    """
    b = bytearray()
    for seg in segments:
        op = seg[0].upper()
        if op == "FILL8":
            _, ln, val = seg
            b.append(FILL8)
            b += struct.pack("<Q", int(ln))
            b.append(int(val) & 0xFF)
        elif op == "COPY":
            _, src_off, ln = seg
            b.append(COPY)
            b += struct.pack("<Q", int(src_off))
            b += struct.pack("<Q", int(ln))
        elif op == "XOR8":
            _, ln, key = seg
            b.append(XOR8)
            b += struct.pack("<Q", int(ln))
            b.append(int(key) & 0xFF)
        elif op == "ANCHOR":
            _, aid, gain = seg
            b.append(ANCHOR)
            b += struct.pack("<I", int(aid))
            b += struct.pack("<f", float(gain))
        elif op == "END":
            b.append(END)
        else:
            raise ValueError(f"unknown op {op}")
    if not b or b[-1] != END:
        b.append(END)
    return bytes(b)

# ────────────────────────────────────────────────────────────────────────────────
# KAT 1: COPY of prior FILL8 (same anchor) → v_ones accessor
# Expectation: dot = FILL8_len + COPY_len
# ────────────────────────────────────────────────────────────────────────────────
def test_copy_of_fill8_simple():
    prog = build_bytes([
        ("FILL8", 100, 7),       # produce 100 logical bytes
        ("COPY",  0,   100),     # copy first 100 bytes
        ("END",),
    ])
    # Sanity on iterator: no errors
    _ = list(CodeIter(prog))

    # Using v_ones, contribution equals total logical length processed
    res = dot_product_code(prog, v_ones)
    assert res == 200.0

# ────────────────────────────────────────────────────────────────────────────────
# KAT 2: COPY chains collapse (COPY of COPY) → v_ones accessor
# Program: [FILL8 50] [COPY 0,50] [COPY 50,50]
# Expectation: dot = 150, no materialization
# ────────────────────────────────────────────────────────────────────────────────
def test_copy_chain_collapse():
    prog = build_bytes([
        ("FILL8", 50, 3),
        ("COPY",  0,  50),   # copies the first 50
        ("COPY",  50, 50),   # copies the second 50 (which itself was a copy)
        ("END",),
    ])
    _ = list(CodeIter(prog))
    res = dot_product_code(prog, v_ones)
    assert res == 150.0

# ────────────────────────────────────────────────────────────────────────────────
# KAT 3: COPY across anchors (v1 semantics use CURRENT anchor)
# Program:
#   ANCHOR (aid=1, gain=2.0); FILL8 30
#   ANCHOR (aid=2, gain=5.0); COPY src_off=0, len=30
# Expectation under v_gain accessor:
#   dot = 30*2.0  (from original FILL8 under gain=2.0)
#       + 30*5.0  (COPY contributes under CURRENT anchor gain=5.0)
#       = 210.0
# ────────────────────────────────────────────────────────────────────────────────
def test_copy_across_anchors_current_anchor_semantics():
    prog = build_bytes([
        ("ANCHOR", 1, 2.0),
        ("FILL8",  30, 9),
        ("ANCHOR", 2, 5.0),
        ("COPY",   0,  30),
        ("END",),
    ])
    _ = list(CodeIter(prog))
    res = dot_product_code(prog, v_gain)
    assert res == 210.0

# ────────────────────────────────────────────────────────────────────────────────
# KAT 4: FILL8 val-invariance (dot-kernel v1 contract)
# val should NOT affect dot contribution (v_accessor is data-aware)
# ────────────────────────────────────────────────────────────────────────────────
def test_fill8_val_invariant():
    """FILL8 val field should not affect dot (v_accessor yields full contribution)."""
    prog1 = build_bytes([("FILL8", 100, 1), ("END",)])
    prog2 = build_bytes([("FILL8", 100, 9), ("END",)])

    res1 = dot_product_code(prog1, v_ones)
    res2 = dot_product_code(prog2, v_ones)

    # Both should give 100 (length only, val ignored)
    assert res1 == 100.0
    assert res2 == 100.0
    assert res1 == res2

# ────────────────────────────────────────────────────────────────────────────────
# Negative KATs: COPY forward reference and XOR length mismatch
# ────────────────────────────────────────────────────────────────────────────────
def test_copy_forward_reference_rejected():
    """COPY with forward reference must raise ValueError."""
    prog = build_bytes([
        ("FILL8", 50, 7),
        ("COPY", 0, 100),  # Tries to copy [0, 100) but only [0, 50) exists
        ("END",),
    ])

    try:
        dot_product_code(prog, v_ones)
        assert False, "Should have raised ValueError for forward reference"
    except ValueError as e:
        assert "forward reference" in str(e).lower()

def test_copy_xor_length_mismatch_rejected():
    """XOR length must equal COPY length (strict v1)."""
    prog = build_bytes([
        ("FILL8", 100, 7),
        ("XOR8", 50, 0x2A),  # XOR says 50
        ("COPY", 0, 100),    # COPY says 100 (mismatch!)
        ("END",),
    ])

    try:
        _ = list(CodeIter(prog))
        assert False, "Should have raised ValueError for XOR length mismatch"
    except ValueError as e:
        assert "xor length" in str(e).lower() or "!=" in str(e)
