import unittest
import os
import signal
import subprocess
import tempfile
import time
from pathlib import Path
import numpy as np
from types import SimpleNamespace
from app.byte_tracking import _prepare_observations


class PreparationTest(unittest.TestCase):
    def test_geometry_score_and_candidate_order_match_previous(self):
        observations = [SimpleNamespace(class_id=i%3,score=.05+i*.02,bbox_xyxy=(.123,.234,.456,.789),
            frame=SimpleNamespace(frame_width=1920+i,frame_height=1080+i)) for i in range(20)]
        groups,boxes,scores=_prepare_observations(observations)
        for cls,indices in groups.items():
            expected=[np.asarray(observations[i].bbox_xyxy)*np.array([
                observations[i].frame.frame_width,observations[i].frame.frame_height,
                observations[i].frame.frame_width,observations[i].frame.frame_height]) for i in indices]
            np.testing.assert_array_equal(boxes[indices],np.asarray(expected,dtype=np.float32))
            np.testing.assert_array_equal(scores[indices],np.asarray([observations[i].score for i in indices],dtype=np.float32))
            self.assertEqual(indices,[i for i,o in enumerate(observations) if o.class_id==cls])

    def test_empty(self):
        groups,boxes,scores=_prepare_observations(())
        self.assertEqual(groups,{})
        self.assertEqual(boxes.shape,(0,4));self.assertEqual(scores.shape,(0,))

    def test_entrypoint_scopes_blas_to_application(self):
        script=Path(__file__).resolve().parents[1]/'docker/entrypoint.sh'
        with tempfile.TemporaryDirectory() as directory:
            root=Path(directory)
            for name,body in {
                'tritonserver': 'echo "triton:${OPENBLAS_NUM_THREADS-unset}" >> "$RESULT"; exec sleep 30',
                'python3': 'if [ "$1 $2" = "-m app.main" ]; then echo "app:${OPENBLAS_NUM_THREADS-unset}" >> "$RESULT"; exec sleep 30; fi; exit 0',
            }.items():
                p=root/name;p.write_text('#!/bin/bash\n'+body+'\n');p.chmod(0o700)
            env=dict(os.environ,PATH=str(root)+':/usr/bin:/bin',PYTHONPATH='',GPU_ID='0',RESULT=str(root/'result'),WATCHDOG_INTERVAL='1')
            env.pop('OPENBLAS_NUM_THREADS',None);env.pop('CONTEXTRA_APP_OPENBLAS_THREADS',None)
            process=subprocess.Popen(['bash',str(script)],env=env,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=True)
            try:
                deadline=time.monotonic()+3
                while time.monotonic()<deadline:
                    if (root/'result').exists() and 'app:1' in (root/'result').read_text():break
                    time.sleep(.02)
                self.assertEqual((root/'result').read_text().splitlines(),['triton:unset','app:1'])
            finally:
                os.killpg(process.pid,signal.SIGTERM)
                try:process.wait(3)
                except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait(3)
            invalid=dict(env,CONTEXTRA_APP_OPENBLAS_THREADS='0')
            result=subprocess.run(['bash',str(script)],env=invalid,capture_output=True,timeout=3)
            self.assertNotEqual(result.returncode,0)
            self.assertEqual((root/'result').read_text().splitlines(),['triton:unset','app:1'])
