"""Generate new multilingual slot-preserving templates, with complete provenance."""
import argparse
import concurrent.futures
import json
import re
import time
import urllib.request
from pathlib import Path
from tinyquery.recipes import RECIPES, LANGUAGES


def slots(s):
    return set(re.findall(r'\{([a-z_]+)\}', s))


def generate(op, seed, base, count):
    source = RECIPES[op]
    required = slots(source)
    prompt = (
        'Create diverse natural user requests that mean EXACTLY the same thing as the source. '
        'Return a JSON object with keys en, noisy_en, hi, hinglish; each value is a list of '
        f'{count} different strings. en: natural English, noisy_en: imperfect English that '
        'still preserves every constraint, hi: natural Devanagari Hindi, hinglish: Romanized Hindi. '
        'Preserve ALL literal {placeholder} slots exactly, including braces. Do not add slots. '
        'Do not replace slots with example values. Do not add or remove conditions, columns, '
        'sorting, limits, negation, or inclusivity. Vary sentence structure and tone. '
        'Do not provide SQL or explanations. Source: ' + source
    )
    payload = {'model': 'Qwen/Qwen3.8-27B-FP8', 'messages': [{'role':'user','content':prompt}],
               'temperature':0.8,'max_tokens':4096,'seed':seed,
               'chat_template_kwargs':{'enable_thinking':False},
               'response_format':{'type':'json_object'}}
    request = urllib.request.Request(base+'/chat/completions', data=json.dumps(payload).encode(),
                                     headers={'Content-Type':'application/json'})
    started = time.time()
    with urllib.request.urlopen(request, timeout=240) as response:
        result = json.load(response)
    raw = result['choices'][0]['message']['content']
    parsed = json.loads(raw)
    accepted = {}; rejected = []
    for lang in LANGUAGES:
        accepted[lang] = []
        for text in parsed.get(lang, []):
            if not isinstance(text,str) or slots(text) != required or len(text)>900:
                rejected.append({'language':lang,'text':text,'reason':'slot_or_length'})
                continue
            if lang == 'hi' and not re.search('[\u0900-\u097f]', text):
                rejected.append({'language':lang,'text':text,'reason':'missing_devanagari'})
                continue
            if text not in accepted[lang]: accepted[lang].append(text)
    return {'op':op,'source':source,'templates':accepted,'rejected':rejected,
            'teacher':payload['model'],'revision':'017b9c7af6b5689d5dd426a76e0bc077eb5ca20a',
            'request':payload,'raw_response':raw,'usage':result.get('usage'),
            'elapsed_seconds':time.time()-started,'finish_reason':result['choices'][0]['finish_reason']}


def main():
    parser=argparse.ArgumentParser()
    parser.add_argument('--out',default='data/tinyquery/templates.jsonl')
    parser.add_argument('--base',default='http://127.0.0.1:18000/v1')
    parser.add_argument('--count',type=int,default=12)
    parser.add_argument('--workers',type=int,default=16)
    args=parser.parse_args()
    path=Path(args.out); path.parent.mkdir(parents=True,exist_ok=True)
    done=set()
    if path.exists():
        done={json.loads(line)['op'] for line in path.read_text().splitlines()}
    started=time.time()
    with path.open('a') as output, concurrent.futures.ThreadPoolExecutor(args.workers) as pool:
        jobs={pool.submit(generate,op,900+i,args.base,args.count):op
              for i,op in enumerate(RECIPES) if op not in done}
        for future in concurrent.futures.as_completed(jobs):
            op=jobs[future]
            try:
                result=future.result()
                output.write(json.dumps(result,ensure_ascii=False)+'\n'); output.flush()
                print(json.dumps({'op':op,'accepted':{k:len(v) for k,v in result['templates'].items()},
                                  'seconds':round(time.time()-started,1)}),flush=True)
            except Exception as exc:
                print(json.dumps({'op':op,'error':str(exc)}),flush=True)


if __name__ == '__main__': main()