"""Fine-tune a LoRA on Tinker from `machine export` JSONL.

Run inside a tinker-cookbook checkout's environment:

    cd ~/code/tinker-cookbook && source .venv/bin/activate
    python ~/code/machine/scripts/train_sft.py data.jsonl ~/code/machine/.machine/training/my-run \\
        --model Qwen/Qwen3.6-35B-A3B --epochs 2

Logging under a workspace's .machine/training/ puts the run on
`machine dashboard`, which charts its losses while it trains.

A line with "whole_run": true is one whole run, trained on every assistant
message marked "trainable" (Inkling only: each step's prompt is a token prefix
of the run's sequence there, so it matches sampling at a fraction of the
tokens of one example per step).

Adapted from tinker-cookbook's recipes/sl_basic.py. The learning rate defaults
to the cookbook's calibrated rate for the model (Llama and Qwen); Inkling has
none, so it defaults to 2e-4, the rate of the cookbook's Inkling SFT recipe
(recipes/audio/asr/sl_train.py). Inkling (the tml_v0 renderer) is rendered at
--effort, which must be the thinking effort runs sampled at: Tinker's OpenAI
endpoint defaults to high, 0.9. Trains
on the last assistant message of each conversation, which is correct for the
Qwen3.5/3.6 renderer on multi-turn data and identical on single-turn data.
Checkpoints, including the sampler path to evaluate with
`machine suite --model tinker://...`, are written to <log_path>/checkpoints.jsonl,
and per-step losses to <log_path>/metrics.jsonl.
"""

import argparse
import asyncio
import json
import os
import shutil

import chz
from tinker_cookbook import hyperparam_utils, model_info
from tinker_cookbook.renderers import TrainOnWhat
from tinker_cookbook.renderers.base import ToolCall
from tinker_cookbook.supervised.common import datum_from_model_input_weights
from tinker_cookbook.supervised.data import conversation_to_datum
from tinker_cookbook.supervised import train
from tinker_cookbook.supervised.data import FromConversationFileBuilder
from tinker_cookbook.supervised.types import ChatDatasetBuilderCommonConfig

parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("data", help="JSONL from `machine export`")
parser.add_argument("log_path", help="directory for checkpoints and metrics, e.g. .machine/training/<name>")
parser.add_argument("--model", default="Qwen/Qwen3.6-35B-A3B")
parser.add_argument("--epochs", type=int, default=2)
parser.add_argument("--batch-size", type=int, default=8)
parser.add_argument("--lr", type=float, default=None, help="default: the cookbook's calibrated rate")
parser.add_argument("--test-size", type=int, default=8, help="examples held out to track loss")
parser.add_argument("--max-length", type=int, default=65536)
parser.add_argument("--dry-run", action="store_true", help="render the first batch, print its token counts, and exit without training")
parser.add_argument("--effort", type=float, default=0.9, help="Inkling only: thinking effort to render at; match the runs' (endpoint default: high, 0.9)")
args = parser.parse_args()

# Keep the training data with the run, so the run can be inspected and reproduced later.
os.makedirs(args.log_path, exist_ok=True)
data_copy = os.path.join(args.log_path, "data.jsonl")
if os.path.abspath(args.data) != os.path.abspath(data_copy):
    shutil.copyfile(args.data, data_copy)

renderer = model_info.get_recommended_renderer_name(args.model)
if args.lr is None:
    try:
        args.lr = hyperparam_utils.get_lr(args.model)
    except NotImplementedError:
        if not args.model.startswith("thinkingmachines/Inkling"):
            raise
        args.lr = 2e-4
    print(f"learning rate {args.lr}")


