"""Filter, deduplicate and decontaminate agent trajectories, and write a training set.

Purpose: the stage that decides what the fine-tune will actually learn. Reads the
    scrubbed episodes, keeps the ones whose outcome and shape are worth imitating,
    removes exact and near-duplicates, drops anything that overlaps the evaluation
    suite, splits by task so that no task appears in both halves, and writes the two
    layouts this part trains from: TRL's conversational format with a `tools` column,
    and mlx-lm's tools format for Track M. Every rejection is counted by reason and a
    sample of each reason is written out, because a filter you cannot inspect is a
    filter you cannot trust.
Platform: all (standard library only; no model, no accelerator, no network)
Minimum memory: 8 GB nominally, and far less in practice: this is text in memory
Assumes: Python 3.10 or newer. The input is the JSON-lines episode file
    scrub-trajectories.py wrote. --tasks takes the evaluation suite you will measure
    with, in the shape of Part 26's agent-tasks.json: an object with a "tasks" list
    whose items carry "id" and "task". agentlog.py sits next to this file.

Usage: python3 trajectories-to-sft.py --in clean/trajectories.jsonl --out-dir . \\
           --tasks agent-tasks.json --report filter-report.json --labbook labbook.md
       python3 trajectories-to-sft.py --in clean/trajectories.jsonl --out-dir . \\
           --include-failures --max-steps 12 --near-duplicate 0.7
       python3 trajectories-to-sft.py --in clean/trajectories.jsonl --print-example

Written under --out-dir:
  data/{train,valid}.jsonl        {"messages": [...], "tools": [...]}, arguments as objects
  data-mlx/{train,valid}.jsonl    the same episodes with tool-call arguments as strings
  filter-report.json              counts by reason, plus a sample of each

The two layouts differ in one field on purpose. The Transformers chat templates and
TRL's SFT trainer expect a tool call's `arguments` to be an object; mlx-lm's documented
tools example carries them as a JSON string. Writing both here means the same episodes
train on every track without anyone editing a file by hand.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import random
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

import agentlog

WORD_RE = re.compile(r"[a-z0-9]+")


# --------------------------------------------------------------------------------------
# Text comparison, the same containment measure Part 13's decontaminate.py uses
# --------------------------------------------------------------------------------------

def words(text: str) -> list[str]:
    return WORD_RE.findall(text.lower())


def ngrams(tokens: list[str], n: int) -> set[tuple[str, ...]]:
    if len(tokens) < n:
        return {tuple(tokens)} if tokens else set()
    return {tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1)}


def containment(a: set[tuple[str, ...]], b: set[tuple[str, ...]]) -> float:
    """How much of a is also in b. Containment rather than Jaccard, because a short
    evaluation task buried inside a long training task is the case that matters."""
    return (len(a & b) / len(a)) if a else 0.0


def episode_text(episode: dict[str, Any]) -> str:
    """Everything the model would see and write, flattened, for hashing and comparison."""
    parts: list[str] = [str(episode.get("task", ""))]
    for message in episode.get("messages", []):
        if message.get("content"):
            parts.append(str(message["content"]))
        for call in message.get("tool_calls") or []:
            function = call.get("function", {})
            parts.append(str(function.get("name", "")))
            parts.append(json.dumps(function.get("arguments", {}), sort_keys=True))
    return "\n".join(parts)


def episode_hash(episode: dict[str, Any]) -> str:
    return hashlib.sha256(episode_text(episode).encode("utf-8")).hexdigest()


# --------------------------------------------------------------------------------------
# Filters
# --------------------------------------------------------------------------------------

def tool_calls_of(episode: dict[str, Any]) -> list[tuple[str, str]]:
    """(name, canonical arguments) for every call in the episode, in order."""
    out: list[tuple[str, str]] = []
    for message in episode.get("messages", []):
        for call in message.get("tool_calls") or []:
            function = call.get("function", {})
            out.append((str(function.get("name", "")),
                        json.dumps(function.get("arguments", {}), sort_keys=True)))
    return out


def shape_reason(episode: dict[str, Any], args: argparse.Namespace) -> str | None:
    """Why this episode is not worth imitating, or None if it is."""
    outcome = episode.get("outcome") or {}
    if not args.include_failures and outcome.get("passed") is not True:
        return "outcome-not-passed"
    if args.drop_lossy and episode.get("lossy"):
        return "lossy-reconstruction"
    if not episode.get("tools"):
        return "no-tool-list"
    calls = tool_calls_of(episode)
    if len(calls) < args.min_tool_calls:
        return "too-few-tool-calls"
    steps = outcome.get("steps")
    if args.max_steps and isinstance(steps, int) and steps > args.max_steps:
        return "over-step-budget"
    errors = outcome.get("tool_errors")
    if isinstance(errors, int) and errors > args.max_tool_errors:
        return "too-many-tool-errors"
    if args.drop_repeated_calls and len(calls) != len(set(calls)):
        return "repeated-identical-call"
    names = {name for name, _ in calls}
    declared = {t.get("function", {}).get("name") for t in episode["tools"]}
    if not names <= declared:
        return "call-to-undeclared-tool"
    if len(episode.get("messages", [])) > args.max_messages:
        return "too-many-messages"
    return None


# --------------------------------------------------------------------------------------
# Output shapes
# --------------------------------------------------------------------------------------

def trl_row(episode: dict[str, Any], keep_system: bool) -> dict[str, Any]:
    """TRL conversational language modelling, with the tools column the SFT trainer reads."""
    messages = [m for m in episode["messages"] if keep_system or m.get("role") != "system"]
    cleaned: list[dict[str, Any]] = []
    for message in messages:
        row: dict[str, Any] = {"role": message["role"]}
        if message.get("content") is not None:
            row["content"] = message["content"]
        if message.get("tool_calls"):
            row["tool_calls"] = [
                {"type": "function",
                 "function": {"name": c["function"]["name"],
                              "arguments": c["function"].get("arguments", {})}}
                for c in message["tool_calls"]
            ]
            row.setdefault("content", "")
        if message.get("name"):
            row["name"] = message["name"]
        cleaned.append(row)
    return {"messages": cleaned, "tools": episode["tools"]}


def mlx_row(row: dict[str, Any]) -> dict[str, Any]:
    """The same episode with tool-call arguments as JSON strings and an id per call,
    which is the shape mlx-lm's documented tools example uses."""
    messages = []
    for index, message in enumerate(row["messages"]):
        copy = dict(message)
        if copy.get("tool_calls"):
            copy["tool_calls"] = [
                {"id": f"call_{index}_{n}", "type": "function",
                 "function": {"name": c["function"]["name"],
                              "arguments": json.dumps(c["function"]["arguments"])}}
                for n, c in enumerate(message["tool_calls"])
            ]
        messages.append(copy)
    return {"messages": messages, "tools": row["tools"]}


