import os
import unittest

from rcl_rlm.modelrelay import (
    DEFAULT_MODEL,
    ModelResolutionPolicy,
    resolve_root_model,
    resolve_subcall_model,
)


class EnvGuard:
    def __init__(self, updates: dict[str, str | None]) -> None:
        self._updates = updates
        self._original: dict[str, str | None] = {}

    def __enter__(self) -> "EnvGuard":
        for key, value in self._updates.items():
            self._original[key] = os.environ.get(key)
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        for key, value in self._original.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value


class TestModelResolution(unittest.TestCase):
    def test_root_model_explicit_over_env(self) -> None:
        policy = ModelResolutionPolicy(allow_env=True)
        with EnvGuard(
            {
                "RECALL_RLM_ROOT_MODEL": "env-root",
                "RECALL_RLM_MODEL": "env-model",
            }
        ):
            self.assertEqual(resolve_root_model("explicit-root", policy=policy), "explicit-root")

    def test_root_model_env_precedence(self) -> None:
        policy = ModelResolutionPolicy(allow_env=True)
        with EnvGuard(
            {
                "RECALL_RLM_ROOT_MODEL": "env-root",
                "RECALL_RLM_MODEL": "env-model",
                "MODELRELAY_MODEL": "env-alt",
            }
        ):
            self.assertEqual(resolve_root_model(None, policy=policy), "env-root")
        with EnvGuard({"RECALL_RLM_MODEL": "env-model"}):
            self.assertEqual(resolve_root_model(None, policy=policy), "env-model")
        with EnvGuard({"MODELRELAY_MODEL": "env-alt"}):
            self.assertEqual(resolve_root_model(None, policy=policy), "env-alt")

    def test_env_disabled_ignores_env(self) -> None:
        policy = ModelResolutionPolicy(allow_env=False)
        with EnvGuard(
            {
                "RECALL_RLM_ROOT_MODEL": "env-root",
                "RECALL_RLM_MODEL": "env-model",
                "MODELRELAY_MODEL": "env-alt",
            }
        ):
            self.assertEqual(resolve_root_model(None, policy=policy), DEFAULT_MODEL)

    def test_subcall_explicit_over_env(self) -> None:
        policy = ModelResolutionPolicy(allow_env=True)
        with EnvGuard(
            {
                "RECALL_RLM_SUBCALL_MODEL": "env-subcall",
                "RECALL_RLM_ROOT_MODEL": "env-root",
            }
        ):
            self.assertEqual(
                resolve_subcall_model("explicit-subcall", root_fallback="explicit-root", policy=policy),
                "explicit-subcall",
            )

    def test_subcall_env_over_root_fallback(self) -> None:
        policy = ModelResolutionPolicy(allow_env=True)
        with EnvGuard(
            {
                "RECALL_RLM_SUBCALL_MODEL": "env-subcall",
                "RECALL_RLM_ROOT_MODEL": "env-root",
            }
        ):
            self.assertEqual(
                resolve_subcall_model(None, root_fallback="explicit-root", policy=policy),
                "env-subcall",
            )

    def test_subcall_root_fallback_over_env_root(self) -> None:
        policy = ModelResolutionPolicy(allow_env=True)
        with EnvGuard({"RECALL_RLM_ROOT_MODEL": "env-root"}):
            self.assertEqual(
                resolve_subcall_model(None, root_fallback="explicit-root", policy=policy),
                "explicit-root",
            )

    def test_subcall_falls_back_to_root_env(self) -> None:
        policy = ModelResolutionPolicy(allow_env=True)
        with EnvGuard({"RECALL_RLM_ROOT_MODEL": "env-root"}):
            self.assertEqual(resolve_subcall_model(None, policy=policy), "env-root")


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