import threading
import time
from unittest import TestCase
from station.tasks import TaskManager


class TasksTest(TestCase):
    def setUp(self): self.manager=TaskManager()
    def tearDown(self): self.manager.close()
    def wait(self,token):
        until=time.monotonic()+3
        while time.monotonic()<until:
            row=self.manager.snapshot(token)
            if row['state'] not in ('running','queued'): return row
            time.sleep(.01)
        self.fail('task did not complete')

    def test_resources_held_until_completion(self):
        gate=threading.Event()
        one=self.manager.submit('move',{'arm:89'},lambda cancel:(gate.wait(2),{'ok':True})[1],context={'arm':'89'})
        self.assertEqual(self.manager.snapshot(one)['resources'],['arm:89'])
        self.assertEqual(self.manager.snapshot(one)['context']['arm'],'89')
        with self.assertRaises(RuntimeError): self.manager.submit('move-again',{'arm:89'},lambda e:{'ok':True})
        other=self.manager.submit('camera',{'camera'},lambda e:{'ok':True})
        self.assertEqual(self.wait(other)['state'],'succeeded')
        gate.set();self.assertEqual(self.wait(one)['state'],'succeeded')
        next_task=self.manager.submit('next',{'arm:89'},lambda e:{'ok':True})
        self.assertEqual(self.wait(next_task)['state'],'succeeded')

    def test_cancel_and_exception_release_ownership(self):
        token=self.manager.submit('long',{'arm'},lambda event:(event.wait(2),{'ok':False})[1])
        self.manager.cancel(token)
        self.assertEqual(self.wait(token)['state'],'cancelled')
        def fail(event): raise ValueError('driver rejected')
        token=self.manager.submit('reject',{'arm'},fail)
        result=self.wait(token)
        self.assertEqual(result['state'],'failed');self.assertIn('driver rejected',result['result']['message'])

    def test_busy_message_identifies_cleanup_and_stop_cancels_only_motion(self):
        gate=threading.Event()
        motion=self.manager.submit('stage_locate',{'robot-control','arm:89'},lambda e:(gate.wait(2),{'ok':True})[1],context={'arm':'89'})
        camera=self.manager.submit('camera_refresh',{'camera-setup:89'},lambda e:(gate.wait(2),{'ok':True})[1])
        try:
            deadline=time.monotonic()+1
            while self.manager.snapshot(motion)['state']=='queued' and time.monotonic()<deadline: time.sleep(.01)
            self.assertEqual(self.manager.snapshot(motion)['state'],'running')
            self.manager.set_phase(self.manager._events[motion],'收尾')
            with self.assertRaisesRegex(RuntimeError,'机械臂 89.*对孔.*收尾'):
                self.manager.submit('reset_home',{'robot-control'},lambda e:{'ok':True})
            self.manager.cancel_resources({'robot-control'})
            self.assertTrue(self.manager._events[motion].is_set())
            self.assertFalse(self.manager._events[camera].is_set())
        finally: gate.set()

    def test_cancelled_task_cannot_report_late_success_or_resume(self):
        started=threading.Event(); release=threading.Event(); calls=[]
        def old_task(event):
            calls.append('old'); started.set(); release.wait(2)
            return {'ok':True}
        token=self.manager.submit('stage_track',{'robot-control'},old_task)
        self.assertTrue(started.wait(1))
        self.manager.cancel_resources({'robot-control'}); release.set()
        result=self.wait(token)
        self.assertEqual(result['state'],'cancelled'); self.assertFalse(result['result']['ok'])
        fresh=self.manager.submit('real_up',{'robot-control'},lambda e:{'ok':True})
        self.assertEqual(self.wait(fresh)['state'],'succeeded')
        self.assertEqual(calls,['old'])