import asyncio
import time
import unittest
from unittest.mock import AsyncMock, MagicMock, patch

from app.ml.runtime import EmailRiskModel, predict_async


def _slow_pipeline(texts, batch_size=None):
    time.sleep(0.4)
    return [[{"label": "NORMAL", "score": 0.9}] for _ in texts]


class NonBlockingInferenceTests(unittest.IsolatedAsyncioTestCase):
    """CPU inference must run OFF the event loop.

    Measured in production: the first request after a cold start took 42.6s and
    warm requests ~2s because torch inference ran synchronously inside `async
    def` handlers, stalling every concurrent request (including /health).
    """

    async def test_predict_async_does_not_block_the_event_loop(self):
        m = EmailRiskModel(repo="user/model")
        m._pipeline = _slow_pipeline
        m._loaded = True

        ticks = 0

        async def ticker():
            nonlocal ticks
            while True:
                await asyncio.sleep(0.02)
                ticks += 1

        t = asyncio.create_task(ticker())
        out = await predict_async(m, ["a [SEP] b.test"])
        t.cancel()

        self.assertEqual(out, [("NORMAL", 0.9)])
        # A blocking implementation would leave the loop frozen for ~0.4s
        # (ticks == 0); off-loop it keeps ticking throughout.
        self.assertGreaterEqual(
            ticks, 5, f"event loop was blocked during inference (ticks={ticks})"
        )

    async def test_load_async_runs_the_factory_off_the_event_loop(self):
        def slow_factory(repo, batch_size):
            time.sleep(0.3)
            return _slow_pipeline

        m = EmailRiskModel(repo="user/model", pipeline_factory=slow_factory)
        ticks = 0

        async def ticker():
            nonlocal ticks
            while True:
                await asyncio.sleep(0.02)
                ticks += 1

        t = asyncio.create_task(ticker())
        self.assertTrue(await m.load_async())
        t.cancel()
        self.assertGreaterEqual(
            ticks, 4, f"model load blocked the event loop (ticks={ticks})"
        )

    async def test_predict_async_is_fail_open(self):
        m = EmailRiskModel(repo="user/model")
        m._pipeline = MagicMock(side_effect=RuntimeError("boom"))
        m._loaded = True
        self.assertIsNone(await predict_async(m, ["a [SEP] b.test"]))

    async def test_predict_async_handles_unloaded_model(self):
        m = EmailRiskModel(repo="user/model")
        self.assertIsNone(await predict_async(m, ["a [SEP] b.test"]))

    async def test_predict_async_matches_sync_results_for_fakes(self):
        # Engine tests inject MagicMock models with a sync predict; the async
        # helper must work with those unchanged.
        fake = MagicMock()
        fake.predict = MagicMock(return_value=[("DISPOSABLE", 0.97)])
        out = await predict_async(fake, ["x [SEP] y.test"])
        self.assertEqual(out, [("DISPOSABLE", 0.97)])
        fake.predict.assert_called_once()

    async def test_concurrent_predictions_are_parallelised(self):
        """Two 0.3s inferences should overlap, not serialise (thread pool)."""
        m = EmailRiskModel(repo="user/model")
        m._pipeline = _slow_pipeline
        m._loaded = True
        t0 = time.perf_counter()
        await asyncio.gather(
            predict_async(m, ["a [SEP] b.test"]),
            predict_async(m, ["c [SEP] d.test"]),
        )
        elapsed = time.perf_counter() - t0
        # Serial execution would take ~0.8 s (two full sleeps); anything
        # comfortably below that proves overlap. 0.7 leaves scheduling slack
        # on a loaded machine while still failing on serialisation.
        self.assertLess(elapsed, 0.7, f"inferences serialised ({elapsed:.2f}s)")


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