import time
import unittest
from unittest.mock import patch

import numpy as np
import torch

from helpers import filter_yolo_boxes_with_scores
from yolo_worker import INFERENCE_DEVICE, INFERENCE_HALF, INFERENCE_SIZE_FULL, YOLOWorker, fixed_letterbox, raw_yolo_boxes


class FailingModel:
    def __call__(self, *args, **kwargs):
        raise RuntimeError("test failure")


class Boxes:
    xyxy = torch.tensor([[0.0, 160.0, 640.0, 480.0]])
    conf = torch.tensor([0.9])
    cls = torch.tensor([0.0])

    def __len__(self):
        return 1


class Result:
    boxes = Boxes()


class YOLOWorkerLifecycleTests(unittest.TestCase):
    def test_cpu_inference_uses_cpu_precision_and_workable_size(self):
        if torch.cuda.is_available():
            self.skipTest("CPU fallback is inactive on CUDA hosts")
        self.assertEqual(INFERENCE_DEVICE, "cpu")
        self.assertFalse(INFERENCE_HALF)
        self.assertLessEqual(INFERENCE_SIZE_FULL, 640)

    def test_fixed_letterbox_maps_boxes_back_without_distortion(self):
        image = np.zeros((100, 200, 3), dtype=np.uint8)
        boxed, scale, pad_x, pad_y = fixed_letterbox(image, 640)
        self.assertEqual(boxed.shape, (640, 640, 3))
        raw = raw_yolo_boxes(Result(), scale=1.0 / scale, pad_x=pad_x, pad_y=pad_y)
        np.testing.assert_allclose(raw[0][:4], [0.0, 0.0, 200.0, 100.0])
        filtered = filter_yolo_boxes_with_scores(
            Result(),
            frame_w=200,
            frame_h=100,
            min_conf=0.1,
            input_scale=scale,
            pad_x=pad_x,
            pad_y=pad_y,
            content_w=200,
            content_h=100,
        )
        np.testing.assert_allclose(filtered[0][:4], [0.0, 0.0, 200.0, 100.0])

    def test_worker_reports_failed_inference_without_hanging(self):
        worker = YOLOWorker(FailingModel())
        with patch("yolo_worker.torch.cuda.is_available", return_value=False):
            worker.start()
            try:
                worker.submit(np.zeros((16, 16, 3), dtype=np.uint8), None, "FULL", 1.0)
                deadline = time.monotonic() + 2.0
                result = None
                while result is None and time.monotonic() < deadline:
                    result = worker.try_get()
                    time.sleep(0.01)
                self.assertIsNotNone(result)
                self.assertEqual(result[0], [])
                self.assertEqual(result[1], 1.0)
            finally:
                worker.stop()
        self.assertFalse(worker.thread.is_alive())


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