"""
Merge all JSONL instruction-tuning datasets into train.json and val.json
for litgpt finetune (JSON data module).

Each row must have: instruction, output, and optionally input.
"""

import json
import os
import random
import argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", type=str, default=r"D:\ASTERIZER 2026\LitGPT\OLD_DATASETS")
    parser.add_argument("--output_dir", type=str, default=r"D:\ASTERIZER 2026\LUNA\Base\Datasets\finetune")
    parser.add_argument("--val_fraction", type=float, default=0.05, help="Fraction for validation split")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    random.seed(args.seed)
    os.makedirs(args.output_dir, exist_ok=True)

    all_samples = []
    file_counts = {}

    for fname in sorted(os.listdir(args.input_dir)):
        if not fname.endswith(".jsonl"):
            continue
        fpath = os.path.join(args.input_dir, fname)
        count = 0
        with open(fpath, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                row = json.loads(line)
                # Keep only the keys litgpt expects
                sample = {
                    "instruction": row.get("instruction", ""),
                    "input": row.get("input", ""),
                    "output": row.get("output", ""),
                }
                # Skip rows with empty instruction AND empty input
                if not sample["instruction"].strip() and not sample["input"].strip():
                    continue
                # Skip rows with empty output
                if not sample["output"].strip():
                    continue
                all_samples.append(sample)
                count += 1
        file_counts[fname] = count
        print(f"  {fname}: {count} samples")

    print(f"\nTotal valid samples: {len(all_samples)}")

    # Shuffle
    random.shuffle(all_samples)

    # Split
    val_size = max(1, int(len(all_samples) * args.val_fraction))
    val_data = all_samples[:val_size]
    train_data = all_samples[val_size:]

    print(f"Train: {len(train_data)}, Val: {val_size}")

    # Write
    train_path = os.path.join(args.output_dir, "train.json")
    val_path = os.path.join(args.output_dir, "val.json")

    with open(train_path, "w", encoding="utf-8") as f:
        json.dump(train_data, f, ensure_ascii=False, indent=None)
    with open(val_path, "w", encoding="utf-8") as f:
        json.dump(val_data, f, ensure_ascii=False, indent=None)

    print(f"\nSaved: {train_path}")
    print(f"Saved: {val_path}")

    # Show a few samples
    print("\n--- Sample train entries ---")
    for s in train_data[:3]:
        print(f"  instruction: {s['instruction'][:80]}")
        print(f"  input:       {s['input'][:80]}")
        print(f"  output:      {s['output'][:80]}")
        print()

if __name__ == "__main__":
    main()