"""Pinned CodeSOTA pilot protocol. Never imports or executes remote Python code."""
from __future__ import annotations
import hashlib, importlib.metadata, json, os, platform, re, subprocess, sys, time, secrets
from pathlib import Path

def code_test(program, image):
    """Execute candidate + hidden checks inside a networkless disposable container."""
    name='codesota-eval-'+secrets.token_hex(8)
    cmd=['docker','run','--name',name,'--rm','--network=none','--read-only','--cap-drop=ALL','--security-opt=no-new-privileges','--pids-limit=32','--memory=128m','--cpus=1','--user=65534:65534','--tmpfs=/tmp:rw,noexec,nosuid,size=16m','-i',image,'python','-I','-']
    try:
        # Docker has an in-container alarm as well as a host deadline.
        guarded='import signal\nsignal.alarm(4)\n'+program
        p=subprocess.run(cmd,input=guarded,text=True,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,timeout=12)
        return p.returncode==0
    except subprocess.TimeoutExpired:return False
    finally:
        subprocess.run(['docker','rm','--force',name],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,timeout=15)

def dataset_for(t):
    if t['adapter']=='commonsense':
        # Read the pinned parquet directly: its newer HF feature metadata uses
        # List, which older datasets versions cannot parse. Arrow preserves row order.
        from huggingface_hub import hf_hub_download
        import pyarrow.parquet as pq
        file=hf_hub_download(t['dataset'],f"data/{t['split']}-00000-of-00001.parquet",repo_type='dataset',revision=t['dataset_revision'])
        return (row for batch in pq.ParquetFile(file).iter_batches(batch_size=128) for row in batch.to_pylist())
    from datasets import load_dataset
    return load_dataset(t['dataset'],t['dataset_config'],split=t['split'],revision=t['dataset_revision'],streaming=True)

