import random
from pathlib import Path

import pytest
import torch

from lightnest import coder


def test_input_limits_are_explicit():
    with pytest.raises(ValueError, match='character'):
        coder.encode_say('返' * 257, None)
    with pytest.raises(ValueError, match='unsupported'):
        coder.encode_say('\U0001f9ea', None)
    with pytest.raises(ValueError, match='limit'):
        coder.encode_explain(b'\0' * 96, None)
    with pytest.raises(ValueError, match='aligned'):
        coder.encode_edit(b'\0', '返回第一个数', None)


def test_batch_fails_instead_of_retrying(monkeypatch):
    calls = []
    def oversize(rng):
        calls.append(1)
        return [0] * 257, 1
    monkeypatch.setattr(coder, '_sample', oversize)
    with pytest.raises(ValueError, match='context'):
        coder.make_batch(random.Random(0), 1, torch.device('cpu'))
    assert len(calls) == 1
    with pytest.raises(ValueError, match='positive'):
        coder.make_batch(random.Random(0), 0, torch.device('cpu'))


def test_generation_rejects_context_overflow():
    model = coder.Coder(dim=8, layers=1, heads=2)
    with pytest.raises(ValueError, match='context'):
        coder._generate(model, [coder.BOS] * 250, coder.BYTE_IDS, 96)
    with pytest.raises(ValueError, match='length'):
        model(torch.zeros((1, 257), dtype=torch.long))


def test_checkpoint_exact_vocab(tmp_path):
    path = tmp_path / 'small.pt'
    model = coder.Coder(dim=8, layers=1, heads=2)
    coder.save_model(model, path)
    loaded = coder.load_model(path, torch.device('cpu'))
    assert loaded.dim == 8
    blob = torch.load(path, weights_only=True)
    blob['vocab'][0], blob['vocab'][1] = blob['vocab'][1], blob['vocab'][0]
    torch.save(blob, path)
    with pytest.raises(ValueError, match='incompatible'):
        coder.load_model(path, torch.device('cpu'))


def test_explain_file_does_not_read_whole_file(tmp_path, monkeypatch):
    path = tmp_path / 'large.bin'
    with path.open('wb') as f:
        f.truncate(1024 * 1024)
    monkeypatch.setattr(Path, 'read_bytes', lambda *a: pytest.fail('unbounded read'))
    with pytest.raises(ValueError):
        coder.explain_file(None, path)


def test_training_subset_applies_to_all_modes(monkeypatch):
    chosen = tuple(coder.JOBS[:2])
    names = {job.name for job in chosen}
    codes = {job.code for job in chosen}
    observed = set()
    def say(sentence, program):
        observed.add('say')
        assert program in codes
        assert sentence in {s for job in chosen for s in job.say_train}
        return [0], 0
    def edit(prior, sentence, program):
        observed.add('edit')
        assert prior in codes and program in codes
        assert sentence in {s for s, name in coder.EDIT_TRAIN if name in names}
        return [0], 0
    def explain(program, sentence):
        observed.add('explain')
        assert program in codes
        assert sentence in {job.explain for job in chosen}
        return [0], 0
    monkeypatch.setattr(coder, 'encode_say', say)
    monkeypatch.setattr(coder, 'encode_edit', edit)
    monkeypatch.setattr(coder, 'encode_explain', explain)
    rng = random.Random(10)
    for _ in range(100):
        coder._sample(rng, chosen)
    assert observed == {'say', 'edit', 'explain'}


def test_semantic_augmentation_changes_bytes_preserves_branches():
    from lightnest.phrase import by_name, heldout_probes
    from lightnest.rv32 import simulate
    order = coder.TEMP_REGISTERS[1:] + coder.TEMP_REGISTERS[:1]
    for name in ('add', 'max', 'sum', 'range', 'count_neg'):
        job = by_name(name)
        transformed = coder.augment_program(job.code, register_order=order, nop_prefix=1)
        assert transformed != job.code
        assert len(transformed) == len(job.code) + 4
        assert all(simulate(transformed, xs) == expected
                   for xs, expected in heldout_probes(name, seed=88212, trials=12))
    with pytest.raises(ValueError, match='permute'):
        coder.augment_program(by_name('add').code, register_order=(0, 10, 11, 12))


def test_augmentation_bank_is_train_only_and_context_bounded():
    withheld = {'abs', 'mul', 'sum_squares', 'count_even', 'last'}
    jobs = tuple(job for job in coder.JOBS if job.name not in withheld)
    bank = coder.training_augmentations(jobs, 91241)
    assert set(bank) == {job.name for job in jobs}
    assert set(bank).isdisjoint(withheld)
    for job in jobs:
        for variant in bank[job.name]:
            assert len(variant) <= coder.MAX_PROGRAM_BYTES
            ids, _ = coder.encode_explain(variant, job.explain)
            assert len(ids) <= coder.MAX_LEN
            for target in jobs:
                sentence = max(target.say_train, key=len)
                ids, _ = coder.encode_edit(variant, '改成' + sentence, target.code)
                assert len(ids) <= coder.MAX_LEN
    for _ in range(10):
        tokens, _ = coder.make_batch(random.Random(15), 8, torch.device('cpu'), jobs, bank)
        assert tokens.shape[1] <= coder.MAX_LEN


def test_augmentation_only_changes_training_inputs(monkeypatch):
    jobs = tuple(coder.JOBS[:2])
    bank = {job.name: (b'\x13\0\0\0' + job.code,) for job in jobs}
    canonical = {job.code for job in jobs}
    variants = {value for values in bank.values() for value in values}
    modes = set()
    def say(sentence, target):
        modes.add('say')
        assert target in canonical
        return [0], 0
    def edit(prior, sentence, target):
        modes.add('edit')
        assert prior in variants and target in canonical
        return [0], 0
    def explain(prior, sentence):
        modes.add('explain')
        assert prior in variants
        return [0], 0
    monkeypatch.setattr(coder, 'encode_say', say)
    monkeypatch.setattr(coder, 'encode_edit', edit)
    monkeypatch.setattr(coder, 'encode_explain', explain)
    rng = random.Random(721)
    for _ in range(100):
        coder._sample(rng, jobs, bank)
    assert modes == {'say', 'edit', 'explain'}


def test_augmentation_rejects_unknown_opcodes_and_oversize():
    with pytest.raises(ValueError, match='opcode'):
        coder.augment_program(b'\xff' * 4)
    with pytest.raises(ValueError, match='limit'):
        coder.augment_program(bytes.fromhex('13000000') * 23, nop_prefix=2)


def test_checkpoint_size_budget_checked_before_load(tmp_path, monkeypatch):
    path = tmp_path / 'large.pt'
    with path.open('wb') as out:
        out.truncate(100 * 1024**2)
    monkeypatch.setattr(torch, 'load', lambda *a, **k: pytest.fail('must reject before tensor load'))
    with pytest.raises(ValueError, match='size budget'):
        coder.load_model(path, torch.device('cpu'))


def test_checkpoint_shape_budget_checked_before_construction(tmp_path, monkeypatch):
    path = tmp_path / 'small.pt'
    coder.save_model(coder.Coder(16, 1, 4), path)
    data = torch.load(path, weights_only=True)
    data['dim'] = 1000000
    torch.save(data, path)
    monkeypatch.setattr(coder, 'Coder', lambda *a, **k: pytest.fail('must reject before model allocation'))
    with pytest.raises(ValueError, match='parameter budget'):
        coder.load_model(path, torch.device('cpu'))