from __future__ import annotations

import asyncio
import os
import unittest
from unittest.mock import patch

from app.main import _app_lifespan, app
from app.simulation.warmup import (
    SimulationWarmupReport,
    _reset_simulation_warmup_for_tests,
    warm_up_simulation_runtime,
)


class SimulationWarmupTests(unittest.TestCase):
    def setUp(self) -> None:
        _reset_simulation_warmup_for_tests()
        self.engine_environment = patch.dict(os.environ, {"SIMULATION_NUMERIC_ENGINE": "native"})
        self.engine_environment.start()
        self.addCleanup(self.engine_environment.stop)

    def tearDown(self) -> None:
        _reset_simulation_warmup_for_tests()

    def test_warmup_runs_only_once_per_process(self) -> None:
        with (
            patch.dict(os.environ, {"SIMULATIONAPP_WARMUP": "on"}),
            patch("app.simulation.warmup._run_native_warmup") as run_warmup,
        ):
            first = warm_up_simulation_runtime()
            second = warm_up_simulation_runtime()

        self.assertIs(second, first)
        self.assertEqual(first.status, "completed")
        run_warmup.assert_called_once_with()

    def test_disabled_warmup_does_not_touch_numerical_runtime(self) -> None:
        with (
            patch.dict(os.environ, {"SIMULATIONAPP_WARMUP": "off"}),
            patch("app.simulation.warmup._run_native_warmup") as run_warmup,
        ):
            report = warm_up_simulation_runtime()

        self.assertEqual(report.status, "disabled")
        run_warmup.assert_not_called()

    def test_regular_failure_is_reported_without_blocking_startup(self) -> None:
        with (
            patch.dict(os.environ, {"SIMULATIONAPP_WARMUP": "on"}),
            patch(
                "app.simulation.warmup._run_native_warmup",
                side_effect=RuntimeError("broken warmup"),
            ),
            self.assertLogs("app.simulation.warmup", level="ERROR"),
        ):
            report = warm_up_simulation_runtime()

        self.assertEqual(report.status, "failed")
        self.assertIn("broken warmup", report.error or "")

    def test_memory_error_remains_fatal(self) -> None:
        with (
            patch.dict(os.environ, {"SIMULATIONAPP_WARMUP": "on"}),
            patch(
                "app.simulation.warmup._run_native_warmup",
                side_effect=MemoryError("out of memory"),
            ),
            self.assertRaises(MemoryError),
        ):
            warm_up_simulation_runtime()

    def test_lifespan_stores_warmup_report_before_serving(self) -> None:
        report = SimulationWarmupReport(status="completed", duration_ms=12.5)

        async def enter_lifespan() -> None:
            with patch(
                "app.simulation.warmup.warm_up_simulation_runtime",
                return_value=report,
            ) as warmup:
                async with _app_lifespan(app):
                    self.assertEqual(app.state.simulation_numeric_engine, "native")
                    self.assertEqual(
                        app.state.simulation_warmup,
                        report.as_dict(),
                    )
                warmup.assert_called_once_with()

        asyncio.run(enter_lifespan())

    def test_default_startup_checks_native_without_python_warmup(self) -> None:
        with (
            patch.dict(os.environ, {"SIMULATIONAPP_WARMUP": "on"}),
            patch("app.simulation.warmup._run_native_warmup") as native,
        ):
            os.environ.pop("SIMULATION_NUMERIC_ENGINE", None)

            async def enter_lifespan() -> None:
                async with _app_lifespan(app):
                    self.assertEqual(app.state.simulation_numeric_engine, "native")
                    self.assertEqual(app.state.simulation_warmup["status"], "completed")

            asyncio.run(enter_lifespan())
        native.assert_called_once_with()

    def test_native_failure_does_not_fall_back_to_python(self) -> None:
        with (
            patch.dict(os.environ, {"SIMULATION_NUMERIC_ENGINE": "native", "SIMULATIONAPP_WARMUP": "on"}),
            patch("app.simulation.warmup._run_native_warmup", side_effect=RuntimeError("C compiler unavailable")),
            self.assertLogs("app.simulation.warmup", level="ERROR"),
        ):
            report = warm_up_simulation_runtime()
        self.assertEqual(report.status, "failed")
        self.assertIn("C compiler unavailable", report.error or "")

    def test_real_numerical_warmup_completes(self) -> None:
        with patch.dict(os.environ, {"SIMULATIONAPP_WARMUP": "on"}):
            report = warm_up_simulation_runtime()

        self.assertEqual(report.status, "completed", report.error)
        self.assertGreater(report.duration_ms, 0.0)



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