"""显式 --run 才加载固定小模型，在真实执行的教学环境中选择动作。"""

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

import agent_environment as lab


def experiment(directory):
    directory.mkdir(parents=True, exist_ok=False)
    plan = lab.make_plan()
    lab.save_new(directory/"plan.json", plan)  # 在加载模型、观察结果前固定。
    lab.save_new(directory/"control.json", lab.control_report(plan))
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer
    settings = plan["settings"]
    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)
    letters = list(lab.ACTIONS)
    ids = [tokenizer.encode(letter, add_special_tokens=False) for letter in letters]
    assert all(len(row) == 1 for row in ids)
    token_ids = [row[0] for row in ids]
    assert len(set(token_ids)) == len(letters)
    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"], settings=settings,
        plan_sha256=lab.fingerprint(plan), action_tokens=dict(zip(letters, token_ids)),
        source_sha256={p.name: lab.digest(p) for p in [Path(__file__), Path(lab.__file__)]},
        model_files_sha256={p.name: lab.digest(p) for p in snapshot.iterdir()
                            if p.is_file() and (p.suffix in {".json", ".safetensors"}
                                               or p.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=[],
    )

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

    checkpoint()
    for mode in plan["feedback"]:
        for case in plan["cases"]:
            env, history = lab.Printer(case), []
            row = dict(id=case["id"], mode=mode, trace=[], stopped="step_limit")
            report["rows"].append(row)
            for _ in range(settings["max_steps"]):
                conversation = lab.messages(history, mode)
                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+1 > settings["input_plus_output_cap"]:
                    row["stopped"] = "context_budget"
                    break
                started = time.perf_counter()
                with torch.inference_mode():
                    scores = model(**inputs).logits[0, -1]
                selected = scores[token_ids]
                probabilities = selected.softmax(0).tolist()
                action = letters[int(selected.argmax())]
                before = env.snapshot()
                result = env.step(action)
                observation = lab.visible(result, mode)
                trace = dict(
                    messages=conversation, input_ids=inputs["input_ids"][0].tolist(),
                    input_tokens=length, action=action, action_token=token_ids[letters.index(action)],
                    logits=selected.tolist(), conditional_probabilities=probabilities,
                    legal_token_mass=scores.softmax(0)[token_ids].sum().item(),
                    seconds=time.perf_counter()-started, before=before,
                    result=result, observation=observation, after=env.snapshot())
                row["trace"].append(trace)
                history.append(dict(action=action, observation=observation))
                row["grade"] = env.grade()
                checkpoint()
                if env.terminal:
                    row["stopped"] = env.terminal
                    break
            row["grade"] = env.grade()
            checkpoint()
            print(f"{mode}/{case['id']}: success={row['grade']['success']} "
                  f"steps={len(row['trace'])} actions="
                  f"{''.join(step['action'] for step in row['trace'])}", flush=True)
    report["complete"] = True
    report["finished_at"] = datetime.now(timezone.utc).isoformat()
    checkpoint()
    print(f"完成 16 条真实模型交互轨迹：{directory}", flush=True)


def main():
    parser = 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 新目录，才运行 16 条交互轨迹。")
        return
    if args.out_dir is None:
        parser.error("需要尚不存在的 --out-dir")
    experiment(args.out_dir)


if __name__ == "__main__":
    main()