"""Environment variable access guard for replay."""

from __future__ import annotations

import os
from typing import Dict

from autodebug_replay.transcripts import ReplayDivergence


class EnvGuard:
    def __init__(self, allowed: Dict[str, str]) -> None:
        self.allowed = allowed
        self._orig_getenv = os.getenv
        self._orig_getitem = os.environ.__getitem__
        self._orig_environ_get = os.environ.get
        # Env vars that are safe to access without capture (infrastructure/config)
        self._safe_keys = {
            # Temp directories
            "TEMP", "TMP", "TMPDIR",
            # System config
            "HOME", "USER", "LANG", "LC_ALL", "LC_CTYPE", "SHELL", "TERM",
            # Python config
            "PATH", "PYTHONPATH", "PYTHONHASHSEED", "PYTHONDONTWRITEBYTECODE",
            # Colors
            "NO_COLOR", "FORCE_COLOR", "CLICOLOR",
            # Protobuf
            "TEMPORARILY_DISABLE_PROTOBUF_VERSION_CHECK",
            # FaaS detection vars (pymongo, boto3, etc check these)
            "FUNCTION_NAME", "K_SERVICE", "K_REVISION",
        }
        # Prefixes that are safe (library configuration read at import time)
        self._safe_prefixes = (
            "GRPC_", "PROTOCOL_", "AWS_", "BOTO_", "MONGO_", "REDIS_",
            "LC_", "LANG_", "PYTHON",
            # FaaS detection (Lambda, Azure Functions, GCP, Vercel)
            "FUNCTIONS_", "AWS_LAMBDA_", "K_", "VERCEL", "GCP_",
            # Container/Kubernetes detection
            "KUBERNETES_", "ECS_", "DOCKER_",
        )

    def _fail(self, key: str) -> None:
        raise ReplayDivergence("UNCAPTURED_ENV_ACCESS", {"var": key})

    def _is_safe(self, key: str) -> bool:
        """Check if an env var is safe to access without capture."""
        if key in self._safe_keys:
            return True
        if key.startswith(self._safe_prefixes):
            return True
        return False

    def getenv(self, key: str, default=None):
        if key not in self.allowed:
            if self._is_safe(key):
                return self._orig_getenv(key, default)
            self._fail(key)
        return self._orig_getenv(key, default)

    def getitem(self, key: str):
        if key not in self.allowed:
            if self._is_safe(key):
                return self._orig_getitem(key)
            self._fail(key)
        return self._orig_getitem(key)

    def environ_get(self, key: str, default=None):
        if key not in self.allowed:
            if self._is_safe(key):
                return self._orig_environ_get(key, default)
            self._fail(key)
        return self._orig_environ_get(key, default)

    def install(self) -> None:
        os.getenv = self.getenv  # type: ignore[assignment]
        os.environ.__getitem__ = self.getitem  # type: ignore[assignment]
        os.environ.get = self.environ_get  # type: ignore[assignment]

    def uninstall(self) -> None:
        """Restore original environment access functions."""
        os.getenv = self._orig_getenv  # type: ignore[assignment]
        os.environ.__getitem__ = self._orig_getitem  # type: ignore[assignment]
        os.environ.get = self._orig_environ_get  # type: ignore[assignment]
