"""Public, generated supervision for the parameter-count experiment.

The assembler supplies the exact changed byte (or NONE).  This is supervised
training data; it is not an execution-feedback or native self-learning loop.
Independent final evaluation lives in a separate, unopened data directory.
"""
from __future__ import annotations

import argparse
from collections import Counter
import hashlib
import json
from pathlib import Path
import random
import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from lightnest import rv32
from lightnest.learned_patch import TRAIN_REGS, TRAIN_TEMPLATES


NEGATION_TEMPLATES = (
    "do not change {a} to {b}",
    "do not turn {a} into {b}",
    "do not replace {a} with {b}",
    "do not switch from {a} to {b}",
)
FAMILY_CYCLE = ("edit", "already_correct", "edit", "negation")
MAX_WINDOW = 256
MAX_COMMAND = 96


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def input_key(command: str, binary: bytes) -> str:
    """Length-delimited identity of the actual model input, not its label."""
    raw = command.encode("utf-8")
    return hashlib.sha256(len(raw).to_bytes(8, "little") + raw + binary).hexdigest()


def signed32(value: int) -> int:
    return (value + 2**31) % 2**32 - 2**31


def expected(operation: str, inputs: list[int]) -> int:
    left, right = inputs
    if operation == "add":
        return signed32(left + right)
    if operation == "sub":
        return signed32(left - right)
    raise ValueError("unsupported arithmetic operation")


def make_row(rng: random.Random, index: int) -> dict:
    family = FAMILY_CYCLE[index % len(FAMILY_CYCLE)]
    source, requested = rng.choice((("add", "sub"), ("sub", "add")))
    actual = requested if family == "already_correct" else source
    desired = actual if family == "negation" else requested
    r1, r2, out = (f"x{x}" for x in rng.sample(TRAIN_REGS, 3))
    prefix = rng.randint(0, 16)
    words = [rv32.addi("x0", "x0", rng.randint(-2048, 2047)) for _ in range(prefix)]
    words.extend((rv32.lw(r1, 0, "x0"), rv32.lw(r2, 4, "x0")))
    location = len(words)
    words.append(getattr(rv32, actual)(out, r1, r2))
    words.extend(rv32.addi("x0", "x0", rng.randint(-2048, 2047)) for _ in range(rng.randint(0, 7)))
    words.extend((rv32.addi("x10", out, 0), rv32.EBREAK))
    target_words = list(words)
    target_words[location] = getattr(rv32, desired)(out, r1, r2)
    pack = lambda ws: b"".join(w.to_bytes(4, "little") for w in ws)
    binary, target = pack(words), pack(target_words)
    different = [i for i, (old, new) in enumerate(zip(binary, target)) if old != new]
    if len(different) != int(family == "edit"):
        raise AssertionError("generator did not produce its declared edit count")
    pointer = different[0] if different else None
    template = rng.choice(NEGATION_TEMPLATES if family == "negation" else TRAIN_TEMPLATES)
    command = template.format(a=source, b=requested)
    if len(binary) > MAX_WINDOW or len(command.encode()) > MAX_COMMAND:
        raise AssertionError("public training input exceeds declared bounds")
    return {
        "id": index,
        "family": family,
        "command": command,
        "binary_hex": binary.hex(),
        "target_hex": target.hex(),
        "pointer": pointer,
        "replacement": target[pointer] if pointer is not None else 0,
        "actual_operation": actual,
        "desired_operation": desired,
        "source_operation": source,
        "requested_operation": requested,
        "registers": [r1, r2, out],
        "prefix_words": prefix,
        "template": template,
        "input_key": input_key(command, binary),
    }


def check_row(row: dict, semantic: bool = False) -> None:
    binary, target = bytes.fromhex(row["binary_hex"]), bytes.fromhex(row["target_hex"])
    if not 0 < len(binary) <= MAX_WINDOW or len(target) != len(binary):
        raise ValueError("invalid training binary size")
    if not 0 < len(row["command"].encode()) <= MAX_COMMAND:
        raise ValueError("invalid training command size")
    if row["input_key"] != input_key(row["command"], binary):
        raise ValueError("training input key mismatch")
    differences = [i for i, (old, new) in enumerate(zip(binary, target)) if old != new]
    pointer = row["pointer"]
    if pointer is None:
        if differences or row["family"] == "edit" or row["replacement"] != 0:
            raise ValueError("invalid NONE supervision")
    elif differences != [pointer] or row["replacement"] != target[pointer] or row["family"] != "edit":
        raise ValueError("invalid byte supervision")
    if semantic:
        edges = (-2**31, -1, 0, 1, 2**31 - 1)
        for left in edges:
            for right in edges:
                inputs = [left, right]
                if rv32.simulate(binary, inputs) != expected(row["actual_operation"], inputs):
                    raise ValueError("source execution disagrees with independent arithmetic")
                if rv32.simulate(target, inputs) != expected(row["desired_operation"], inputs):
                    raise ValueError("target execution disagrees with independent arithmetic")


def freeze(output: Path, seed: int, count: int) -> dict:
    if not 0 < count <= 200_000:
        raise ValueError("training count must be between 1 and 200000")
    output.mkdir(parents=True, exist_ok=True)
    if any(output.iterdir()):
        raise ValueError("refusing to overwrite a nonempty dataset directory")
    rng = random.Random(seed)
    rows, identities, keys = [], [], set()
    attempts = 0
    for index in range(count):
        while True:
            row = make_row(rng, index)
            attempts += 1
            if row["input_key"] not in keys:
                break
        check_row(row, semantic=index < 32)
        keys.add(row["input_key"])
        rows.append(row)
        identities.append({"id": index, "input_key": row["input_key"], "command": row["command"],
                           "binary_sha256": hashlib.sha256(bytes.fromhex(row["binary_hex"])).hexdigest()})
    dataset = output / "train.jsonl"
    dataset.write_text("".join(json.dumps(r, sort_keys=True) + "\n" for r in rows))
    key_file = output / "seen-inputs.jsonl"
    key_file.write_text("".join(json.dumps(r, sort_keys=True) + "\n" for r in identities))
    meta = {
        "seed": seed,
        "count": count,
        "attempts": attempts,
        "families": dict(Counter(r["family"] for r in rows)),
        "operations": dict(Counter(r["desired_operation"] for r in rows)),
        "train_registers": list(TRAIN_REGS),
        "edit_templates": list(TRAIN_TEMPLATES),
        "negation_templates": list(NEGATION_TEMPLATES),
        "max_window": MAX_WINDOW,
        "max_command": MAX_COMMAND,
        "observed_max_binary_bytes": max(len(r["binary_hex"]) // 2 for r in rows),
        "observed_max_command_bytes": max(len(r["command"].encode()) for r in rows),
        "dataset_sha256": sha256(dataset),
        "input_keys_sha256": sha256(key_file),
        "generator_sha256": sha256(Path(__file__)),
        "supervision": "assembler-supplied pointer/NONE and absolute replacement byte; not execution-only RL",
        "ordering": "stored rows are shared by every model and seed; the training protocol fixes their execution order",
        "validation": "all exact labels checked; first 32 sources and targets each executed on 25 independent arithmetic inputs",
    }
    (output / "dataset.json").write_text(json.dumps(meta, indent=2))
    print(json.dumps(meta), flush=True)
    return meta


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, required=True)
    parser.add_argument("--count", type=int, required=True)
    args = parser.parse_args()
    freeze(args.output, args.seed, args.count)


if __name__ == "__main__":
    main()