import importlib.util,json,tempfile,unittest,sys,subprocess,os
from pathlib import Path
from unittest.mock import patch
P=Path(__file__).resolve().parents[1]
s=importlib.util.spec_from_file_location('inf',P/'inference.py');inf=importlib.util.module_from_spec(s);s.loader.exec_module(inf)
Q={'q_no':1,'question':'Example?','subject':'Demo','category':'stem','choices':['a','b','c','d'],'_level':'easy','_file':'easy/demo.json'}
class PipelineTests(unittest.TestCase):
 def test_rotation_mapping_and_fallback(self):
  k=inf.kof(Q)
  with tempfile.TemporaryDirectory() as d:
   with patch.object(inf,'chat',return_value=('no answer','')) as c:
    result=inf.stage_rot([Q],['http://127.0.0.1:11434'],1,Path(d),1,'pass2','p2.jsonl',{k:'4'},{k:'3'})
   self.assertEqual(result[k],'2'); self.assertEqual(c.call_args.args[-1],300);self.assertEqual(c.call_count,1)
   with patch.object(inf,'chat',return_value=('no answer','')):
    self.assertEqual(inf.stage_rot([Q],['http://127.0.0.1:11434'],1,Path(d),2,'p4','p4.jsonl',{k:'4'},{k:'3'},False),{})
 def test_pass3_preserves_actual_sampling(self):
  with tempfile.TemporaryDirectory() as d,patch.object(inf,'chat',return_value=('Answer: 2','')) as c:
   inf.stage_pass3([Q],['http://127.0.0.1:11434'],1,Path(d),{inf.kof(Q):'4'})
   self.assertEqual(c.call_args.args[4],{'temperature':0.1,'top_p':0.95,'num_predict':4096})
 def test_transport_does_not_invent_seed(self):
  from io import BytesIO
  with patch.object(inf._OPENER,'open',return_value=BytesIO(b'{"message":{"content":"Answer: 2"}}')) as op:
   inf.chat('http://127.0.0.1:11434','s','u',True,{'temperature':0.1})
  payload=json.loads(op.call_args.args[0].data)
  self.assertNotIn('seed',payload['options'])
 def test_loader_row_order_and_duplicates(self):
  rows=[{**Q,'_file':'hard/b.json','_level':'hard'}, {**Q,'_file':'easy/a.json'}]
  self.assertEqual([q['_file'] for q in inf.prepare(rows)],['hard/b.json','easy/a.json'])
  with self.assertRaises(ValueError): inf.prepare([dict(Q),dict(Q)])
 def test_missing_or_duplicate_checkpoint_fails(self):
  with tempfile.TemporaryDirectory() as d:
   p=Path(d)/'p.jsonl'
   with self.assertRaises(FileNotFoundError):inf.load_ck(p,required=True)
   r={'q_no':1,'subject':'Demo','category':'stem','level':'easy','answer':'1'};p.write_text((json.dumps(r)+'\n')*2)
   with self.assertRaises(ValueError):inf.load_ck(p)
 def test_protect_submitted_file(self):
  old=(P/'predictions.jsonl').read_bytes()
  r=subprocess.run([sys.executable,'-B',str(P/'inference.py'),'--output',str(P/'predictions.jsonl')],capture_output=True,text=True)
  self.assertNotEqual(r.returncode,0);self.assertIn('Do not overwrite',r.stderr);self.assertEqual(old,(P/'predictions.jsonl').read_bytes())
 def test_full_orchestration_with_scripted_responses(self):
  # Tests all stage orchestration on fabricated data. No actual model or competition answers.
  with tempfile.TemporaryDirectory() as td:
   root=Path(td);data=root/'data'
   for level in ['easy','medium','hard']:
    (data/level).mkdir(parents=True)
    (data/level/'demo.json').write_text(json.dumps([{'q_no':1,'subject':'Demo','category':'stem','question':'Example?','choices':['a','b','c','d']}]))
   calls=[]
   def chat(host,system,user,think,options,timeout=420):
    calls.append(options)
    if not think:return ('Answer: 1' if options['num_predict']==260 else '2','')
    return ('Answer: 3','')
   env={'WORKDIR':str(root/'ck'),'WORKERS':'1','OLLAMA_HOSTS':'http://127.0.0.1:11434'}
   with patch.dict(os.environ,env),patch.object(sys,'argv',['inference.py','--local-data',str(data),'--output',str(root/'out.jsonl')]),patch.object(inf,'ensure_ollama'),patch.object(inf,'chat',side_effect=chat):
    inf.main()
   answers=[json.loads(l) for l in (root/'out.jsonl').read_text().splitlines()]
   self.assertEqual(len(answers),3); self.assertTrue(any(x.get('temperature')==0.1 for x in calls))
   before=len(calls)
   with patch.dict(os.environ,env),patch.object(sys,'argv',['inference.py','--local-data',str(data),'--output',str(root/'out2.jsonl')]),patch.object(inf,'ensure_ollama'),patch.object(inf,'chat',side_effect=chat):inf.main()
   self.assertEqual(len(calls),before)
if __name__=='__main__':unittest.main(verbosity=2)