"""Freeze user-split labeled JSONL before training; no runtime filtering."""
import argparse
import json
import sys
import unicodedata
from collections import Counter
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from transformers import AutoConfig, AutoTokenizer
from mugi_decision.encoding import DecisionEncoder, TooLongError
from mugi_decision.frozen import SCHEMA, encode_frozen, file_sha256, model_encoding_signature
from mugi_decision.runtime import save_json


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--model', default=str(Path(__file__).resolve().parents[1]))
    p.add_argument('--train', type=Path, required=True)
    p.add_argument('--validation', type=Path, required=True)
    p.add_argument('--test', type=Path)
    p.add_argument('--output', type=Path, required=True)
    p.add_argument('--max-length', type=int, default=1024)
    p.add_argument('--order-variants', type=int, default=2)
    p.add_argument('--seed', type=int, default=42)
    args = p.parse_args()
    if min(args.max_length, args.order_variants) < 1:
        p.error('Length and number of variants must be positive')
    if args.output.exists() and any(args.output.iterdir()):
        p.error('Output is not empty; use a new dataset directory')
    cfg = AutoConfig.from_pretrained(args.model, trust_remote_code=True, local_files_only=True)
    tokenizer = AutoTokenizer.from_pretrained(args.model, config=cfg, local_files_only=True)
    encoder = DecisionEncoder(tokenizer, cfg.include_enum, args.max_length)
    args.output.mkdir(parents=True, exist_ok=True)
    files, ids, stems = {}, set(), {}
    # Check provided heldout splits first. Repeated material across splits is an error.
    for split, path in [('test', args.test), ('validation', args.validation), ('train', args.train)]:
        if path is None:
            continue
        counts, skipped, rows = Counter(), 0, 0
        dest = args.output / (split + '.jsonl')
        with path.open(encoding='utf-8-sig') as src, dest.open('w', encoding='utf-8', newline='\n') as dst:
            for number, line in enumerate(src, 1):
                if not line.strip():
                    continue
                r = json.loads(line)
                if not isinstance(r.get('id'), str) or not r['id'] or r['id'] in ids:
                    raise ValueError(f'{path}:{number}: unique nonempty id required')
                ids.add(r['id'])
                if not isinstance(r.get('language'), str) or not r['language']:
                    raise ValueError(f'{path}:{number}: language required')
                if not isinstance(r.get('source'), dict) or not isinstance(r['source'].get('dataset'), str):
                    raise ValueError(f'{path}:{number}: source.dataset required')
                stem = ' '.join(unicodedata.normalize('NFKC', r['description']).casefold().split())
                if stem in stems:
                    raise ValueError(f'{path}:{number}: repeated description, first seen in {stems[stem]}; deduplicate/group-split material first')
                stems[stem] = split
                try:
                    frozen = encode_frozen(r, encoder, args.order_variants if split == 'train' else 1, args.seed)
                except TooLongError:
                    if split != 'train':
                        raise
                    skipped += 1
                    continue
                dst.write(json.dumps(frozen, ensure_ascii=False, separators=(',', ':')) + '\n')
                counts[r['language']] += 1
                rows += 1
        if not rows:
            raise ValueError(f'{split}: no retained records')
        files[split] = {'path': dest.name, 'rows': rows, 'sha256': file_sha256(dest),
                        'languages': dict(counts), 'offline_rejected_overlength': skipped}
        print(split, json.dumps(files[split], ensure_ascii=False), flush=True)
    save_json(args.output/'manifest.json', {'schema': SCHEMA, 'max_length': args.max_length,
        'encoding_signature': model_encoding_signature(args.model), 'files': files,
        'seed': args.seed, 'order_variants': args.order_variants})


if __name__ == '__main__':
    main()