import unittest
from comfy_progress import _hp_event, _hp_snapshot, _hp_history_estimate

class ProgressTests(unittest.TestCase):
    def state(self):
        return {'prompt':'one','classes':{'12':'SamplerCustomAdvanced','13':'VAEDecode'},'samples':[], 'connected':True}
    def event(self,state,value,t,node='12',prompt='one',maximum=20):
        _hp_event(state,{'type':'progress','data':{'prompt_id':prompt,'node':node,'value':value,'max':maximum}},t)
    def test_job_isolation_and_node_reset(self):
        s=self.state();self.event(s,1,100);self.event(s,2,110,prompt='other')
        self.assertEqual(s['value'],1)
        self.event(s,2,110);self.event(s,3,120)
        self.assertEqual(_hp_snapshot(s,120)['samplingMinSeconds'],136)
        _hp_event(s,{'type':'executing','data':{'prompt_id':'one','node':'13'}},125)
        r=_hp_snapshot(s,125);self.assertEqual(r['phase'],'decoding');self.assertNotIn('step',r)
    def test_stale_disconnect_reset_no_false_countdown(self):
        s=self.state()
        for i in range(1,4):self.event(s,i,100+i*10)
        self.assertNotIn('samplingMinSeconds',_hp_snapshot(s,200))
        self.event(s,1,201)
        self.assertNotIn('samplingMinSeconds',_hp_snapshot(s,201))
        s['connected']=False
        self.assertEqual(_hp_snapshot(s,202),{'phase':'running'})
    def test_sampling_completion_is_not_video_completion(self):
        s=self.state()
        for i in [18,19,20]:self.event(s,i,100+i*10)
        r=_hp_snapshot(s,300)
        self.assertEqual(r['phase'],'sampling');self.assertEqual(r['step'],20)
        self.assertNotIn('samplingMinSeconds',r)
    def test_history_exact_settings_and_cache_exclusion(self):
        graph={'1':{'class_type':'SamplerCustomAdvanced','inputs':{'steps':20,'noise_seed':123}},
               '2':{'class_type':'MiniMaxH3ReferenceToVideo','inputs':{'width':1344,'height':768,'text':'private'}}}
        def history(duration=1000,change=None,cached=False,failed=False):
            import copy
            other=copy.deepcopy(graph);other['1']['inputs']['noise_seed']=456;other['2']['inputs']['text']='another'
            if change:other['2']['inputs'].update(change)
            return {'prompt':[1,'id',other], 'status':{'status_str':'error' if failed else 'success','completed':True,
                'messages':[['execution_start',{'timestamp':1000}],['execution_cached',{'nodes':['1'] if cached else []}],
                            ['execution_success',{'timestamp':1000+duration*1000}]]}}
        r=_hp_history_estimate(graph,{'a':history()})
        self.assertEqual(r,{'historySamples':1})
        r=_hp_history_estimate(graph,{'a':history(),'b':history(1200),'c':history(5,cached=True),
                                    'd':history(500,change={'width':2560}),'e':history(2,failed=True)})
        self.assertEqual(r,{'historySamples':2,'historyMinSeconds':800,'historyMaxSeconds':1440})

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