def evaluate(job, output):
    import numpy as np
    import torch
    from datasets import load_dataset
    from transformers import AutoTokenizer, AutoModelForCausalLM
    from huggingface_hub import HfApi
    t=job['task'];adapter=t['adapter'];revision=job['revision'];model_id=job['model']
    torch.manual_seed(42);np.random.seed(42);torch.set_num_threads(4)
    device='cuda:0' if torch.cuda.is_available() else 'cpu'
    dataset=dataset_for(t)
    rows=[]
    for index,row in enumerate(dataset):
        if adapter=='entailment' and row['label']<0:continue
        rows.append((index,row))
        if len(rows)==t['limit']:break
    if len(rows)!=t['limit']:raise ValueError('Dataset did not yield the declared sample count')
    samples=[];started=time.monotonic();special={};peak=0
    def save(index,reference,prediction,score=None,**extra):
        samples.append(dict(id=str(index),reference=reference,prediction=prediction,score=score,**extra))
    def load_causal():
        tok=AutoTokenizer.from_pretrained(model_id,revision=revision,trust_remote_code=False)
        model=AutoModelForCausalLM.from_pretrained(model_id,revision=revision,trust_remote_code=False,torch_dtype=torch.float16 if device.startswith('cuda') else torch.float32).to(device).eval()
        return tok,model
    if adapter in ('reading','commonsense','sentiment','entailment','math','code'):
        tokenizer,model=load_causal()
        def choice(prompt,choices):
            scores=[]
            for c in choices:
                # Tokenize the prompt and answer independently to define an unambiguous boundary.
                prefix=tokenizer.encode(prompt,add_special_tokens=False)[-1800:]
                suffix=tokenizer.encode(' '+c,add_special_tokens=False)
                ids=torch.tensor([prefix+suffix],device=device)
                with torch.inference_mode(): logits=model(ids).logits[0,len(prefix)-1:-1].float()
                targets=ids[0,len(prefix):]
                score=torch.log_softmax(logits,dim=-1).gather(1,targets[:,None]).mean().item()
                scores.append(score)
            return int(np.argmax(scores)),scores
        def generate(prompt,max_tokens):
            ids=tokenizer(prompt,return_tensors='pt',truncation=True,max_length=1800).to(device)
            with torch.inference_mode(): out=model.generate(**ids,max_new_tokens=max_tokens,do_sample=False,pad_token_id=tokenizer.eos_token_id)
            return tokenizer.decode(out[0,ids['input_ids'].shape[1]:],skip_special_tokens=True)
        if adapter=='code':
            image=os.environ.get('CODESOTA_SANDBOX_IMAGE')
            if not image or '@sha256:' not in image:raise ValueError('Pinned Docker sandbox image required')
            subprocess.run(['docker','image','inspect',image],check=True,stdout=subprocess.DEVNULL)
            special['sandbox_image']=image
        for index,row in rows:
            before=time.monotonic()
            if adapter=='reading':
                prompt=f"Passage: {row['passage']}\nQuestion: {row['question']}\nAnswer:";choices=['no','yes'];target=int(row['answer'])
            elif adapter=='commonsense':prompt=row['ctx'];choices=row['endings'];target=int(row['label'])
            elif adapter=='sentiment':prompt=f"Review: {row['sentence']}\nSentiment:";choices=['negative','positive'];target=int(row['label'])
            elif adapter=='entailment':
                prompt=f"Premise: {row['premise']}\nHypothesis: {row['hypothesis']}\nRelationship:";choices=['entailment','neutral','contradiction'];target=int(row['label'])
            elif adapter=='math':
                raw=generate(f"Question: {row['question']}\nSolve the problem. End with #### followed by the numerical answer.\nAnswer:",256)
                matches=re.findall(r'####\s*([-+]?\d[\d,]*(?:\.\d+)?)',raw)
                guess=matches[-1].replace(',','') if matches else None
                answer=row['answer'].split('####')[-1].strip().replace(',','')
                try:correct=float(guess)==float(answer)
                except (TypeError,ValueError):correct=False
                save(index,answer,raw,float(correct),extracted=guess,seconds=time.monotonic()-before);continue
            elif adapter=='code':
                completion=generate(row['prompt'],256)
                # Completion-only protocol. Stop before the next top-level definition / commentary.
                completion=re.split(r'\n(?=def |class |if __name__|print\(|#)',completion)[0]
                code=row['prompt']+completion+'\n'+row['test']+f"\ncheck({row['entry_point']})\n"
                correct=code_test(code,special['sandbox_image'])
                save(row['task_id'], 'passes hidden unit tests',completion,float(correct),seconds=time.monotonic()-before);continue
            pred,likelihoods=choice(prompt,choices)
            save(index,target,pred,float(pred==target),choices=choices,log_likelihood_mean=likelihoods,seconds=time.monotonic()-before)
        score=float(np.mean([s['score'] for s in samples]))
    elif adapter=='similarity':
        from sentence_transformers import SentenceTransformer
        from scipy.stats import spearmanr
        model=SentenceTransformer(model_id,revision=revision,device=device,trust_remote_code=False)
        left=model.encode([r['sentence1'] for _,r in rows],normalize_embeddings=True,convert_to_numpy=True)
        right=model.encode([r['sentence2'] for _,r in rows],normalize_embeddings=True,convert_to_numpy=True)
        pred=np.sum(left*right,axis=1)
        refs=[float(r['score']) for _,r in rows]
        score=float(spearmanr(pred,refs).statistic)
        for (i,r),p in zip(rows,pred):save(i,float(r['score']),float(p))
    elif adapter=='ocr':
        from transformers import TrOCRProcessor, VisionEncoderDecoderModel
        from jiwer import cer
        processor=TrOCRProcessor.from_pretrained(model_id,revision=revision)
        model=VisionEncoderDecoderModel.from_pretrained(model_id,revision=revision).to(device).eval()
        refs=[];preds=[]
        for index,row in rows:
            image=row['image'].convert('RGB');ref=row['text']
            pixels=processor(images=image,return_tensors='pt').pixel_values.to(device)
            with torch.inference_mode():out=model.generate(pixels,max_new_tokens=128,do_sample=False)
            pred=processor.batch_decode(out,skip_special_tokens=True)[0]
            # Case and punctuation preserved; strip only outer whitespace.
            ref=ref.strip();pred=pred.strip();refs.append(ref);preds.append(pred);save(index,ref,pred,cer(ref,pred))
        score=float(cer(refs,preds))
    elif adapter=='asr':
        from transformers import WhisperProcessor, WhisperForConditionalGeneration
        from jiwer import wer
        from datasets import Audio
        dataset=dataset.cast_column('audio',Audio(sampling_rate=16000))
        rows=[]
        for i,row in enumerate(dataset):
            rows.append((i,row))
            if len(rows)==t['limit']:break
        processor=WhisperProcessor.from_pretrained(model_id,revision=revision)
        model=WhisperForConditionalGeneration.from_pretrained(model_id,revision=revision).to(device).eval()
        from transformers.models.whisper.english_normalizer import BasicTextNormalizer
        normalize=BasicTextNormalizer();refs=[];preds=[]
        for index,row in rows:
            audio=row['audio'];inputs=processor(audio['array'],sampling_rate=audio['sampling_rate'],return_tensors='pt').input_features.to(device)
            with torch.inference_mode():out=model.generate(inputs,max_new_tokens=128,do_sample=False)
            pred=normalize(processor.batch_decode(out,skip_special_tokens=True)[0]);ref=normalize(row['text'])
            refs.append(ref);preds.append(pred);save(row.get('id',index),ref,pred,wer(ref,pred))
        score=float(wer(refs,preds));special['normalizer']='transformers Whisper BasicTextNormalizer'
    elif adapter=='vision':
        from transformers import CLIPProcessor, CLIPModel
        labels=['airplane','automobile','bird','cat','deer','dog','frog','horse','ship','truck']
        prompts=[f'a photo of a {label}.' for label in labels]
        processor=CLIPProcessor.from_pretrained(model_id,revision=revision)
        model=CLIPModel.from_pretrained(model_id,revision=revision).to(device).eval()
        for index,row in rows:
            inputs=processor(text=prompts,images=row['img'].convert('RGB'),return_tensors='pt',padding=True).to(device)
            with torch.inference_mode(): pred=int(model(**inputs).logits_per_image[0].argmax())
            save(index,int(row['label']),pred,float(pred==row['label']))
        score=float(np.mean([s['score'] for s in samples]));special['class_prompts']=prompts;special['training_regime']='zero-shot CLIP; compare only within this protocol'
    else:raise ValueError('Unknown adapter')
    if torch.cuda.is_available():peak=torch.cuda.max_memory_allocated()
    duration=time.monotonic()-started
    packages={d.metadata['Name']:d.version for d in importlib.metadata.distributions()}
    result=dict(score=score,n=len(samples),model_revision=revision,dataset_revision=t['dataset_revision'],runner_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),seconds=duration,peak_gpu_bytes=peak,selection='first 32 valid examples in pinned split order',seed=42,device=device,protocol=t,measured_cost_usd=None,notes='Pilot subset, custom CodeSOTA protocol. Not an official full-benchmark score.',**special)
    output.parent.mkdir(parents=True,exist_ok=True)
    output.write_text(json.dumps(dict(result=result,samples=samples,environment=dict(python=sys.version,platform=platform.platform(),torch=torch.__version__,cuda=torch.version.cuda,gpu=torch.cuda.get_device_name() if torch.cuda.is_available() else None,packages=packages)),allow_nan=False))

if __name__=='__main__':evaluate(json.loads(Path(sys.argv[1]).read_text()),Path(sys.argv[2]))