#!/usr/bin/env python3
"""Merge curated + augmented router data into a balanced train/val SFT dataset."""
import json
import os
import random

ROOT = "/workspace/TASK18/data/router"
CUR = os.path.join(ROOT, "router_train.jsonl")        # 12093 grounded single-intent
AUG = os.path.join(ROOT, "router_augmented.jsonl")    # alert + compound

random.seed(42)


def load(p):
    return [json.loads(l) for l in open(p, encoding="utf-8")]


cur = load(CUR)
aug = load(AUG)

# split augmented: hold out some alert & compound for val
random.shuffle(aug)
n_aug_val = max(30, int(len(aug) * 0.2))
aug_val, aug_train = aug[:n_aug_val], aug[n_aug_val:]

# from curated, keep a representative grounded sample for train, plus a small val set
random.shuffle(cur)
val_cur = cur[:300]
pool_cur = cur[300:]

# limit grounded pool in train to balance (alert ~200, keep grounded ~ few thousand)
GROUND_TRAIN_CAP = 3000
train_cur = pool_cur[:GROUND_TRAIN_CAP]

train = aug_train + train_cur
val = aug_val + val_cur

# dedup train by utterance
seen = set()
train_dedup = []
for r in train:
    u = r["utterance"]
    if u in seen:
        continue
    seen.add(u)
    train_dedup.append(r)

from collections import Counter


def mode_of(r):
    return json.loads(r["messages"][1]["content"]).get("mode", "grounded")


print("train:", len(train_dedup), Counter(mode_of(r) for r in train_dedup))
print("val  :", len(val), Counter(mode_of(r) for r in val))

for name, recs in (("router_sft_train.jsonl", train_dedup), ("router_sft_val.jsonl", val)):
    with open(os.path.join(ROOT, name), "w", encoding="utf-8") as f:
        for r in recs:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print("saved ->", os.path.join(ROOT, name))