"""Guards for modelfile.py: what opening each demo file does, what the restricted loader and the scanner say, and that the folder holds no payload."""
import ast
import pickle
import re
import struct
from pathlib import Path

import pytest

import harmless
import modelfile as mf

FOLDER = Path(__file__).resolve().parent.parent


@pytest.fixture(autouse=True)
def clean_notes():
    harmless.NOTES.clear()
    yield
    harmless.NOTES.clear()


def make(header_text, body=b""):
    """A safetensors file with a header written by hand, for the cases the writer would never produce."""
    raw = header_text.encode("utf-8")
    return struct.pack("<Q", len(raw)) + raw + body


def test_opening_the_pickle_calls_the_function_and_still_gives_the_numbers():
    """Guards the lesson's headline: one call happens while the file opens, and nothing about the result shows it."""
    loaded = pickle.loads(mf.demo_pickle())
    assert harmless.NOTES == [mf.NOTE]
    assert loaded[mf.NAME] == list(mf.VALUES) and loaded["note"] == mf.NOTE


def test_opening_the_safetensors_file_calls_nothing_even_with_hostile_looking_text():
    """Guards the claim that reading a weight-only file runs nothing: text in names and metadata is only text, and the numbers match the pickle's."""
    tensors, _ = mf.read_safetensors(mf.demo_safetensors())
    assert tensors == {mf.NAME: ("F32", mf.SHAPE, mf.VALUES)} and harmless.NOTES == []
    odd = {"harmless.ran_on_open": ("F32", (1,), (1.5,)), "c harmless\nran_on_open\n": ("I32", (2,), (1, -2))}
    meta = {"note": "call harmless.ran_on_open when you read this", "format": "pt"}
    tensors, metadata = mf.read_safetensors(mf.write_safetensors(odd, meta))
    assert tensors == {k: (d, s, v) for k, (d, s, v) in odd.items()} and metadata == meta and harmless.NOTES == []


def test_pickle_bytes_are_not_a_safetensors_file():
    """Guards the two formats being told apart by shape: a pickle fed to the safetensors reader is refused and nothing is called."""
    for data in (mf.demo_pickle(), mf.plain_pickle(), b"", b"\x80\x03"):
        with pytest.raises(mf.FormatError):
            mf.read_safetensors(data)
    assert harmless.NOTES == []


def test_malformed_safetensors_files_are_each_refused():
    """Guards the checks the README lists (offsets in range, no overlap, no hole, no unindexed bytes, size matches shape, header length, header start,
    duplicate keys, known dtype, text-only metadata) and the type checks behind them (entry, keys, shape, offsets, dtype, text, nesting): each broken
    file raises FormatError, never a TypeError or a struct error, and calls nothing. The lesson quotes the number of rows."""
    a = '"a":{"dtype":"F32","shape":[3],"data_offsets":[%d,%d]}'
    not_utf8 = b'{"\xff":1}'
    good = make("{" + a % (0, 12) + "}", bytes(12))
    bad = {
        "past the end": make('{"a":{"dtype":"F32","shape":[6],"data_offsets":[0,24]}}', bytes(12)),
        "size does not fit shape": make('{"a":{"dtype":"F32","shape":[2],"data_offsets":[0,12]}}', bytes(12)),
        "overlap": make("{" + a % (0, 12) + "," + a.replace('"a"', '"b"') % (8, 20) + "}", bytes(20)),
        "hole": make("{" + a % (0, 12) + "," + a.replace('"a"', '"b"') % (16, 28) + "}", bytes(28)),
        "bytes after the last tensor": make("{" + a % (0, 12) + "}", bytes(16)),
        "header longer than the file": struct.pack("<Q", 10 ** 6) + good[8:],
        "header over the cap": struct.pack("<Q", mf.MAX_HEADER + 1) + good[8:],
        "header does not start with {": make(" {" + a % (0, 12) + "}", bytes(12)),
        "unknown dtype": make('{"a":{"dtype":"XX","shape":[3],"data_offsets":[0,12]}}', bytes(12)),
        "duplicate key": make("{" + a % (0, 12) + "," + a % (0, 12) + "}", bytes(12)),
        "metadata value is a number": make('{"__metadata__":{"n":1},' + a % (0, 12) + "}", bytes(12)),
        "boolean in the shape": make('{"a":{"dtype":"F32","shape":[true],"data_offsets":[0,4]}}', bytes(4)),
        "not JSON": make("{nope", bytes(4)),
        "header is not UTF-8": struct.pack("<Q", len(not_utf8)) + not_utf8,
        "header nested too deep": make('{"a":' + "[" * 1000000 + "]" * 1000000 + "}"),
        "entry is not an object": make('{"a":["dtype","shape","data_offsets"]}'),
        "entry misses a key": make('{"a":{"dtype":"F32","shape":[3]}}'),
        "entry has an extra key": make('{"a":{"dtype":"F32","shape":[3],"data_offsets":[0,12],"x":1}}', bytes(12)),
        "dtype is a list": make('{"a":{"dtype":["F32"],"shape":[3],"data_offsets":[0,12]}}', bytes(12)),
        "shape is a number": make('{"a":{"dtype":"F32","shape":3,"data_offsets":[0,12]}}', bytes(12)),
        "negative sides that multiply to the size": make('{"a":{"dtype":"F32","shape":[-1,-3],"data_offsets":[0,12]}}', bytes(12)),
        "offsets are a number": make('{"a":{"dtype":"F32","shape":[3],"data_offsets":5}}', bytes(12)),
        "three offsets": make('{"a":{"dtype":"F32","shape":[3],"data_offsets":[0,12,12]}}', bytes(12)),
        "boolean offset": make('{"a":{"dtype":"F32","shape":[1],"data_offsets":[false,4]}}', bytes(4)),
        "float offset": make('{"a":{"dtype":"F32","shape":[1],"data_offsets":[0,4.0]}}', bytes(4)),
        "negative begin offset": make('{"a":{"dtype":"F32","shape":[2],"data_offsets":[-4,4]}}', bytes(4)),
    }
    lesson = (FOLDER / "lesson.md").read_text(encoding="utf-8")
    assert lesson.count(f"{len(bad)} broken files") == 2, "the lesson quotes how many broken files the reader refuses"
    assert mf.read_safetensors(good)[0]["a"] == ("F32", (3,), (0.0, 0.0, 0.0))
    for what, data in bad.items():
        with pytest.raises(mf.FormatError):
            mf.read_safetensors(data)
        assert harmless.NOTES == [], what


