#!/usr/bin/env python3
"""Invariant R: a read-only carve-out can neither disappear nor be shadowed.

The broker's carve-out is an authentication bypass for UNVERIFIED runtimes, so
this test imports the live ``READ_ONLY_TOOLS`` object instead of copying it.
Claude MCP names are namespaced, making an exact proven allow-list sufficient.
Pi extensions are unnamespaced and may override built-ins, so the Pi half boots
the installed runtime and requires every carve-out winner to retain built-in
provenance.
"""

from __future__ import annotations

import importlib
import json
import os
import shutil
import subprocess
import sys
import tempfile
import time
import unittest
from pathlib import Path
from typing import Final


PACKAGE_ROOT = Path(__file__).parents[2]
FRAMEWORK = PACKAGE_ROOT / "framework"
LEASE_BROKER = FRAMEWORK / "tools/lease-broker"
PI_EXTENSION = FRAMEWORK / "runtime/pi/mosaic-extension.ts"
sys.path.insert(0, str(LEASE_BROKER))
daemon = importlib.import_module("daemon")
READ_ONLY_TOOLS = daemon.READ_ONLY_TOOLS

# Claude Code's measured, bare built-ins that are both registered and incapable
# of filesystem mutation or subprocess execution. MCP tools are namespaced as
# mcp__<server>__<tool>, so they cannot replace these bare identities.
CLAUDE_PROVEN_READ_ONLY_TOOLS: Final = frozenset({"Read", "Grep", "Glob"})

# W-B measured Pi 0.84.1 through getAllTools(), observed every tool_call name,
# and cross-checked dist/core/tools/index.js:18. Keep every measured built-in
# here so a runtime registry change forces the security classification to be
# revisited even when a built-in is deliberately excluded from the carve-out.
PI_VERSION: Final = "0.84.1"
PI_PROBE_ATTEMPTS: Final = 3
PI_PROBE_TIMEOUT_SECONDS: Final = 45
PI_PROBE_BACKOFF_SECONDS: Final = 0.25
PI_PROVEN_READ_ONLY_TOOLS: Final = frozenset({"read", "ls"})
PI_SUBPROCESS_TOOLS: Final = frozenset({"grep", "find"})
PI_MUTATING_TOOLS: Final = frozenset({"bash", "edit", "write"})
PI_MEASURED_BUILTINS: Final = (
    PI_PROVEN_READ_ONLY_TOOLS | PI_SUBPROCESS_TOOLS | PI_MUTATING_TOOLS
)

# Pi 0.84.1 built-ins individually proven incapable of subprocess execution or
# filesystem writes on their default path:
# - read: dist/core/tools/read.js:26-29 dispatches only read/access operations.
# - ls: dist/core/tools/ls.js:19-22 dispatches only exists/stat/readdir operations.
# grep and find are deliberately absent: grep.js:99/148 and find.js:161/203
# reach ensureTool(..., true) and spawn(), including the cold-cache download,
# write, chmod, and exec path in dist/utils/tools-manager.js:285-313.
PI_CAPABILITY_SAFE_TOOLS: Final = frozenset({"read", "ls"})

# Falsifier-only inputs. They are intentionally undocumented outside this test:
# normal CI leaves them unset; the W-A evidence run uses them to prove that the
# suite turns red for a nonexistent Claude carve-out or a Pi built-in override.
CLAUDE_EXTRA_TOOL_ENV: Final = "MOSAIC_INVARIANT_R_CLAUDE_EXTRA_TOOL"
PI_EXTRA_EXTENSION_ENV: Final = "MOSAIC_INVARIANT_R_PI_EXTRA_EXTENSION"


def run_pi_registry_command(
    command: list[str],
    environ: dict[str, str],
    *,
    runner=subprocess.run,
    sleeper=time.sleep,
) -> subprocess.CompletedProcess[str]:
    """Run the registry probe with bounded retries for concurrent-Pi stalls."""

    for attempt in range(1, PI_PROBE_ATTEMPTS + 1):
        try:
            return runner(
                command,
                check=False,
                capture_output=True,
                text=True,
                env=environ,
                timeout=PI_PROBE_TIMEOUT_SECONDS,
            )
        except subprocess.TimeoutExpired as error:
            if attempt == PI_PROBE_ATTEMPTS:
                raise AssertionError(
                    "Pi registry probe could not complete after "
                    f"{PI_PROBE_ATTEMPTS} attempts (concurrent pi?); this is a "
                    "probe/infra failure, NOT an Invariant R violation"
                ) from error
            sleeper(PI_PROBE_BACKOFF_SECONDS * attempt)

    raise AssertionError("unreachable Pi registry retry state")


