"""Instrument an external callable harness and export an explicitly synthetic trace.

No API calls, GPU, or training take place. Replace BytePolicy with inference that
returns its actual token IDs and log probabilities before collecting real data.
"""

from __future__ import annotations

import argparse
import copy
import hashlib
import json
from collections.abc import Callable

from hessian_recipes.trajectories import SCHEMA_VERSION, write_trajectories


CHAT_TEMPLATE = "json-context-utf8-bytes-v1"


class BytePolicy:
    """Deterministic fixture policy; UTF-8 bytes are its native token vocabulary."""

    def __call__(self, context: list[dict]) -> dict:
        prompt = json.dumps(context, sort_keys=True, ensure_ascii=False)
        response = "42" if context[-1]["role"] == "tool" else "multiply(6, 7)"
        tokens = list(response.encode("utf-8"))
        return {
            "text": response,
            "prompt_token_ids": list(prompt.encode("utf-8")),
            "response_token_ids": tokens,
            "response_loss_mask": [1] * len(tokens),
            "behavior_logprobs": [0.0] * len(tokens),
            "finish_reason": "stop" if response == "42" else "tool_calls",
        }


def external_harness(policy: Callable[[list[dict]], dict]) -> str:
    """A team's existing tool loop: instrumentation does not replace this loop."""
    context = [{"role": "user", "content": "Use the calculator to multiply 6 by 7."}]
    first = policy(context)
    if first["text"] != "multiply(6, 7)":
        return "unsupported tool request"
    context.append({"role": "assistant", "content": first["text"]})
    context.append({"role": "tool", "content": str(6 * 7)})
    return policy(context)["text"]


def build_episode() -> dict:
    calls = []
    policy = BytePolicy()

    def instrumented(context: list[dict]) -> dict:
        sampled = policy(context)
        calls.append({
            "call_id": f"call-{len(calls) + 1}",
            "context": copy.deepcopy(context),
            "response_role": "assistant",
            **{key: value for key, value in sampled.items() if key != "text"},
        })
        return sampled

    answer = external_harness(instrumented)
    return {
        "schema_version": SCHEMA_VERSION,
        "episode_id": "calculator-fixture-0",
        "group_id": "calculator-fixture-policy-v1",
        "task_id": "multiply-6-7",
        "seed": 0,
        "source_kind": "fixture",
        "policy": {
            "model_id": "fixture/byte-policy",
            "revision": "byte-policy-v1",
            "tokenizer_id": "fixture/utf8-bytes",
            "tokenizer_revision": "utf8-v1",
            "chat_template_sha256": hashlib.sha256(CHAT_TEMPLATE.encode()).hexdigest(),
        },
        "harness": {"name": "external-calculator", "version": "1.0", "revision": "fixture-v1"},
        "reward": {"value": float(answer == "42"), "scorer": "answer-equals-42", "revision": "v1"},
        "calls": calls,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", default="artifacts/trajectories/calculator.jsonl")
    args = parser.parse_args()
    count = write_trajectories(args.output, [build_episode()])
    print(f"Wrote {count} fixture episode to {args.output}; no model training was performed.")


if __name__ == "__main__":
    main()