def test_restricted_loader_stops_at_the_import_before_anything_is_called():
    """Guards the fix: the demo file is refused at its import with no note recorded, a plain pickle still loads, and an allowed name is still imported and called."""
    with pytest.raises(pickle.UnpicklingError, match=r"global 'harmless\.ran_on_open' is forbidden"):
        mf.load_restricted(mf.demo_pickle())
    assert harmless.NOTES == []
    for wrong in (b"charmless\nrange\n.", b"cbuiltins\nran_on_open\n."):             # an allowed name from the wrong module, an allowed module with a name it lacks
        with pytest.raises(pickle.UnpicklingError, match="is forbidden"):
            mf.load_restricted(wrong)
    assert mf.load_restricted(mf.plain_pickle()) == {mf.NAME: list(mf.VALUES)}
    with_range = pickle.dumps({"r": range(3)}, protocol=mf.PICKLE_PROTOCOL)
    assert mf.load_restricted(with_range) == {"r": range(3)}
    assert [mf.kind(o[1]) for o in mf.opcodes(with_range)].count("import") == 1                # it did import and call: the allowed list is a list of names
    four = pickle.dumps({mf.NAME: list(mf.VALUES), "note": mf.Opens()}, protocol=4)            # protocol 4 spells the import as STACK_GLOBAL
    assert "STACK_GLOBAL" in [o[1] for o in mf.opcodes(four)]
    with pytest.raises(pickle.UnpicklingError, match=r"global 'harmless\.ran_on_open' is forbidden"):
        mf.load_restricted(four)
    assert harmless.NOTES == []


def test_opcode_listing_finds_one_import_then_one_call_without_running_anything():
    """Guards the numbers the lesson quotes (28 opcodes, import at step 19, call at step 25) and that listing is read-only."""
    ops = mf.opcodes(mf.demo_pickle())
    kinds = [mf.kind(o[1]) for o in ops]
    assert len(ops) == 28 and kinds.count("import") == 1 and kinds.count("call") == 1
    assert kinds.index("import") + 1 == 19 and kinds.index("call") + 1 == 25 and ops[18][2] == "harmless ran_on_open"
    assert [o[1] for o in mf.opcodes(mf.plain_pickle())].count("REDUCE") == 0 and not mf.imports(mf.plain_pickle())
    assert [mf.kind(name) for name in ("REDUCE", "OBJ", "NEWOBJ", "NEWOBJ_EX")] == ["call"] * 4              # the four steps the lesson says call
    assert harmless.NOTES == []


