import copy
import json
from pathlib import Path
import tempfile
import unittest

from tools.stream_benchmark import Benchmark, load_config
from tools.stream_benchmark_metrics import summarize
from tools.stream_benchmark_runtime import Runtime, subset


def observation(at, count, batch=1):
    stat = {k:dict(count=count,ns=count*1000000) for k in ('success','queue','compute_input','compute_infer','compute_output')}
    return dict(at=at,channels={'a':dict(run_id='run',generation=1,pid=1,state='running',error=None,
        processed=count,queue=0,stale=0,drop=0,counts=dict(infer=count),publisher=dict(replaced=0,send_error=0),
        redis=dict(at=at,entries=count,length=100,pttl=9000))},
        triton=dict(at=at,pending=0,stats=dict(inference_count=count,execution_count=count//batch,
            inference_stats=stat,batch_stats=[dict(batch_size=batch,compute_infer=dict(count=count//batch,ns=count*1000000))])),
        resource=dict(at=at,memory=1024,cpu_us=count*100,gpu_mib=100,gpu_util=10))


class MetricsTest(unittest.TestCase):
    def test_batch_counts_and_request_denominators(self):
        result=summarize([observation(1,0,4),observation(11,100,4)],True)
        self.assertEqual(10,result['rows'][0]['result_fps'])
        self.assertEqual(4,result['triton']['average_batch'])
        self.assertEqual({'4':25},result['triton']['batch_executions'])
        self.assertEqual(1,result['triton']['durations']['compute_infer']['mean_ms'])
        self.assertTrue(result['accepted'])

    def test_missing_redis_and_stage_are_not_zero(self):
        a,b=observation(1,0),observation(11,100)
        a['channels']['a']['redis'].pop('entries')
        result=summarize([a,b],True)
        self.assertIsNone(result['rows'][0]['redis_fps'])
        self.assertIsNone(result['rows'][0]['selected_fps'])
        self.assertFalse(result['accepted'])

    def test_reset_run_and_partial(self):
        a,b=observation(1,100),observation(11,0)
        with self.assertRaises(ValueError):summarize([a,b],True)
        a,b=observation(1,0),observation(11,100)
        b['channels']['a']['run_id']='other'
        with self.assertRaises(ValueError):summarize([a,b],True)
        b['channels']['a']['run_id']='run'
        self.assertFalse(summarize([a,b],False)['accepted'])

    def test_settings_subset(self):
        self.assertTrue(subset({'pose':{'enabled':False}},{'pose':{'enabled':False,'model':'x'}}))
        self.assertFalse(subset({'pose':{'enabled':False}},{'pose':{'enabled':True}}))


class ControlTest(unittest.TestCase):
    def check_failure_cleanup(self, stop_all):
        with tempfile.TemporaryDirectory() as tmp:
            class Fake:
                config={'stages':[dict(channels=3,seconds=20)],'failure_stop_all':stop_all}
                model={'name':'fixture'}
                evidence=Path(tmp)
                def __init__(self):
                    self.live={'original':True,'one':False,'two':False,'foreign':True};self.started=[];self.stopped=[];self.saved={}
                def resources(self):return {}
                def start(self,sid):
                    if sid=='two':raise RuntimeError('First Frame failed')
                    if not self.live[sid]:self.started.append(sid);self.live[sid]=True
                    return {'child_pid':1}
                def stop(self,sid):
                    if self.live[sid]:self.stopped.append(sid);self.live[sid]=False
                    return {'state':'stopped','child_pid':None}
                def states(self):return {s:dict(state='running' if v else 'stopped') for s,v in self.live.items()}
                def settings(self,sid):return {'tracking':False}
                def save(self,name,value):self.saved[name]=value
            fake=Fake();rows=[dict(stream_id=s,original=s=='original') for s in ('original','one','two')]
            result=Benchmark(fake).run(dict(rows=rows,settings={r['stream_id']:{'tracking':False} for r in rows}))
            self.assertEqual(2,result)
            self.assertEqual(['one'],fake.started);self.assertEqual(['one','original'] if stop_all else ['one'],fake.stopped)
            self.assertEqual(not stop_all,fake.live['original']);self.assertTrue(fake.live['foreign'])

    def test_only_target_additions_are_stopped_on_failure(self):
        self.check_failure_cleanup(False)

    def test_explicit_failure_policy_stops_original_targets(self):
        self.check_failure_cleanup(True)

    def test_mutation_marker_prevents_duplicate_request(self):
        with tempfile.TemporaryDirectory() as tmp:
            runtime=Runtime({'api_base':'http://127.0.0.1:1','container':'fixture','model':{}},tmp)
            class API:
                calls=0
                def call(self,*args):self.calls+=1;return {}
            runtime.api=API()
            runtime.mutate('a','start')
            with self.assertRaises(FileExistsError):runtime.mutate('a','start')
            self.assertEqual(1,runtime.api.calls)

    def test_config_rejects_unknown_or_nonlocal(self):
        path=Path(__file__).resolve().parents[1]/'tools/examples/stream_benchmark.example.yaml'
        config=load_config(path)
        self.assertEqual(6,config['model']['instances'])
        with tempfile.TemporaryDirectory() as tmp:
            target=Path(tmp)/'bad.json'
            for patch in ({'unknown':1},{'api_base':'http://example.com'}):
                value=copy.deepcopy(config);value.update(patch);target.write_text(json.dumps(value))
                with self.assertRaises(ValueError):load_config(target)


if __name__=='__main__':unittest.main()