@chz.chz
class ToolAwareConversationFileBuilder(FromConversationFileBuilder):
    """FromConversationFileBuilder for `machine export` lines that carry tools.

    A line's "tools" (the tools the model was offered) are rendered into the
    system prompt with the renderer's own tool declaration, as the sampling
    endpoint did, and OpenAI-style tool_calls become the renderer's ToolCall
    objects. Lines without tools are rendered as before.
    """

    def __call__(self):
        train_ds, test_ds = super().__call__()
        renderer_obj = self.renderer
        max_length = self.common_config.max_length
        train_on_what = TrainOnWhat(self.common_config.train_on_what)

        def convert(row: dict) -> list:
            messages = [dict(m) for m in row["messages"]]
            for m in messages:
                for key in [k for k, v in m.items() if v is None]:
                    del m[key]
                if m.get("tool_calls"):
                    m["tool_calls"] = [
                        ToolCall(id=tc["id"], function=ToolCall.FunctionBody(name=tc["function"]["name"], arguments=tc["function"]["arguments"]))
                        for tc in m["tool_calls"]
                    ]
            tools = json.loads(row["tools_json"]) if row.get("tools_json") else None
            if tools and messages and messages[0]["role"] == "system":
                system = messages[0]["content"]
                text = system if isinstance(system, str) else "".join(p.get("text", "") for p in system if p.get("type") == "text")
                # Qwen's renderer declares the tools as given, so they go in wrapped as OpenAI tools
                # ({"type": "function", "function": ...}), which is how Tinker's OpenAI-compatible endpoint
                # renders them. tml_v0 (Inkling) wraps them itself. Either way the prompt is token for token
                # what the endpoint sampled from (checked against its prompt_tokens on dev runs).
                declared = tools if renderer == "tml_v0" else [{"type": "function", "function": t} for t in tools]
                messages = renderer_obj.create_conversation_prefix_with_tools(declared, system_prompt=text) + messages[1:]
            return messages

        def map_fn(row: dict):
            if row.get("whole_run") and renderer != "tml_v0":
                raise ValueError(f"whole-run examples are checked for Inkling's tml_v0 renderer only, not {renderer}; train on step examples")
            if renderer == "tml_v0":
                # A whole-run example trains every assistant message marked trainable; a step example, its last.
                what = TrainOnWhat.CUSTOMIZED if row.get("whole_run") else train_on_what
                model_input, weights = renderer_obj.build_supervised_example(convert(row), train_on_what=what, effort=args.effort)
                return datum_from_model_input_weights(model_input, weights, max_length, reduction="mean")
            return conversation_to_datum(convert(row), renderer_obj, max_length, train_on_what)

        for ds in (train_ds, test_ds):
            if ds is not None:
                ds.map_fn = map_fn
        return train_ds, test_ds


# Tools go in as a JSON string: their nested schemas do not fit one Arrow column type across lines.
data_path = args.data
with open(args.data) as f:
    lines = [json.loads(line) for line in f if line.strip()]
if any("tools" in line for line in lines):
    data_path = os.path.join(args.log_path, "data.tools.jsonl")
    with open(data_path, "w") as f:
        for line in lines:
            tools = line.pop("tools", None)
            line["tools_json"] = json.dumps(tools) if tools else ""
            f.write(json.dumps(line) + "\n")

dataset = ToolAwareConversationFileBuilder(
    common_config=ChatDatasetBuilderCommonConfig(
        model_name_for_tokenizer=args.model,
        renderer_name=renderer,
        max_length=args.max_length,
        batch_size=args.batch_size,
        train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE,
    ),
    file_path=data_path,
    test_size=args.test_size,
)
if args.dry_run:
    train_ds, _ = dataset()
    for datum in train_ds.get_batch(0):
        weights = datum.loss_fn_inputs["weights"].to_numpy()
        print(f"example: {datum.model_input.length} tokens, {int((weights > 0).sum())} trained")
    raise SystemExit(0)

config = chz.Blueprint(train.Config).apply({
    "log_path": args.log_path,
    "recipe_name": "machine_sft",
    "model_name": args.model,
    "renderer_name": renderer,
    "dataset_builder": dataset,
    "learning_rate": args.lr,
    "num_epochs": args.epochs,
    "eval_every": 5,
}).make()
asyncio.run(train.main(config))