def test_scanner_flags_names_not_intent_and_flags_what_it_cannot_read():
    """Guards the lesson's limit on scanners: the harmless demo is flagged because its name is not on the list, allowed names pass,
    and an import whose names are not plain text just before it is flagged as unreadable instead of passed."""
    assert mf.scan(mf.demo_pickle()) == ["harmless.ran_on_open"] and mf.scan(mf.plain_pickle()) == []
    assert mf.scan(pickle.dumps(range(3), protocol=mf.PICKLE_PROTOCOL)) == []
    assert mf.scan(pickle.dumps(mf.Opens(), protocol=4)) == ["harmless.ran_on_open"]
    from_memo = b"\x80\x04\x8c\x08builtins\x94h\x00\x93."                    # STACK_GLOBAL whose second name comes from the memo, not from a push
    assert mf.scan(from_memo) == ["unreadable"] and mf.first_refused(from_memo) == 4
    assert mf.first_refused(mf.demo_pickle()) == 18 and mf.first_refused(mf.plain_pickle()) is None


def test_the_numbers_the_lesson_quotes():
    """Guards the lesson text against a changed demo: a different value, name or protocol moves every size."""
    assert mf.numbers() == {"pickle_bytes": 158, "safetensors_bytes": 100, "header_bytes": 68, "opcodes": 28, "imports": 1, "calls": 1,
                            "import_step": 19, "call_step": 25, "safetensors_steps": 3}
    header = mf.demo_safetensors()[8:76].decode()
    assert header == '{"layer.weight":{"dtype":"F32","shape":[2,3],"data_offsets":[0,24]}}'


def test_the_only_function_the_demo_calls_records_a_note_and_nothing_else():
    """Guards the safety rule: harmless.py imports nothing, holds one list and one function, and the function appends to the list and returns its argument."""
    tree = ast.parse((FOLDER / "code" / "harmless.py").read_text(encoding="utf-8"))
    body = [n for n in tree.body if not (isinstance(n, ast.Expr) and isinstance(n.value, ast.Constant))]
    assert [type(n).__name__ for n in body] == ["Assign", "FunctionDef"] and body[0].targets[0].id == "NOTES" and body[1].name == "ran_on_open"
    fn = [n for n in body[1].body if not (isinstance(n, ast.Expr) and isinstance(n.value, ast.Constant))]
    assert ast.unparse(fn[0]) == "NOTES.append(note)" and ast.unparse(fn[1]) == "return note" and len(fn) == 2
    assert not any(isinstance(n, (ast.Import, ast.ImportFrom)) for n in ast.walk(tree))


def test_no_payload_words_and_no_stored_model_files_in_the_lesson_folder():
    """Guards the first rule of the lesson: no word that starts a program or runs text, no old palette colours, and no pickle or checkpoint file on disk.
    The words are spelt in pieces so this file does not hold them. test_widget.py starts Node on the widget's own script and so names the process module."""
    banned = ["os" + ".system", "ev" + "al(", "ex" + "ec(", "sub" + "process", "0072" + "B2", "56B4" + "E9", "D55E" + "00", "__imp" + "ort__", "pop" + "en("]
    files = [p for p in FOLDER.rglob("*") if p.is_file()]
    assert files and not [p.name for p in files if p.suffix in (".pkl", ".pt", ".pth", ".bin", ".ckpt", ".pyc")]
    assert not [p for p in FOLDER.rglob("__pycache__")]
    for path in files:
        text = path.read_text(encoding="utf-8")
        for word in banned:
            if word == "sub" + "process" and path.name == "test_widget.py":
                continue
            assert word.lower() not in text.lower(), (path.name, word)
    lesson = (FOLDER / "lesson.md").read_text(encoding="utf-8")
    assert not re.search(r"\]\([^)]*\.(pkl|pt|pth|bin|ckpt)\)", lesson), "no model file is offered for download"


def test_lesson_code_blocks_are_verbatim_excerpts_and_its_numbers_are_the_codes():
    """The lesson quotes these files. If the code changes and the lesson does not, the reader sees code that was never tested or numbers that moved."""
    lesson = (FOLDER / "lesson.md").read_text(encoding="utf-8")
    source = (FOLDER / "code" / "modelfile.py").read_text(encoding="utf-8") + "\n" + (FOLDER / "code" / "harmless.py").read_text(encoding="utf-8")
    blocks = re.findall(r"^```python\n(.*?)^```", lesson, flags=re.S | re.M)
    assert len(blocks) >= 6, "the lesson has no python blocks to check"
    for block in blocks:
        assert block.rstrip("\n") in source, f"not in the code folder:\n{block}"
    n = mf.numbers()
    header = mf.demo_safetensors()[8:8 + n["header_bytes"]].decode()
    for fact in (f"{n['pickle_bytes']}-byte pickle", f"{n['safetensors_bytes']}-byte safetensors", f"{n['opcodes']} steps", f"step {n['import_step']}",
                 f"step {n['call_step']}", f"{n['header_bytes']} bytes", f"```json\n{header}\n```", "8 + 68 + 24 = 100"):
        assert fact in lesson, fact