"""Guards for the memory layout described in docs/memory.md.

* The server process must not import the math libraries (they live in the workers).
* Workers must not import the MCP/HTTP stack (spawn re-imports the main module).
"""

from __future__ import annotations

import builtins
import io
import json
import os
import subprocess
import sys
import tempfile
import textwrap
import typing
import unittest
from pathlib import Path
from unittest import mock

from mcp_calculator import executor as executor_module
from mcp_calculator.config import load_settings
from mcp_calculator.executor import ComputationError, Executor
from mcp_calculator.ops import spec, stats
from tests.support import slow

MATH_LIBS = ("sympy", "numpy", "scipy", "pint", "mpmath")
SERVER_LIBS = ("mcp", "uvicorn", "starlette", "pydantic", "anyio")


def _run_python(code: str, modules: tuple[str, ...]) -> dict:
    """Run ``code`` as a script (MODULES is replaced by the module names); return its last JSON line."""
    with tempfile.TemporaryDirectory() as tmp:
        script = Path(tmp) / "probe.py"
        script.write_text(textwrap.dedent(code).replace("MODULES", repr(modules)))
        out = subprocess.run([sys.executable, str(script)], capture_output=True, text=True, timeout=120, check=True)
    return json.loads(out.stdout.strip().splitlines()[-1])


@slow
class ImportFootprintTest(unittest.TestCase):
    def test_server_process_does_not_import_math_libraries(self):
        loaded = _run_python(
            """
            import json, sys
            import mcp_calculator.__main__, mcp_calculator.server, mcp_calculator.http_app, mcp_calculator.executor
            print(json.dumps({m: m in sys.modules for m in MODULES}))
            """,
            MATH_LIBS,
        )
        self.assertFalse(any(loaded.values()), loaded)

    def test_workers_do_not_import_server_stack(self):
        # Mirrors the console script: the main module imports mcp_calculator.__main__.
        loaded = _run_python(
            """
            import json, sys
            from mcp_calculator.__main__ import main  # noqa: F401

            def probe():
                import os
                loaded = {m: m in sys.modules for m in MODULES}
                loaded["single_threaded_blas"] = os.environ.get("OPENBLAS_NUM_THREADS") == "1"
                return loaded

            if __name__ == "__main__":
                from mcp_calculator.config import load_settings
                from mcp_calculator.executor import Executor

                executor = Executor(load_settings([], env={"MCP_CALC_WORKERS": "1"}))
                executor.start()
                result = executor._pool.schedule(probe, timeout=60).result()
                executor.stop()
                print(json.dumps(result))
            """,
            MATH_LIBS + SERVER_LIBS,
        )
        self.assertTrue(loaded["sympy"] and loaded["scipy"], "workers should preload the math libraries")
        self.assertFalse(any(loaded[m] for m in ("mcp", "uvicorn", "starlette")), loaded)
        self.assertTrue(loaded["single_threaded_blas"])


class MemoryLimitTest(unittest.TestCase):
    def test_limit_is_relative_to_loaded_size(self):
        """RLIMIT_AS = current address space + configured extra (Linux path, simulated)."""
        with (
            mock.patch.object(builtins, "open", lambda *a, **k: io.StringIO("100000 5000 0 0 0 0 0")),
            mock.patch.object(executor_module.resource, "setrlimit") as setrlimit,
        ):
            executor_module._limit_address_space(512)
        expected = 100000 * os.sysconf("SC_PAGE_SIZE") + 512 * 1024 * 1024
        setrlimit.assert_called_once_with(executor_module.resource.RLIMIT_AS, (expected, expected))

    def test_failure_is_not_fatal(self):
        with mock.patch.object(builtins, "open", side_effect=OSError("no /proc")), self.assertLogs("mcp_calculator.executor", "WARNING") as logs:
            executor_module._limit_address_space(512)
        self.assertIn("could not set memory limit", logs.output[0])


class TaskDispatchTest(unittest.IsolatedAsyncioTestCase):
    async def test_only_public_ops_functions_can_be_called(self):
        executor = Executor(load_settings([], env={"MCP_CALC_EXECUTOR": "inline"}))
        for task in ("os.system", "algebra._equivalent", "algebra", "parsing.parse", "algebra.__class__"):
            with self.subTest(task=task), self.assertRaisesRegex(ComputationError, "unknown task"):
                await executor.run(task)


class SpecConsistencyTest(unittest.TestCase):
    def test_distribution_tables_are_consistent(self):
        names = set(typing.get_args(spec.Distribution))
        self.assertEqual(names, set(spec.DISTRIBUTION_PARAMETERS))
        self.assertEqual(names, set(stats._DISTRIBUTIONS))