import argparse
import ast
import json
import platform
import re
import time
from collections import Counter
from importlib.metadata import version
from common import ROOT, inside, rows, save_json, append, stable_seed, permit_compute
from evaluate import model_fingerprint,recover_generations
from inference_queue import completed_outputs
from resources import snapshot
from code_sandbox import check_code


def extract_code(text):
    text = text.split('</think>')[-1].strip()
    blocks = re.findall(r'```(?:python|py)?\s*\n(.*?)```',text,flags=re.S|re.I)
    if blocks:
        return blocks[0].strip()
    if text.startswith('```'):
        return '\n'.join(text.split('\n')[1:]).removesuffix('```').strip()
    return text


def public_prelude(row):
    if row['source'] != 'humaneval':
        return row.get('prelude','')
    tree = ast.parse(row['problem'])
    return '\n'.join(ast.unparse(node) for node in tree.body if isinstance(node,(ast.Import,ast.ImportFrom)))


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--model',default=str(ROOT/'models/original'))
    p.add_argument('--data',required=True)
    p.add_argument('--out',required=True)
    p.add_argument('--config',default=str(ROOT/'configs/coding_protocol.json'))
    p.add_argument('--limit',type=int)
    p.add_argument('--samples',type=int)
    args = p.parse_args()
    permit_compute()
    from transformers import AutoTokenizer
    from vllm import LLM,SamplingParams
    model,data,out = inside(args.model),inside(args.data),inside(args.out)
    cfg = json.loads(inside(args.config).read_text())
    records = rows(data)
    if args.limit:
        records = records[:args.limit]
    samples = args.samples or cfg['samples']
    out.mkdir(parents=True,exist_ok=True)
    job = {'model':str(model),'fingerprint':model_fingerprint(model),'data':str(data),
           'data_sha256':__import__('hashlib').sha256(data.read_bytes()).hexdigest(),'protocol':cfg,
           'samples':samples,'max_tokens':cfg['max_tokens'],'set':'coding','limit':args.limit,
           'runtime_versions':{x:version(x) for x in ('vllm','torch','transformers')},'platform':platform.platform(),
           'seed_scheme':'sha256(base_seed:coding/item_id:sample)31'}
    if (out/'job.json').exists() and json.loads((out/'job.json').read_text()) != job:
        raise ValueError('Existing job differs')
    save_json(out/'job.json',job)
    completed = recover_generations(out/'generations.jsonl')
    done = {(r['id'],r['sample']) for r in completed}
    if len(done)!=len(completed):
        raise ValueError('Duplicate generations')
    pending = [(row,sample) for row in records for sample in range(samples) if (row['id'],sample) not in done]
    if pending:
        tok = AutoTokenizer.from_pretrained(model,local_files_only=True)
        engine = LLM(model=str(model),dtype='bfloat16',seed=cfg['seed'],max_model_len=cfg['max_model_len'],
                     gpu_memory_utilization=cfg['gpu_memory_utilization'],kv_cache_memory_bytes=cfg['kv_cache_memory_bytes'],
                     max_num_seqs=cfg['max_num_seqs'],max_num_batched_tokens=cfg['max_num_batched_tokens'],
                     enable_prefix_caching=False,enforce_eager=False,limit_mm_per_prompt={'image':0,'video':0})
        append(out/'resources.jsonl',snapshot('code_engine_initialized'))
        prompts,params = [],[]
        for row,sample in pending:
            prompt = tok.apply_chat_template([{'role':'user','content':cfg['instruction']+'\n\n'+row['problem']}],
                                            tokenize=False,add_generation_prompt=True,enable_thinking=cfg['enable_thinking'])
            prompts.append(prompt)
            params.append(SamplingParams(n=1,seed=stable_seed('coding/'+row['id'],sample,cfg['seed']),
                                         temperature=cfg['temperature'],top_p=cfg['top_p'],top_k=cfg['top_k'],
                                         presence_penalty=cfg['presence_penalty'],max_tokens=cfg['max_tokens']))
        start = time.time()
        for index,output in completed_outputs(engine,prompts,params,stream=True):
            row,sample = pending[index]
            generation = output.outputs[0]
            if generation.finish_reason not in ('stop','length'):
                raise RuntimeError('Incomplete code generation')
            code = public_prelude(row)+'\n'+extract_code(generation.text)
            verdict = check_code(code,row['entry_point'],row['test_code'])
            result = {'id':row['id'],'source':row['source'],'sample':sample,
                      'seed':stable_seed('coding/'+row['id'],sample,cfg['seed']),'correct':verdict['passed'],
                      'verdict':verdict,'code':code,'text':generation.text,'tokens':len(generation.token_ids),
                      'prompt':prompts[index],
                      'prompt_token_ids':tok.encode(prompts[index],add_special_tokens=False),
                      'token_ids':list(generation.token_ids),'finish_reason':generation.finish_reason,
                      'queue_elapsed_seconds':time.time()-start}
            append(out/'generations.jsonl',result)
            completed.append(result)
            print(json.dumps({k:result[k] for k in ('id','sample','correct','tokens','finish_reason')},ensure_ascii=False),flush=True)
        append(out/'resources.jsonl',snapshot('code_queue_finished'))
    n = len(completed)
    if not n:
        raise ValueError('No coding results')
    by_source = {}
    for source in sorted({r['source'] for r in completed}):
        subset = [r for r in completed if r['source']==source]
        by_source[source] = {'correct':sum(r['correct'] for r in subset),'generations':len(subset),
                             'accuracy':sum(r['correct'] for r in subset)/len(subset)}
    by_sample = {}
    for sample in range(samples):
        subset = [r for r in completed if r['sample']==sample]
        by_sample[str(sample)] = {'correct':sum(r['correct'] for r in subset),'generations':len(subset)}
    task_results = {}
    for result in completed:
        task_results.setdefault(result['id'],[]).append(result['correct'])
    save_json(out/'summary.json',{'generations':n,'planned_generations':len(records)*samples,
                                'complete':n==len(records)*samples,
                                'problems':len(task_results),'samples_per_problem':samples,
                                'correct':sum(r['correct'] for r in completed),'accuracy':sum(r['correct'] for r in completed)/n,
                                'mean_tokens':sum(r['tokens'] for r in completed)/n,'by_source':by_source,
                                'by_sample':by_sample,
                                'tasks_correct_all_samples':sum(len(v)==samples and all(v) for v in task_results.values()),
                                'verdict_counts':dict(Counter('passed' if r['correct'] else r['verdict'].get('error_type','unknown') for r in completed)),
                                'guard_rejections':sum(not r['correct'] and r['verdict'].get('error','').startswith('unsupported_') for r in completed),
                                'cut_share':sum(r['finish_reason']=='length' for r in completed)/n,
                                'scope':'Compact fixed HumanEval/MBPP subsets, custom chat prompt and restricted stdlib sandbox; not official leaderboard scores.'})


if __name__ == '__main__':
    main()