"""Native neural self-data candidates; no generated host payload is executed."""
from pathlib import Path
import json
import os
import random
import shutil
import struct
import subprocess

import pytest
import torch

from lightnest import rv32
from lightnest.learned_patch import encode, load_predictor, make_example
from lightnest.pack_patch import build, export_weights

ROOT = Path(__file__).resolve().parents[1]
CHECKPOINT = Path(os.environ.get('LIGHTNEST_PATCH_CHECKPOINT', ROOT / 'work/learned-patch.pt'))

def run(exe, *args):
    return subprocess.run([str(exe), *map(str, args)], capture_output=True, timeout=20)

def seed_bytes():
    return b''.join(w.to_bytes(4, 'little') for w in [rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'), rv32.add('x3', 'x1', 'x2'), rv32.addi('x10', 'x3', 0), rv32.EBREAK])

@pytest.fixture(scope='module')
def native(tmp_path_factory):
    if not CHECKPOINT.exists():
        pytest.skip('run the bounded learned_patch training experiment first')
    if not shutil.which('gcc'):
        pytest.skip('gcc is required')
    directory = tmp_path_factory.mktemp('native-patch')
    executable = build(CHECKPOINT, directory / 'patch-nest')
    return executable, load_predictor(CHECKPOINT)

def seeded(native, tmp_path, binary):
    raw = tmp_path / 'seed.rv32'
    raw.write_bytes(binary)
    exe = tmp_path / 'seeded'
    result = run(native[0], '--init', raw, '-o', exe)
    assert result.returncode == 0, result.stderr
    return exe

@torch.inference_mode()
def reference(predictor, command, binary):
    c,b=encode(command,binary)
    logits,bits=predictor.model(torch.tensor([c]),torch.tensor([b]))
    probs=logits.softmax(-1)[0]
    pos=int(probs.argmax());confidence=float(probs[pos]);byte=0
    if pos==256:
        return pos,byte,confidence,False,'none'
    byte=sum(int(bits[0,pos,k]>=0)<<k for k in range(8))
    if confidence<predictor.threshold:
        return pos,byte,confidence,False,'uncertain'
    bitp=bits[0,pos].sigmoid()
    confidence=min(confidence,float(torch.maximum(bitp,1-bitp).min()))
    if confidence<predictor.threshold:
        return pos,byte,confidence,False,'uncertain_byte'
    if byte==binary[pos]:
        return pos,byte,confidence,False,'no_change'
    return pos,byte,confidence,True,'unvalidated_candidate'

def test_native_python_parity_fixed_examples(native,tmp_path):
    # Frozen seeds, no checkpoint selection and no repair from expected labels.
    rng=random.Random(984176)
    seen=set()
    for split in ('iid','register','offset','language','combined'):
        for index in range(6):
            e=make_example(rng,split)
            directory=tmp_path/f'{split}-{index}';directory.mkdir()
            exe=seeded(native,directory,e.binary)
            result=run(exe,'--inspect',e.command)
            assert result.returncode==0,result.stderr
            actual=json.loads(result.stdout)
            p,b,c,proposed,status=reference(native[1],e.command,e.binary)
            assert (actual['position'],actual['replacement'],actual['proposed'],actual['status'])==(p,b,proposed,status)
            assert actual['confidence']==pytest.approx(c,abs=2e-5)
            patch=native[1].propose(e.command,e.binary)
            assert actual['proposed']==(patch is not None)
            seen.add(status)
    assert 'none' in seen
    assert 'unvalidated_candidate' in seen

def test_native_child_grandchild_independent_semantics(native,tmp_path):
    original=seed_bytes();parent=seeded(native,tmp_path,original)
    before=parent.read_bytes();child=tmp_path/'child';grandchild=tmp_path/'grandchild'
    assert run(parent,'--propose','change add to sub','-o',child).returncode==0
    assert parent.read_bytes()==before
    candidate=run(child,'--dump-self').stdout
    # Candidate host executable differs at precisely one data byte.
    assert len(child.read_bytes())==len(before)
    assert sum(a!=b for a,b in zip(before,child.read_bytes()))==1
    assert sum(a!=b for a,b in zip(original,candidate))==1
    cases=[(0,0),(3,5),(-12,8),(2147483647,-1),(-2147483648,1),(12345,-6789)]
    signed=lambda n:(n+2**31)%2**32-2**31
    assert all(rv32.simulate(candidate,list(pair))==signed(pair[0]-pair[1]) for pair in cases)
    assert run(child,'--propose','change sub to add','-o',grandchild).returncode==0
    assert grandchild.read_bytes()==before
    restored=run(grandchild,'--dump-self').stdout
    assert restored==original
    assert all(rv32.simulate(restored,list(pair))==signed(pair[0]+pair[1]) for pair in cases)
    # No-op model decision must create no output.
    absent=tmp_path/'absent'
    result=run(child,'--propose','change add to sub','-o',absent)
    assert result.returncode==3
    assert not absent.exists()

def test_input_output_bounds_and_immutability(native,tmp_path):
    exe=seeded(native,tmp_path,seed_bytes());before=exe.read_bytes()
    for request in ('','a'*97):
        assert run(exe,'--inspect',request).returncode==2
    assert run(exe,'--propose','change add to sub','-o',exe).returncode==2
    alias=tmp_path/'alias';os.link(exe,alias)
    assert run(exe,'--propose','change add to sub','-o',alias).returncode==2
    assert exe.read_bytes()==before
    raw=tmp_path/'huge';raw.write_bytes(b'X'*257)
    assert run(native[0],'--init',raw,'-o',tmp_path/'bad').returncode==2
    assert not (tmp_path/'bad').exists()
    assert run(native[0],'--inspect','change add to sub').returncode==2

def test_corrupt_embedded_weights_fail_closed(native,tmp_path):
    encoded=tmp_path/'weights';count=export_weights(CHECKPOINT,encoded)
    assert count==107994
    blob=encoded.read_bytes();image=native[0].read_bytes()
    location=image.find(blob)
    assert location>=0 and image.find(blob,location+1)<0
    for name,offset,value in [('version',8,struct.pack('<I',2)),('count',20,struct.pack('<I',0)),('threshold',24,struct.pack('<f',0)),('nan',32,struct.pack('<f',float('nan')))]:
        corrupt=bytearray(image);corrupt[location+offset:location+offset+4]=value
        exe=tmp_path/name;exe.write_bytes(corrupt);exe.chmod(0o700)
        result=run(exe,'--dump-self')
        assert result.returncode==2
        assert b'weight' in result.stderr or b'confidence' in result.stderr

def test_untrained_weights_abstain_without_fallback(tmp_path):
    # A different (untrained) policy changes behavior; there is no keyword/opcode
    # fallback silently supplying the learned task's known answer.
    from lightnest.learned_patch import PatchNetwork, Predictor
    torch.manual_seed(991)
    model=PatchNetwork().eval()
    checkpoint=tmp_path/'untrained.pt'
    torch.save({'state_dict':model.state_dict(),'threshold':.55},checkpoint)
    exe=build(checkpoint,tmp_path/'untrained')
    native=(exe,Predictor(model))
    parent=seeded(native,tmp_path,seed_bytes())
    checkpoint.unlink()  # All runtime weights are really embedded.
    expected=reference(native[1],'change add to sub',seed_bytes())
    target=tmp_path/'must-not-exist'
    result=run(parent,'--propose','change add to sub','-o',target)
    actual=json.loads(result.stdout)
    assert result.returncode==3 and not target.exists()
    assert actual['status'] in ('uncertain','uncertain_byte','none')
    assert (actual['position'],actual['replacement'],actual['proposed'],actual['status'])==(expected[0],expected[1],expected[3],expected[4])
    assert actual['confidence']==pytest.approx(expected[2],abs=2e-5)

def test_combined_holdout_native_generation_chain(native,tmp_path):
    rng=random.Random(20261016)
    while True:
        e=make_example(rng,'combined')
        if e.present:
            break
    assert e.command=='please replace sub with add'
    assert e.pointer==115 and e.replacement==1
    parent=seeded(native,tmp_path,e.binary);before=parent.read_bytes()
    child=tmp_path/'ood-child';grandchild=tmp_path/'ood-grandchild'
    first=run(parent,'--propose',e.command,'-o',child)
    assert first.returncode==0,first.stderr
    assert json.loads(first.stdout)['position']==115
    second=run(child,'--propose','change add to sub','-o',grandchild)
    assert second.returncode==0,second.stderr
    candidate=run(child,'--dump-self').stdout
    restored=run(grandchild,'--dump-self').stdout
    inputs=random.Random(2026101602)
    signed=lambda n:(n+2**31)%2**32-2**31
    for _ in range(32):
        a,b=[inputs.randint(-2**31,2**31-1) for _ in range(2)]
        assert rv32.simulate(candidate,[a,b])==signed(a+b)
        assert rv32.simulate(restored,[a,b])==signed(a-b)
    assert parent.read_bytes()==before==grandchild.read_bytes()
    assert sum(a!=b for a,b in zip(before,child.read_bytes()))==1

def test_running_inode_used_after_path_replacement(native, tmp_path):
    original = seed_bytes()
    parent = seeded(native, tmp_path, original)
    before = parent.read_bytes()
    fd = os.open(parent, os.O_RDONLY)
    try:
        parent.unlink()
        parent.write_bytes(b'concurrently replaced pathname')
        result = subprocess.run(
            [f'/proc/self/fd/{fd}', '--propose', 'change add to sub', '-o', str(tmp_path / 'child')],
            pass_fds=(fd,), capture_output=True, timeout=20,
        )
    finally:
        os.close(fd)
    assert result.returncode == 0, result.stderr
    child = tmp_path / 'child'
    assert len(child.read_bytes()) == len(before)
    assert sum(a != b for a, b in zip(before, child.read_bytes())) == 1
    assert parent.read_bytes() == b'concurrently replaced pathname'
    assert rv32.simulate(run(child, '--dump-self').stdout, [3, 5]) == -2


def test_fifo_seed_rejected_without_blocking(native, tmp_path):
    fifo = tmp_path / 'seed-fifo'
    os.mkfifo(fifo)
    result = subprocess.run(
        [str(native[0]), '--init', str(fifo), '-o', str(tmp_path / 'child')],
        capture_output=True, timeout=2,
    )
    assert result.returncode == 2
    assert not (tmp_path / 'child').exists()


def test_packer_failure_leaves_no_output(native, tmp_path, monkeypatch):
    from lightnest import pack_patch
    output = tmp_path / 'final'

    def failed_compiler(command, **kwargs):
        staged = Path(command[command.index('-o') + 1])
        assert staged != output
        assert staged.parent.parent == output.parent
        staged.write_bytes(b'partial compiler output')
        raise subprocess.CalledProcessError(1, command)

    monkeypatch.setattr(pack_patch.subprocess, 'run', failed_compiler)
    with pytest.raises(subprocess.CalledProcessError):
        build(CHECKPOINT, output)
    assert not output.exists()
    assert list(tmp_path.iterdir()) == []


def test_packer_publish_race_preserves_existing_output(native, tmp_path, monkeypatch):
    from lightnest import pack_patch
    output = tmp_path / 'final'

    def racing_compiler(command, **kwargs):
        staged = Path(command[command.index('-o') + 1])
        assert staged != output
        staged.write_bytes(b'compiled candidate')
        output.write_bytes(b'concurrently created destination')
        return subprocess.CompletedProcess(command, 0)

    monkeypatch.setattr(pack_patch.subprocess, 'run', racing_compiler)
    with pytest.raises(FileExistsError):
        build(CHECKPOINT, output)
    assert output.read_bytes() == b'concurrently created destination'
    assert list(tmp_path.iterdir()) == [output]


def test_packer_does_not_follow_destination_symlinks(native, tmp_path):
    target = tmp_path / 'absent-target'
    output = tmp_path / 'output'
    output.symlink_to(target)
    with pytest.raises(FileExistsError):
        build(CHECKPOINT, output)
    assert output.is_symlink()
    assert not target.exists()


@pytest.mark.parametrize('threshold', [float('nan'), float('inf'), -float('inf')])
def test_export_rejects_nonfinite_threshold(native, tmp_path, threshold):
    saved = torch.load(CHECKPOINT, map_location='cpu', weights_only=True)
    saved['threshold'] = threshold
    checkpoint = tmp_path / 'invalid.pt'
    torch.save(saved, checkpoint)
    output = tmp_path / 'weights.bin'
    with pytest.raises(ValueError, match='confidence'):
        export_weights(checkpoint, output)
    assert not output.exists()