def probe_pi_registry() -> list[dict[str, object]]:
    """Boot Pi's real registry and return the final winning tool definitions."""

    pi = shutil.which("pi")
    if pi is None:
        raise AssertionError("installed Pi runtime is required for Invariant R")

    version = subprocess.run(
        [pi, "--version"],
        check=False,
        capture_output=True,
        text=True,
        timeout=10,
    )
    if version.returncode != 0:
        raise AssertionError(f"Pi version probe failed: {version.stderr.strip()}")
    if version.stdout.strip() != PI_VERSION:
        raise AssertionError(
            f"Pi runtime changed from measured {PI_VERSION} to {version.stdout.strip()!r}; "
            "remeasure its registry before updating Invariant R"
        )

    with tempfile.TemporaryDirectory() as temporary:
        root = Path(temporary)
        output = root / "registry.json"
        observer = root / "registry-observer.ts"
        observer.write_text(
            "import { writeFileSync } from 'node:fs';\n"
            "export default function register(pi: any) {\n"
            "  pi.on('session_start', () => {\n"
            f"    writeFileSync({json.dumps(str(output))}, JSON.stringify(pi.getAllTools()));\n"
            "    process.exit(0);\n"
            "  });\n"
            "}\n",
            encoding="utf-8",
        )

        command = [
            pi,
            "--mode",
            "text",
            "--no-session",
            "--no-approve",
            "--no-context-files",
            "--no-skills",
            "--no-prompt-templates",
            "--no-extensions",
            "-e",
            str(observer),
            "-e",
            str(PI_EXTENSION),
        ]
        extra_extension = os.environ.get(PI_EXTRA_EXTENSION_ENV)
        if extra_extension:
            command.extend(("-e", extra_extension))
        command.append("Invariant R registry probe")

        completed = run_pi_registry_command(
            command,
            {**os.environ, "PI_OFFLINE": "1"},
        )
        if completed.returncode != 0 or not output.is_file():
            raise AssertionError(
                "Pi registry probe failed "
                f"(status {completed.returncode}): {completed.stderr.strip()}"
            )
        value = json.loads(output.read_text(encoding="utf-8"))
        if not isinstance(value, list) or not value:
            raise AssertionError("Pi registry probe returned no tools; control failed")
        return value


class InvariantRTest(unittest.TestCase):
    def test_live_carve_out_has_only_supported_runtimes(self) -> None:
        self.assertEqual(set(READ_ONLY_TOOLS), {"claude", "pi"})

    def test_claude_carve_out_is_registered_and_proven(self) -> None:
        carve_out = set(READ_ONLY_TOOLS["claude"])
        falsifier = os.environ.get(CLAUDE_EXTRA_TOOL_ENV)
        if falsifier:
            carve_out.add(falsifier)

        self.assertEqual(
            carve_out,
            set(CLAUDE_PROVEN_READ_ONLY_TOOLS),
            "every Claude carve-out must exist and be in the exact proven read-only allow-list",
        )

    def test_pi_carve_out_has_no_exec_or_write_capability(self) -> None:
        carve_out = set(READ_ONLY_TOOLS["pi"])

        capability_unsafe = carve_out - set(PI_CAPABILITY_SAFE_TOOLS)
        self.assertFalse(
            capability_unsafe,
            f"capability-unsafe Pi carve-out tools: {sorted(capability_unsafe)!r}; "
            "Pi 0.84.1 grep.js:99/148 and find.js:161/203 reach "
            "ensureTool(..., true) and spawn(), whose cold-cache path downloads, "
            "writes, chmods, and execs",
        )

    def test_pi_carve_out_resolves_to_real_unshadowed_builtins(self) -> None:
        carve_out = set(READ_ONLY_TOOLS["pi"])
        self.assertEqual(
            carve_out,
            set(PI_PROVEN_READ_ONLY_TOOLS),
            "Pi carve-out drift requires a new runtime measurement and classification",
        )
        self.assertTrue(carve_out.isdisjoint(PI_MUTATING_TOOLS))

        registry = probe_pi_registry()
        by_name: dict[str, dict[str, object]] = {}
        for entry in registry:
            name = entry.get("name")
            if not isinstance(name, str):
                self.fail(f"Pi registry entry has no string name: {entry!r}")
            by_name[name] = entry

        builtin_names = {
            name
            for name, entry in by_name.items()
            if isinstance(entry.get("sourceInfo"), dict)
            and entry["sourceInfo"].get("source") == "builtin"
        }
        self.assertEqual(
            builtin_names,
            set(PI_MEASURED_BUILTINS),
            "Pi's real built-in registry drifted from the positive-control W-B measurement",
        )

        for name in sorted(carve_out):
            with self.subTest(tool=name):
                self.assertIn(name, by_name, "Pi carve-out names must exist in the real registry")
                source = by_name[name].get("sourceInfo")
                self.assertIsInstance(source, dict)
                if isinstance(source, dict):
                    self.assertEqual(
                        source.get("source"),
                        "builtin",
                        f"Pi extension or SDK tool shadowed read-only carve-out {name!r}",
                    )
                    self.assertEqual(source.get("path"), f"<builtin:{name}>")

    def test_pi_probe_retries_timeouts_before_succeeding(self) -> None:
        attempts: list[float] = []
        backoffs: list[float] = []

        def timeout_twice(command, **kwargs):
            attempts.append(kwargs["timeout"])
            if len(attempts) < 3:
                raise subprocess.TimeoutExpired(command, kwargs["timeout"])
            return subprocess.CompletedProcess(command, 0, "", "")

        completed = run_pi_registry_command(
            ["pi", "probe"],
            {},
            runner=timeout_twice,
            sleeper=backoffs.append,
        )

        self.assertEqual(completed.returncode, 0)
        self.assertEqual(attempts, [45, 45, 45])
        self.assertEqual(backoffs, [0.25, 0.5])

    def test_pi_probe_labels_exhausted_timeouts_as_infrastructure_failure(self) -> None:
        attempts = 0

        def always_timeout(command, **kwargs):
            nonlocal attempts
            attempts += 1
            raise subprocess.TimeoutExpired(command, kwargs["timeout"])

        with self.assertRaisesRegex(
            AssertionError,
            "Pi registry probe could not complete .* NOT an Invariant R violation",
        ) as caught:
            run_pi_registry_command(
                ["pi", "probe"],
                {},
                runner=always_timeout,
                sleeper=lambda _delay: None,
            )

        self.assertEqual(attempts, 3)
        self.assertIsInstance(caught.exception.__cause__, subprocess.TimeoutExpired)


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