from __future__ import annotations

import asyncio
from contextlib import contextmanager
from functools import lru_cache
import importlib
import os
import unittest
from unittest.mock import patch

import app.simulation.performance as performance


@contextmanager
def profiling_mode(mode: str):
    previous = os.environ.get("SIMULATIONAPP_PROFILE")
    os.environ["SIMULATIONAPP_PROFILE"] = mode
    module = importlib.reload(performance)
    try:
        yield module
    finally:
        if previous is None:
            os.environ.pop("SIMULATIONAPP_PROFILE", None)
        else:
            os.environ["SIMULATIONAPP_PROFILE"] = previous
        importlib.reload(performance)


class SimulationPerformanceTests(unittest.TestCase):

    def test_standard_mode_records_nested_inclusive_and_self_time(self) -> None:
        with profiling_mode("standard") as module:
            @module.profile_phase("inner")
            def inner() -> None:
                return None

            @module.profile_phase("outer")
            def outer() -> None:
                inner()

            clock = iter((0, 10, 20, 30, 50, 80))
            with patch.object(module, "perf_counter_ns", side_effect=clock):
                with module.profile_run() as trace:
                    outer()

            snapshot = trace.snapshot()
            self.assertEqual(
                snapshot["phases"]["inner"],
                {
                    "calls": 1,
                    "inclusiveNs": 10,
                    "selfNs": 10,
                    "maxNs": 10,
                    "errors": 0,
                },
            )
            self.assertEqual(snapshot["phases"]["outer"]["inclusiveNs"], 40)
            self.assertEqual(snapshot["phases"]["outer"]["selfNs"], 30)
            self.assertEqual(
                snapshot["phases"]["simulation.total"]["inclusiveNs"],
                80,
            )
            self.assertEqual(
                snapshot["phases"]["simulation.total"]["selfNs"],
                40,
            )

    def test_errors_are_recorded_and_propagated(self) -> None:
        with profiling_mode("standard") as module:
            @module.profile_phase("explode")
            def explode() -> None:
                raise RuntimeError("expected")

            with self.assertRaisesRegex(RuntimeError, "expected"):
                with module.profile_run() as trace:
                    explode()

            snapshot = trace.snapshot()
            self.assertEqual(snapshot["phases"]["explode"]["errors"], 1)
            self.assertEqual(
                snapshot["phases"]["simulation.total"]["errors"],
                1,
            )

    def test_manual_performance_span_uses_the_current_trace(self) -> None:
        with profiling_mode("standard") as module:
            clock = iter((0, 10, 20, 30))
            with patch.object(module, "perf_counter_ns", side_effect=clock):
                with module.profile_run() as trace:
                    with module.performance_span("manual"):
                        pass

            snapshot = trace.snapshot()
            self.assertEqual(snapshot["phases"]["manual"]["inclusiveNs"], 10)
            self.assertEqual(snapshot["phases"]["simulation.total"]["selfNs"], 20)




    def test_async_profile_runs_are_context_isolated(self) -> None:
        with profiling_mode("standard") as module:
            @module.profile_phase("work")
            async def work() -> None:
                await asyncio.sleep(0)

            async def one_run() -> dict[str, object]:
                with module.profile_run() as trace:
                    await work()
                return trace.snapshot()

            async def exercise() -> list[dict[str, object]]:
                return await asyncio.gather(one_run(), one_run())

            snapshots = asyncio.run(exercise())
            self.assertEqual(
                [snapshot["phases"]["work"]["calls"] for snapshot in snapshots],
                [1, 1],
            )


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

class NativeOrchestrationPerformanceTests(unittest.TestCase):
    def test_off_mode_keeps_the_original_phase_callable(self):
        with profiling_mode('off') as module:
            def work(): return 1
            self.assertIs(module.profile_phase('work')(work), work)
            with module.profile_run() as trace:
                work()
            self.assertEqual(trace.snapshot(), {'mode': 'off', 'phases': {}})