"""Fast contract checks; long learned lineage experiments are saved separately."""
import hashlib
import json
from pathlib import Path
import subprocess

import pytest
import torch

from lightnest.learned_patch import PatchNetwork, encode
from lightnest.pack_recursive import build


@pytest.fixture(scope="module")
def native(tmp_path_factory):
    directory = tmp_path_factory.mktemp("recursive")
    torch.set_num_threads(2)
    torch.manual_seed(716031)
    model = PatchNetwork().eval()
    checkpoint = directory / "model.pt"
    torch.save({"state_dict": model.state_dict(), "threshold": .55}, checkpoint)
    executable = build(checkpoint, directory / "parent")
    return executable, checkpoint, model


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


def test_proposals_match_python(native, tmp_path):
    executable, _, model = native
    raw = bytes([3, 7, 10, 19, 31])
    path = tmp_path / "window"
    path.write_bytes(raw)
    result = run(executable, "--propose", "change add to sub", path)
    assert result.returncode == 0, result.stderr
    actual = json.loads(result.stdout)
    command, binary = encode("change add to sub", raw)
    with torch.no_grad():
        scores, bits = model(torch.tensor([command]), torch.tensor([binary]))
    position = int(scores[0].argmax())
    byte = 0 if position == 256 else sum(int(bits[0, position, k] >= 0) << k for k in range(8))
    assert (actual["position"], actual["byte"]) == (position, byte)
    assert actual["pointer_confidence"] == pytest.approx(float(scores[0].softmax(0)[position]), abs=2e-6)


@pytest.mark.parametrize("generations,updates,reason", [(0, 256, "generation_budget"), (3, 0, "zero_updates")])
def test_budget_stops_without_child_or_parent_change(native, tmp_path, generations, updates, reason):
    executable, _, _ = native
    before = hashlib.sha256(executable.read_bytes()).digest()
    directory = tmp_path / "children"
    result = run(executable, "--evolve", directory, generations, updates)
    assert result.returncode == 0, result.stderr
    events = [json.loads(line) for line in result.stdout.splitlines()]
    assert events[-1]["reason"] == reason
    assert not list(directory.iterdir())
    assert hashlib.sha256(executable.read_bytes()).digest() == before


def test_uninitialized_continuation_and_invalid_limits_reject(native, tmp_path):
    executable, _, _ = native
    assert run(executable, "--continue", tmp_path / "child").returncode != 0
    for generations, updates in [(9, 1), (3, 2049), (-1, 4), (2.5, 4), (3, "nan")]:
        result = run(executable, "--evolve", tmp_path / "child", generations, updates)
        assert result.returncode != 0
    assert not (tmp_path / "child").exists()


def test_checkpoint_semantics_and_output_publication(native, tmp_path):
    executable, checkpoint, model = native
    before = executable.read_bytes()
    with pytest.raises(FileExistsError):
        build(checkpoint, executable)
    link = tmp_path / "link"
    link.symlink_to(executable)
    with pytest.raises(FileExistsError):
        build(checkpoint, link)
    unsupported = tmp_path / "xor.pt"
    torch.save({"state_dict": model.state_dict(), "threshold": .55,
                "architecture": "xor-mask"}, unsupported)
    with pytest.raises(ValueError, match="architecture"):
        build(unsupported, tmp_path / "wrong-model")
    assert not (tmp_path / "wrong-model").exists()
    assert executable.read_bytes() == before