"""
Prepare dataset for fine-tuning: read cleaned JSONL, optionally apply prompt formatting,
create train/val/test splits (stratified if labels exist), and write HF-compatible JSON/JSONL.

Usage (bash):
python scripts/prepare_dataset.py --cleaned data/cleaned/cleaned_extracted_text.jsonl --out_dir data/prepared_run

"""
import argparse
import json
import os
from collections import Counter
from sklearn.model_selection import train_test_split

# Compact prompt template for training (IDs + short names)
COMPACT_PROMPT = (
    "CATEGORIES: 1:MaliciousData 2:Ransomware 3:Marketplace 4:MalwareTools 5:ThreatActor "
    "6:VulnIntel 7:InfraOps 8:Credentials 9:SocialEng 10:General\n\nPOST:\n{txt}\n\nCATEGORY_ID:"
)


def format_prompt(text):
    return COMPACT_PROMPT.format(txt=text)


def read_jsonl(path):
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            yield json.loads(line)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--cleaned', required=True)
    parser.add_argument('--out_dir', default='data/prepared_run')
    parser.add_argument('--use_prompt', action='store_true')
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    samples = list(read_jsonl(args.cleaned))
    print(f"Loaded {len(samples)} cleaned samples")

    # If labels exist, create stratified splits
    labels = [s.get('label') for s in samples]
    has_labels = all([l is not None for l in labels])

    if has_labels:
        # map label names to numeric ids
        unique = sorted(list(set(labels)))
        label2id = {name: i for i, name in enumerate(unique)}
        for s in samples:
            s['label_id'] = label2id[s['label']]

        # stratified split
        train, temp = train_test_split(samples, test_size=0.3, stratify=[s['label_id'] for s in samples], random_state=42)
        val, test = train_test_split(temp, test_size=0.5, stratify=[s['label_id'] for s in temp], random_state=42)
    else:
        # no labels: split randomly
        train, val_test = train_test_split(samples, test_size=0.3, random_state=42)
        val, test = train_test_split(val_test, test_size=0.5, random_state=42)
        label2id = {}

    # Optionally apply prompt formatting
    def maybe_prompt(s):
        txt = s['text']
        return format_prompt(txt) if args.use_prompt else txt

    def write_jsonl(l, path):
        with open(path, 'w', encoding='utf-8') as f:
            for s in l:
                out = {
                    'text': maybe_prompt(s),
                }
                if 'label_id' in s:
                    out['label'] = s['label_id']
                f.write(json.dumps(out, ensure_ascii=False) + '\n')

    write_jsonl(train, os.path.join(args.out_dir, 'train.jsonl'))
    write_jsonl(val, os.path.join(args.out_dir, 'val.jsonl'))
    write_jsonl(test, os.path.join(args.out_dir, 'test.jsonl'))

    # Write labels mapping
    with open(os.path.join(args.out_dir, 'labels.json'), 'w', encoding='utf-8') as f:
        json.dump({'label2id': label2id}, f, ensure_ascii=False, indent=2)

    print(f"Prepared splits -> train:{len(train)} val:{len(val)} test:{len(test)}")


if __name__ == '__main__':
    main()