"""显式 --run 才运行长上下文与记忆专题的小模型探测；CPU，无静默截断。"""

import argparse
from datetime import datetime, timezone
from importlib.metadata import version
import json
from pathlib import Path
import platform
import time

import context_memory as lab


def experiment(directory):
    directory.mkdir(parents=True, exist_ok=False)
    plan = lab.make_plan()
    lab.save_new(directory / "plan.json", plan)  # 在下载、加载、生成之前固定。
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer
    settings = plan["generation"]
    torch.set_num_threads(settings["threads"])
    torch.manual_seed(settings["seed"])
    snapshot = Path(snapshot_download(
        plan["model"], revision=plan["revision"],
        allow_patterns=["*.json", "*.safetensors", "merges.txt", "README.md", "LICENSE"]))
    tokenizer = AutoTokenizer.from_pretrained(snapshot, trust_remote_code=False)
    encoded = []
    for case in plan["cases"]:
        conversation = lab.messages(case)
        rendered = tokenizer.apply_chat_template(
            conversation, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer.apply_chat_template(
            conversation, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt")
        length = inputs["input_ids"].shape[1]
        if length + settings["max_new_tokens"] > settings["input_plus_output_cap"]:
            raise ValueError(f"{case['id']} 超出预先固定的预算，停止而不截断")
        encoded.append((case, conversation, rendered, inputs))
    model = AutoModelForCausalLM.from_pretrained(
        snapshot, dtype=torch.float32, attn_implementation=settings["attention"],
        trust_remote_code=False).eval()
    report = dict(
        complete=False, started_at=datetime.now(timezone.utc).isoformat(),
        model=plan["model"], revision=plan["revision"], generation=settings,
        plan_sha256=lab.fingerprint(plan),
        source_sha256={path.name: lab.digest(path) for path in
                       (Path(lab.__file__), Path(__file__))},
        model_files_sha256={path.name: lab.digest(path) for path in snapshot.iterdir()
                            if path.is_file() and (path.suffix in {".json", ".safetensors"}
                                                   or path.name == "merges.txt")},
        environment=dict(python=platform.python_version(), platform=platform.platform(),
                         packages={name: version(name) for name in
                                   ("torch", "transformers", "huggingface-hub",
                                    "tokenizers", "safetensors")}),
        rows=[],
    )
    target = directory / "model.json"

    def checkpoint():
        temporary = directory / "model.pending.json"
        temporary.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n")
        temporary.replace(target)

    checkpoint()
    eos = model.generation_config.eos_token_id
    eos = [eos] if isinstance(eos, int) else eos
    for case, conversation, rendered, inputs in encoded:
        length = inputs["input_ids"].shape[1]
        started = time.perf_counter()
        with torch.inference_mode():
            output = model.generate(
                **inputs, max_new_tokens=settings["max_new_tokens"],
                do_sample=False, temperature=None, top_p=None, top_k=None,
                pad_token_id=tokenizer.eos_token_id)
        seconds = time.perf_counter() - started
        ids = output[0, length:].tolist()
        text = tokenizer.decode(ids, skip_special_tokens=True)
        row = dict(id=case["id"], messages=conversation, rendered_prompt=rendered,
                   input_ids=inputs["input_ids"][0].tolist(), input_tokens=length,
                   output_ids=ids, output_tokens=len(ids), text=text,
                   stopped_on_eos=bool(ids and ids[-1] in eos), seconds=seconds,
                   **lab.grade(case, text))
        report["rows"].append(row)
        checkpoint()
        print(f"{case['id']}: {text!r}, 输入 {length}, "
              f"约定正确 {row['contract_correct']}", flush=True)
    report["complete"] = True
    report["finished_at"] = datetime.now(timezone.utc).isoformat()
    checkpoint()
    lab.save_new(directory / "memory.json", lab.memory_report())
    print(f"已完成全部 {len(report['rows'])} 次探测：{target}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--out-dir", type=Path)
    args = parser.parse_args()
    if not args.run:
        print("显式传入 --run --out-dir 新目录，才下载并运行固定的 27 次模型探测。")
        return
    if args.out_dir is None:
        parser.error("请用 --out-dir 指定尚不存在的目录")
    experiment(args.out_dir)


if __name__ == "__main__":
    main()