"""Teach exact project-scope and discovery-table copying across identifier formats."""
import argparse
import hashlib
import json
from pathlib import Path
import random
import re
import shutil
import string
from tokenizers import Tokenizer
from sqlglot import Tokenizer as SQLTokenizer
from tinyquery.data import compact,ddls,serialize


def main():
    p=argparse.ArgumentParser();p.add_argument('--source',default='data/tinyquery-v4')
    p.add_argument('--seed-source',default='data/tinyquery');p.add_argument('--out',default='data/tinyquery-v5')
    args=p.parse_args();source=Path(args.source);seeds=Path(args.seed_source);out=Path(args.out);out.mkdir(parents=True,exist_ok=True)
    for name in ['validation.jsonl','test.jsonl','manual.jsonl','tokenizer.json']:shutil.copy2(source/name,out/name)
    tokenizer=Tokenizer.from_file(str(source/'tokenizer.json'))
    chunks=sorted({tokenizer.decode([i]) for i in range(tokenizer.get_vocab_size()) if re.fullmatch('[a-z]{1,8}',tokenizer.decode([i]))})
    rng=random.Random(99335);counts={'replay':0,'scope':0,'discovery':0};seen=set()
    with (out/'train.jsonl').open('w') as stream:
        def emit(row,kind):
            key=hashlib.sha256(row['prompt'].encode()).hexdigest()
            if key in seen:return
            seen.add(key);counts[kind]+=1;stream.write(json.dumps(row,ensure_ascii=False)+'\n')
        for line in (source/'train.jsonl').open():emit(json.loads(line),'replay')
        for line in (seeds/'train.jsonl').open():
            row=json.loads(line)
            if row['target']['action']!='call':continue
            a=row['target']['arguments'];kind=None
            if 'project_id' in a:
                style=rng.randrange(3)
                project=(''.join(rng.choices(string.ascii_lowercase,k=20)) if style==0 else
                         ''.join(rng.choices(chunks,k=4)) if style==1 else
                         '-'.join(''.join(rng.choices(string.hexdigits[:16],k=k)) for k in [8,4,4,4,12]))
                row['context']['project_id']=project;a['project_id']=project;kind='scope'
            if 'table' in a:
                old=a['table']
                while True:
                    new=''.join(rng.choices(chunks,k=4))
                    if len(new)>=6 and new.upper() not in SQLTokenizer.KEYWORDS:break
                row['question']=re.sub(r'(?<![A-Za-z0-9_])'+re.escape(old)+r'(?![A-Za-z0-9_])',new,row['question'])
                row['slots']['table']=new;a['table']=new
                if row['context']['schema']:row['context']['schema']=ddls(row['slots'])
                kind='discovery'
            if kind:
                row['id']+='_scope_99335';row['sample_weight']=4
                row['provenance']+='; exact copying of arbitrary project scope and discovery table identifiers'
                row['response']=compact(row['target']);row['prompt']=serialize(row['context'],row['question'])
                emit(row,kind)
    (out/'grounding-stats.json').write_text(json.dumps(counts,indent=2));print(counts)


if __name__=='__main__':main()