# --------------------------------------------------------------------------------------
# Main
# --------------------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--in", dest="input", default="clean/trajectories.jsonl")
    parser.add_argument("--out-dir", default=".")
    parser.add_argument("--tasks", default=None,
                        help="the evaluation suite to decontaminate against, in Part 26's shape")
    parser.add_argument("--include-failures", action="store_true",
                        help="keep episodes whose outcome was not a pass; read the lesson first")
    parser.add_argument("--keep-lossy", dest="drop_lossy", action="store_false",
                        help="keep episodes rebuilt from logs with no assistant text")
    parser.add_argument("--min-tool-calls", type=int, default=1)
    parser.add_argument("--max-steps", type=int, default=0, help="0 disables the step filter")
    parser.add_argument("--max-tool-errors", type=int, default=0)
    parser.add_argument("--drop-repeated-calls", action="store_true",
                        help="drop an episode that made the same call with the same arguments twice")
    parser.add_argument("--max-messages", type=int, default=80)
    parser.add_argument("--drop-system", dest="keep_system", action="store_false",
                        help="leave the system prompt out of the training rows")
    parser.add_argument("--near-duplicate", type=float, default=0.8,
                        help="containment above this counts as a duplicate; 0 disables")
    parser.add_argument("--contamination", type=float, default=0.6,
                        help="containment against an evaluation task above this is contamination")
    parser.add_argument("--n", type=int, default=8, help="n-gram size for both comparisons")
    parser.add_argument("--valid-fraction", type=float, default=0.15)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--report", default=None)
    parser.add_argument("--print-example", action="store_true",
                        help="print one converted training row and exit")
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--notes", default=None)
    args = parser.parse_args()

    path = Path(args.input)
    if not path.is_file():
        sys.exit(f"{path} does not exist; run scrub-trajectories.py first")
    episodes = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()
                if line.strip().startswith("{")]
    if not episodes:
        sys.exit(f"{path} holds no episodes")

    rejected: Counter[str] = Counter()
    examples: dict[str, list[str]] = defaultdict(list)

    def reject(episode: dict[str, Any], reason: str) -> None:
        rejected[reason] += 1
        if len(examples[reason]) < 5:
            examples[reason].append(f"{episode.get('id')}: {str(episode.get('task', ''))[:120]}")

    # 1. Shape and outcome.
    kept: list[dict[str, Any]] = []
    for episode in episodes:
        reason = shape_reason(episode, args)
        if reason:
            reject(episode, reason)
        else:
            kept.append(episode)

    # 2. Exact duplicates.
    seen_hash: set[str] = set()
    unique: list[dict[str, Any]] = []
    for episode in kept:
        digest = episode_hash(episode)
        if digest in seen_hash:
            reject(episode, "exact-duplicate")
            continue
        seen_hash.add(digest)
        unique.append(episode)

    # 3. Near-duplicates, compared against what has already been accepted.
    deduped: list[dict[str, Any]] = []
    accepted_grams: list[set[tuple[str, ...]]] = []
    if args.near_duplicate > 0:
        for episode in unique:
            grams = ngrams(words(episode_text(episode)), args.n)
            if any(containment(grams, other) >= args.near_duplicate for other in accepted_grams):
                reject(episode, "near-duplicate")
                continue
            accepted_grams.append(grams)
            deduped.append(episode)
    else:
        deduped = unique

    # 4. Decontamination against the suite this fine-tune will be measured with.
    eval_ids: set[str] = set()
    eval_grams: list[tuple[str, set[tuple[str, ...]]]] = []
    if args.tasks:
        suite = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
        for task in suite.get("tasks", []):
            eval_ids.add(str(task.get("id")))
            eval_grams.append((str(task.get("id")), ngrams(words(str(task.get("task", ""))), args.n)))

    clean: list[dict[str, Any]] = []
    contaminated: list[dict[str, str]] = []
    for episode in deduped:
        if str(episode.get("task_id")) in eval_ids:
            reject(episode, "same-task-id-as-evaluation")
            contaminated.append({"id": str(episode.get("id")), "why": "task id is in the suite"})
            continue
        grams = ngrams(words(str(episode.get("task", ""))), args.n)
        hit = next((tid for tid, other in eval_grams
                    if containment(other, grams) >= args.contamination), None)
        if hit:
            reject(episode, "overlaps-evaluation-task")
            contaminated.append({"id": str(episode.get("id")),
                                 "why": f"overlaps evaluation task {hit}"})
            continue
        clean.append(episode)

    if not clean:
        sys.exit("nothing survived the filters. Read filter-report.json: the commonest cause "
                 "is that no episode passed, and the second is that every task in the "
                 "collection set is also in the evaluation suite.")

    # 5. Split by task, so that no task id appears in both halves. A model that saw the
    #    same task in training is not being evaluated on it.
    by_task: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for episode in clean:
        by_task[str(episode.get("task_id"))].append(episode)
    task_ids = sorted(by_task)
    random.Random(args.seed).shuffle(task_ids)
    want_valid = max(1, round(len(task_ids) * args.valid_fraction)) if len(task_ids) > 1 else 0
    valid_ids = set(task_ids[:want_valid])
    train_episodes = [e for tid in task_ids if tid not in valid_ids for e in by_task[tid]]
    valid_episodes = [e for tid in task_ids if tid in valid_ids for e in by_task[tid]]
    if not valid_episodes:
        print("WARNING: every episode came from one task, so there is no held-out split. "
              "Collect trajectories from more tasks before you believe any number here.")

    rows_train = [trl_row(e, args.keep_system) for e in train_episodes]
    rows_valid = [trl_row(e, args.keep_system) for e in valid_episodes]

    if args.print_example:
        if not rows_train:
            sys.exit("no training rows to print")
        print(json.dumps(rows_train[0], indent=2))
        return

    out_dir = Path(args.out_dir)
    (out_dir / "data").mkdir(parents=True, exist_ok=True)
    (out_dir / "data-mlx").mkdir(parents=True, exist_ok=True)
    for name, rows in (("train", rows_train), ("valid", rows_valid)):
        (out_dir / "data" / f"{name}.jsonl").write_text(
            "".join(json.dumps(r) + "\n" for r in rows), encoding="utf-8")
        (out_dir / "data-mlx" / f"{name}.jsonl").write_text(
            "".join(json.dumps(mlx_row(r)) + "\n" for r in rows), encoding="utf-8")

    assistant_turns = sum(1 for r in rows_train for m in r["messages"] if m["role"] == "assistant")
    calls = sum(len(m.get("tool_calls") or []) for r in rows_train for m in r["messages"])
    report = {
        "input": args.input,
        "episodes_in": len(episodes),
        "train_rows": len(rows_train),
        "valid_rows": len(rows_valid),
        "train_tasks": len(task_ids) - len(valid_ids),
        "valid_tasks": len(valid_ids),
        "assistant_turns_train": assistant_turns,
        "tool_calls_train": calls,
        "rejected_by_reason": dict(sorted(rejected.items())),
        "rejected_examples": {k: v for k, v in sorted(examples.items())},
        "contaminated": contaminated,
        "settings": {"include_failures": args.include_failures, "max_steps": args.max_steps,
                     "max_tool_errors": args.max_tool_errors,
                     "drop_repeated_calls": args.drop_repeated_calls,
                     "near_duplicate": args.near_duplicate, "contamination": args.contamination,
                     "n": args.n, "valid_fraction": args.valid_fraction, "seed": args.seed,
                     "tasks": args.tasks},
    }
    report_path = Path(args.report) if args.report else out_dir / "filter-report.json"
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")

    print(f"episodes in:      {len(episodes)}")
    print(f"train rows:       {len(rows_train)}   from {report['train_tasks']} task(s)")
    print(f"valid rows:       {len(rows_valid)}   from {report['valid_tasks']} task(s)")
    print(f"assistant turns:  {assistant_turns}")
    print(f"tool calls:       {calls}")
    if rejected:
        print("\nrejected by reason")
        for reason, count in sorted(rejected.items()):
            print(f"  {reason:<28} {count:>5}")
    print(f"\nreport written to {report_path}")
    print("Read five rejected examples and five accepted ones before you train. The filter "
          "decides what the model imitates and it is easier to fix here than afterwards.")

    if args.labbook:
        record = agentlog.record(
            labbook=args.labbook,
            lab="part-27/trajectories-to-sft",
            model=None,
            dataset={"path": str(out_dir / "data" / "train.jsonl"),
                     "sha256": agentlog.file_sha256(out_dir / "data" / "train.jsonl"),
                     "train_examples": len(rows_train), "validation_examples": len(rows_valid)},
            data_lineage=agentlog.lineage(trajectories=args.input, filter_report=str(report_path),
                                          evaluation_suite=args.tasks),
            hyperparameters=report["settings"],
            seed=args.seed, losses=None,
            scores={"train_rows": len(rows_train), "valid_rows": len(rows_valid),
                    "rejected": sum(rejected.values())},
            config_path=__file__, notes=args.notes,
        )
        print(f"recorded run {record['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()