server/env.py

import json
import os
import subprocess
import time
from pathlib import Path
from typing import Any

from .constants import (
    DEFAULT_CONFIGS,
    NO_COMMAND_PROVIDED_SENTINEL,
    TASK_MAX_STEPS,
    TaskName,
)
from .fault_injector import inject_fault
from .graders import grade_task
from .metrics_poller import MetricsPoller
from .models import Action, Observation, StepResult
from .process_manager import ProcessManager


class DistributedDebugEnv:
    """OpenEnv-compatible distributed systems debugging environment."""

    def __init__(
        self, project_root: Path | None = None, mesh_root: Path | None = None
    ) -> None:
        self.project_root = (
            project_root or Path(__file__).resolve().parent.parent
        ).resolve()
        self.mesh_root = (
            mesh_root or Path(os.getenv("MESH_ROOT", self.project_root / "mesh"))
        ).resolve()

        self._process_manager = ProcessManager(
            project_root=self.project_root, mesh_root=self.mesh_root
        )
        self._metrics_poller = MetricsPoller(poll_interval_s=2.0)

        self.current_task: TaskName | None = None
        self.max_steps: int = 0
        self.step_count: int = 0
        self.last_exit_code: int = 0
        self.prev_observation: Observation | None = None
        self._baselines: dict[str, int] = {
            "baseline_worker_restart_count": 0,
            "baseline_consumer_stall_count": 0,
        }
        self._seen_diagnostic_signatures: set[str] = set()
        self._command_counts: dict[str, int] = {}
        self._last_grader_score: float = 0.0

    def start(self) -> None:
        if not self._metrics_poller.is_alive():
            self._metrics_poller.start()

    def close(self) -> None:
        self._metrics_poller.stop()

    def _write_json(self, path: Path, payload: dict[str, Any]) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")

    def _restore_defaults(self) -> None:
        self._write_json(
            self.mesh_root / "registry.json",
            {
                "services": {
                    "auth": {"host": "localhost", "port": 3001, "protocol": "http"},
                    "redis": {"host": "localhost", "port": 6379, "protocol": "tcp"},
                    "worker": {
                        "host": "localhost",
                        "port": None,
                        "protocol": "internal",
                    },
                }
            },
        )
        self._write_json(
            self.mesh_root / "auth" / "config.json", DEFAULT_CONFIGS["auth"]
        )
        self._write_json(
            self.mesh_root / "gateway" / "config.json", DEFAULT_CONFIGS["gateway"]
        )
        self._write_json(
            self.mesh_root / "gateway" / "blocked_routes.json",
            DEFAULT_CONFIGS["blocked_routes"],
        )
        self._write_json(
            self.mesh_root / "worker" / "config.json", DEFAULT_CONFIGS["worker"]
        )
        self._write_json(
            self.mesh_root / "worker" / "job_generator_config.json",
            DEFAULT_CONFIGS["job_generator"],
        )

    def _truncate_logs(self) -> None:
        for service in ["gateway", "auth", "worker", "job_gen"]:
            Path(f"/tmp/{service}.log").write_text("", encoding="utf-8")

    def _reset_runtime_counters(self) -> None:
        Path("/tmp/worker_restart_count").write_text("0", encoding="utf-8")
        Path("/tmp/consumer_stall_count").write_text("0", encoding="utf-8")

    def _redis_flush(self) -> None:
        subprocess.run(
            ["redis-cli", "FLUSHDB"], check=True, capture_output=True, text=True
        )

    def _read_float(self, value: str, default: float = 0.0) -> float:
        try:
            return float(value)
        except (TypeError, ValueError):
            return default

    def _is_route_blocked(self) -> bool:
        blocked_file = self.mesh_root / "gateway" / "blocked_routes.json"
        try:
            payload = json.loads(blocked_file.read_text(encoding="utf-8"))
            blocked = payload.get("blocked", [])
            return "gateway->redis" in blocked
        except Exception:
            return False

    def _is_lock_present(self) -> bool:
        result = subprocess.run(
            ["redis-cli", "EXISTS", "LOCK:job_processor"],
            capture_output=True,
            text=True,
            timeout=2,
            check=False,
        )
        return result.stdout.strip() == "1"

    def _is_cascading_timeout_resolved(self) -> bool:
        auth_config_file = self.mesh_root / "auth" / "config.json"
        gateway_config_file = self.mesh_root / "gateway" / "config.json"
        try:
            auth_payload = json.loads(auth_config_file.read_text(encoding="utf-8"))
            gateway_payload = json.loads(
                gateway_config_file.read_text(encoding="utf-8")
            )
        except Exception:
            return False

        auth_delay_ms = self._read_float(auth_payload.get("delay_ms"), default=0.0)
        auth_timeout_ms = self._read_float(
            gateway_payload.get("auth_timeout_ms"), default=0.0
        )
        if auth_timeout_ms <= 0:
            return False
        return auth_delay_ms <= auth_timeout_ms

    def _is_registry_auth_default(self) -> bool:
        registry_file = self.mesh_root / "registry.json"
        try:
            payload = json.loads(registry_file.read_text(encoding="utf-8"))
            auth_service = payload["services"]["auth"]
        except Exception:
            return False

        return (
            auth_service.get("host") == "localhost"
            and int(auth_service.get("port", 0)) == 3001
            and auth_service.get("protocol") == "http"
        )

    def _job_generator_interval_ms(self) -> int:
        config_file = self.mesh_root / "worker" / "job_generator_config.json"
        try:
            payload = json.loads(config_file.read_text(encoding="utf-8"))
        except Exception:
            return 0

        try:
            return int(payload.get("interval_ms", 0))
        except (TypeError, ValueError):
            return 0

    def _is_job_generator_rate_resolved(self) -> bool:
        return self._job_generator_interval_ms() >= int(
            DEFAULT_CONFIGS["job_generator"]["interval_ms"]
        )

    def _build_grader_context(self) -> dict[str, Any]:
        return {
            **self._baselines,
            "route_blocked": self._is_route_blocked(),
            "lock_exists": self._is_lock_present(),
            "cascading_timeout_resolved": self._is_cascading_timeout_resolved(),
            "registry_auth_matches_default": self._is_registry_auth_default(),
            "job_generator_interval_ms": self._job_generator_interval_ms(),
            "job_generator_rate_resolved": self._is_job_generator_rate_resolved(),
        }

    def _blocked_command(self, command: str) -> bool:
        dangerous_patterns = [
            "rm -rf /",
            "kill -9 1",
            "pkill -f uvicorn",
            "> /tmp/gateway.log",
            "> /tmp/auth.log",
            "> /tmp/worker.log",
        ]
        normalized = command.strip().lower()
        return any(pattern in normalized for pattern in dangerous_patterns)

    def _run_command(self, command: str) -> tuple[str, str | None]:
        if command.strip() == NO_COMMAND_PROVIDED_SENTINEL:
            self.last_exit_code = 2
            return (
                "No command provided by model. Expected JSON with a command field.",
                "no_command_provided",
            )

        if self._blocked_command(command):
            self.last_exit_code = 1
            return (
                "BLOCKED: This command would damage the environment infrastructure.",
                "blocked_command",
            )

        try:
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=10,
                cwd="/",
                env={
                    **os.environ,
                    "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
                },
                check=False,
            )
            self.last_exit_code = result.returncode
            output = (result.stdout + result.stderr).strip() or "(no output)"
            return output, None
        except subprocess.TimeoutExpired:
            self.last_exit_code = 124
            return "Command timed out after 10 seconds.", "timeout"
        except Exception as exc:
            self.last_exit_code = 1
            return f"Command execution error: {exc}", str(exc)

    def _command_signature(self, command: str) -> str:
        return " ".join(command.strip().lower().split())

    def _is_diagnostic_command(self, command: str) -> bool:
        diagnostic_keywords = [
            "cat",
            "curl",
            "redis-cli",
            "ps",
            "ls",
            "grep",
            "tail",
            "jq",
            "lrange",
            "llen",
            "keys",
            "ttl",
            "get",
        ]
        normalized = command.lower()
        return any(keyword in normalized for keyword in diagnostic_keywords)

    def _is_state_change_command(self, command: str) -> bool:
        normalized = command.lower()
        state_change_patterns = [
            "kill -hup",
            "redis-cli del",
            "redis-cli lrem",
            "redis-cli set",
            "redis-cli flushdb",
            "echo '{",
            "> /mesh/",
            "tee /mesh/",
        ]
        return any(pattern in normalized for pattern in state_change_patterns)

    def _compute_reward(
        self,
        command: str,
        current: Observation,
        previous: Observation,
        grader_score: float,
        previous_grader_score: float,
        command_error: str | None,
    ) -> float:
        if command_error == "no_command_provided":
            return 0.01

        if grader_score >= 0.95:
            return 0.99

        reward = grader_score * 0.75
        signature = self._command_signature(command)
        signature_count = self._command_counts.get(signature, 0) + 1
        self._command_counts[signature] = signature_count

        if (
            self._is_diagnostic_command(command)
            and signature not in self._seen_diagnostic_signatures
        ):
            reward += 0.02
            self._seen_diagnostic_signatures.add(signature)

        if self._is_state_change_command(command):
            reward += 0.03

        if grader_score > previous_grader_score + 1e-4:
            reward += 0.15
        else:
            reward -= 0.05

        if (
            current.metrics.gateway_success_rate
            > previous.metrics.gateway_success_rate + 1e-3
        ):
            reward += 0.05

        if current.metrics.queue_depth < previous.metrics.queue_depth:
            reward += 0.05

        if current.metrics.worker_restart_count < previous.metrics.worker_restart_count:
            reward += 0.03

        if current.metrics.consumer_stall_count < previous.metrics.consumer_stall_count:
            reward += 0.03

        if signature_count > 1:
            reward -= min(0.12, 0.04 * (signature_count - 1))

        if command.strip().lower() in {
            "echo",
            "pwd",
            "whoami",
            "date",
            "true",
            "false",
        }:
            reward -= 0.08

        if self.last_exit_code != 0 and command_error not in {
            "blocked_command",
            "no_command_provided",
        }:
            reward -= 0.08

        if command_error == "blocked_command":
            reward -= 0.25

        return max(0.01, min(0.99, reward))

    def _status_block(self, metrics: Any) -> str:
        return (
            "=== pipeline status after reset ===\n"
            "gateway:  running\n"
            "auth:     running\n"
            "worker:   running\n"
            f"queue_depth: {metrics.queue_depth}\n"
            f"gateway_success_rate: {metrics.gateway_success_rate:.2f}"
        )

    def reset(self, task_name: TaskName | str) -> Observation:
        task = TaskName.parse(task_name) if isinstance(task_name, str) else task_name

        self.current_task = task
        self.max_steps = TASK_MAX_STEPS[task]
        self.step_count = 0
        self._seen_diagnostic_signatures = set()
        self._command_counts = {}
        self._last_grader_score = 0.0

        self._truncate_logs()
        self._restore_defaults()
        self._redis_flush()
        self._reset_runtime_counters()

        Path("/tmp/current_task").write_text(task.value, encoding="utf-8")

        self._process_manager.restart_all()
        if not self._process_manager.wait_healthy(timeout_s=30):
            raise RuntimeError("Services failed health checks after reset")

        inject_fault(task, self._process_manager)
        time.sleep(1.0)

        self._metrics_poller.poll_once()
        metrics = self._metrics_poller.get_current_metrics()

        self._baselines = {
            "baseline_worker_restart_count": metrics.worker_restart_count,
            "baseline_consumer_stall_count": metrics.consumer_stall_count,
        }
        self._last_grader_score = grade_task(
            task, metrics, self._build_grader_context()
        )

        observation = Observation(
            command_output=self._status_block(metrics),
            metrics=metrics,
            process_status=self._process_manager.get_status(),
        )
        self.prev_observation = observation
        return observation

    def step(self, action: Action) -> StepResult:
        if not self.current_task:
            raise RuntimeError(
                "Environment not initialized. Call reset(task_name) first."
            )

        self.step_count += 1
        command_output, command_error = self._run_command(action.command)

        self._metrics_poller.poll_once()
        metrics = self._metrics_poller.get_current_metrics()

        observation = Observation(
            command_output=command_output,
            metrics=metrics,
            process_status=self._process_manager.get_status(),
        )

        previous = self.prev_observation or observation
        previous_grader_score = self._last_grader_score
        grader_score = grade_task(
            self.current_task, metrics, self._build_grader_context()
        )
        reward = self._compute_reward(
            action.command,
            observation,
            previous,
            grader_score,
            previous_grader_score,
            command_error,
        )
        if command_error == "no_command_provided":
            done = self.step_count >= self.max_steps
        else:
            done = grader_score >= 0.95 or self.step_count >= self.max_steps

        self._last_grader_score = grader_score
        self.prev_observation = observation

        info: dict[str, Any] = {
            "grader_score": round(grader_score, 4),
            "error": command_error,
            "exit_code": self.last_exit_code,
            "task": self.current_task.value if self.current_task else None,
        }

        return StepResult(observation=observation, reward=reward, done=done, info=info)

    def state(self) -> dict[str, Any]:
        self._metrics_poller.poll_once()
        metrics = self._metrics_poller.get_current_metrics()
        return {
            "task": self.current_task.value if self.current_task else None,
            "step_count": self.step_count,
            "max_steps": self.max_steps,
            "metrics": metrics.model_dump(),
            "process_status": self._process_manager.get_status(),
            "baselines": dict(self._baselines),
        }