"""Generate ongoing-chat token fixtures through MTPLX's actual server encoder.

Use MTPLX's installed Python environment. Reads only local tokenizer files and
the supplied README; no inference, server, tools or model downloads are started.
"""

import argparse
import hashlib
import json
import os
from pathlib import Path

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

from mtplx.runtime import _load_tokenizer_resilient
from mtplx.server import openai as server


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model", type=Path)
    parser.add_argument("readme", type=Path)
    parser.add_argument("--source-run", type=Path,
                        help="Use exact prompts and generated history from a completed model-eval JSONL")
    args = parser.parse_args()
    tokenizer = _load_tokenizer_resilient(
        args.model.resolve(), json.loads((args.model / "config.json").read_text()),
    )
    readme = args.readme.read_text()
    prompts = [
        "Give a summary of the following text:\n\n" + readme,
        "Tell me a complete short fictional story. Choose the setting yourself; do not ask questions.",
        "Return one Python code block defining has_close_elements(numbers: list[float], threshold: float) -> bool.",
    ]
    replies = [
        {"role": "assistant", "content": text,
         "reasoning_content": "**Plan:** answer the request directly."}
        for text in [
            "DS4Server runs local language models on Apple silicon.",
            "A lighthouse keeper found a letter in an empty bottle. She wrote back, and the sea carried her reply home.",
            "```python\ndef has_close_elements(numbers, threshold):\n    return any(abs(a-b) < threshold for i, a in enumerate(numbers) for b in numbers[i+1:])\n```",
        ]
    ]
    if args.source_run:
        records = [json.loads(line) for line in args.source_run.read_text().splitlines()]
        start = next(record for record in records if record["event"] == "start")
        results = [record for record in records if record["event"] == "result"]
        if not (start["plain_chat"] and not start["system_prompt"]
                and start["settings"]["reasoning"] == "low"
                and len(start["prompts"]) == 3
                and [record["turn"] for record in results] == [1, 2, 3]
                and all(record["ok"] and record["finish_reason"] == "stop" for record in results)
                and readme in start["prompts"][0]):
            parser.error("source must be a complete three-turn Low plain chat using the supplied README")
        prompts = start["prompts"]
        replies = [{"role": "assistant", "content": record["content"],
                    "reasoning_content": record["reasoning"] or ""} for record in results]
    history = []
    cases = []
    for index, prompt in enumerate(prompts):
        history.append({"role": "user", "content": prompt})
        ids = server._encode_messages_uncached(
            tokenizer, [server.ChatMessage(**message) for message in history],
            enable_thinking=True, reasoning_effort="low",
            preserve_reasoning_history=True, tools=None,
        )
        cases.append({"messages": list(history), "token_ids": ids,
                      "rendered": tokenizer.decode(ids, skip_special_tokens=False)})
        history.append(replies[index])
    source = Path(server.__file__).resolve()
    print(json.dumps({"reference": "MTPLX server _encode_messages_uncached",
                      "source": str(source),
                      "source_sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
                      "tokenizer_sha256": hashlib.sha256((args.model / "tokenizer.json").read_bytes()).hexdigest(),
                      "chat_template_sha256": hashlib.sha256(tokenizer.chat_template.encode()).hexdigest(),
                      "readme_sha256": hashlib.sha256(readme.encode()).hexdigest(),
                      "source_run_sha256": hashlib.sha256(args.source_run.read_bytes()).hexdigest() if args.source_run else None,
                      "cases": cases}))


if __name__ == "__main__":
    main()