from __future__ import annotations

import pytest

from investibook.clients.run9 import Run9ApiError
from investibook.core import run9_sandbox
from investibook.core.run9_sandbox import Run9SandboxProvider
from investibook.core.run9_workspace_storage import Run9WorkspaceStorage
from investibook.core.sandbox import SandboxAgentProcessSpec


class FakeClient:
    def __init__(self, *, created: bool = True) -> None:
        self.closed = False
        self.created = created
        self.ensure_kwargs = None
        self.validated = None
        self.secret_calls = []

    async def ensure_box(self, **kwargs):
        self.ensure_kwargs = kwargs
        labels = kwargs["create_labels"] if self.created else kwargs["labels"]
        return {"id": kwargs["box_id"], "labels": labels}

    async def get_box(self, box_id):
        if self.created and "-secret-" not in box_id:
            raise Run9ApiError(
                method="GET",
                path=f"/boxes/{box_id}",
                status_code=404,
                detail="box not found",
            )
        return {"id": box_id, "labels": {}}

    def validate_box_labels(self, **kwargs):
        self.validated = kwargs

    async def list_box_secrets(self, box_id):
        return []

    async def upsert_box_secret(self, **kwargs):
        self.secret_calls.append(kwargs)
        return {"placeholder": "run9secret_test123"}

    async def aclose(self):
        self.closed = True


def test_run9_provider_resolves_image_commands(monkeypatch):
    provider = Run9SandboxProvider()
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_OPENCODE_BIN", "/bin/opencode")
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_REMOTE_AGENTKIT_DIR", "/agentkit")
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_PI_CORE_BIN", "/bin/pi-core")
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_NODE_BIN", "/bin/node")

    assert provider.agent_command(
        runtime_name="opencode",
        binary_names=("opencode",),
        command_args=("acp",),
    ) == ["/bin/opencode", "acp"]
    pi_command = provider.agent_command(
        runtime_name="pi",
        binary_names=("node",),
        command_args=(),
    )
    assert pi_command[-3:] == [
        "/bin/pi-core",
        "/bin/node",
        "/agentkit/runtime/pi-core/index.mjs",
    ]
    assert 'if [ -x "$pi_bin" ]' in pi_command[2]
    assert "sandbox_transport_ready" not in pi_command[2]


@pytest.mark.asyncio
async def test_run9_provider_owns_box_bootstrap_and_process_start(monkeypatch):
    provider = Run9SandboxProvider()
    client = FakeClient()
    bootstrap_calls = []
    process_calls = []

    monkeypatch.setattr(run9_sandbox, "_run9_client", lambda: client)
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_NETWORK_MODE", "managed")
    monkeypatch.setattr(
        run9_sandbox.cfg,
        "RUN9_BOX_SECRET_RULES_JSON",
        """[{"name":"openai","env_names":["OPENAI_API_KEY"],"allowed_hosts":["api.openai.com"],"header_name":"Authorization"}]""",
    )
    monkeypatch.setattr(
        run9_sandbox.SharedAgentProviderConfig,
        "load",
        lambda: None,
    )

    async def bootstrap(**kwargs):
        bootstrap_calls.append(kwargs)

    async def start_process(**kwargs):
        process_calls.append(kwargs)
        return object()

    monkeypatch.setattr(run9_sandbox, "bootstrap_run9_workspace", bootstrap)
    monkeypatch.setattr(run9_sandbox.Run9ProcessTransport, "start", start_process)

    process = await provider.spawn_agent_process(
        SandboxAgentProcessSpec(
            workspace_id="ws-1",
            owner_user_id="user-1",
            runtime_name="opencode",
            runtime_transport="acp",
            command=("opencode", "acp"),
            env={"OPENAI_API_KEY": "test-key", "UNRELATED": "ignored"},
            stdout_limit=4096,
        )
    )

    assert process is not None
    assert client.ensure_kwargs["labels"]["owner_user_id"] == "user-1"
    assert client.ensure_kwargs["box_id"] == "ib-dev-secret-ws-1"
    assert bootstrap_calls[0]["workspace_id"] == "ws-1"
    assert process_calls[0]["command"] == ("opencode", "acp")
    assert process_calls[0]["env"]["OPENAI_API_KEY"] == "run9secret_test123"
    assert client.secret_calls[0]["value"] == "test-key"
    assert "test-key" not in process_calls[0]["env"].values()
    assert "UNRELATED" not in process_calls[0]["env"]


def test_run9_environment_never_forwards_sensitive_values_without_secret_rules(
    monkeypatch,
):
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_BOX_SECRET_RULES_JSON", "[]")
    monkeypatch.setattr(
        run9_sandbox.cfg,
        "RUN9_FORWARD_ENV_VARS",
        ("CUSTOM_VALUE", "CUSTOM_API_KEY", "JIRA_PASSWORD"),
    )

    forwarded = run9_sandbox._forwarded_environment(
        {
            "CUSTOM_VALUE": "safe",
            "CUSTOM_API_KEY": "custom-secret",
            "JIRA_PASSWORD": "jira-secret",
            "GITHUB_TOKEN": "github-secret",
        }
    )

    assert forwarded == {"CUSTOM_VALUE": "safe"}


def test_legacy_run9_environment_preserves_previous_forwarding(monkeypatch):
    monkeypatch.setattr(
        run9_sandbox.cfg,
        "RUN9_FORWARD_ENV_VARS",
        ("CUSTOM_API_KEY",),
    )

    forwarded = run9_sandbox._legacy_forwarded_environment(
        {"CLINIC_API_KEY": "clinic-secret", "CUSTOM_API_KEY": "custom-secret"}
    )

    assert forwarded == {
        "CLINIC_API_KEY": "clinic-secret",
        "CUSTOM_API_KEY": "custom-secret",
    }


@pytest.mark.asyncio
async def test_existing_unlabelled_box_does_not_create_secrets(monkeypatch):
    provider = Run9SandboxProvider()
    client = FakeClient(created=False)
    process_calls = []
    monkeypatch.setattr(run9_sandbox, "_run9_client", lambda: client)
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_NETWORK_MODE", "managed")
    monkeypatch.setattr(
        run9_sandbox.cfg,
        "RUN9_BOX_SECRET_RULES_JSON",
        """[{"name":"clinic","env_names":["CLINIC_API_KEY"],"allowed_hosts":["clinic.pingcap.com"],"header_name":"Authorization"}]""",
    )
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_WORKSPACE_PREINITIALIZED", True)
    monkeypatch.setattr(run9_sandbox.SharedAgentProviderConfig, "load", lambda: None)

    async def start_process(**kwargs):
        process_calls.append(kwargs)
        return object()

    monkeypatch.setattr(run9_sandbox.Run9ProcessTransport, "start", start_process)

    await provider.spawn_agent_process(
        SandboxAgentProcessSpec(
            workspace_id="ws-old",
            owner_user_id="user-1",
            runtime_name="opencode",
            runtime_transport="acp",
            command=("opencode", "acp"),
            env={"CLINIC_API_KEY": "legacy-secret"},
            stdout_limit=4096,
        )
    )

    assert process_calls[0]["env"]["CLINIC_API_KEY"] == "legacy-secret"
    assert client.secret_calls == []
    assert client.ensure_kwargs["box_id"] == "ib-dev-ws-old"


@pytest.mark.asyncio
async def test_run9_box_secrets_require_managed_network(monkeypatch):
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_NETWORK_MODE", "normal")
    monkeypatch.setattr(
        run9_sandbox.cfg,
        "RUN9_BOX_SECRET_RULES_JSON",
        """[{"name":"clinic","env_names":["CLINIC_API_KEY"],"allowed_hosts":["clinic.pingcap.com"],"header_name":"Authorization"}]""",
    )

    with pytest.raises(RuntimeError, match="RUN9_NETWORK_MODE=managed"):
        await run9_sandbox._secret_safe_environment(
            client=FakeClient(),
            box_id="box-1",
            env={"CLINIC_API_KEY": "real-key"},
        )


@pytest.mark.asyncio
async def test_run9_provider_skips_bootstrap_for_preinitialized_snapshot(monkeypatch):
    provider = Run9SandboxProvider()
    client = FakeClient()
    process_calls = []

    monkeypatch.setattr(run9_sandbox, "_run9_client", lambda: client)
    monkeypatch.setattr(run9_sandbox.cfg, "RUN9_WORKSPACE_PREINITIALIZED", True)
    monkeypatch.setattr(
        run9_sandbox.SharedAgentProviderConfig,
        "load",
        lambda: None,
    )

    async def fail_bootstrap(**kwargs):
        raise AssertionError("bootstrap should not run")

    async def start_process(**kwargs):
        process_calls.append(kwargs)
        return object()

    monkeypatch.setattr(run9_sandbox, "bootstrap_run9_workspace", fail_bootstrap)
    monkeypatch.setattr(run9_sandbox.Run9ProcessTransport, "start", start_process)

    await provider.spawn_agent_process(
        SandboxAgentProcessSpec(
            workspace_id="ws-1",
            owner_user_id="user-1",
            runtime_name="pi",
            runtime_transport="pi_rpc",
            command=("node", "pi-core.mjs"),
            env={},
            stdout_limit=4096,
        )
    )

    assert process_calls[0]["env"]["INVESTIBOOK_PI_CORE_CONFIG_JSON"] == "{}"
    command = process_calls[0]["command"]
    assert command[:2] == ("sh", "-c")
    assert command[3] == "investibook-preinitialized"
    assert "/workspace/.investibook/runtime-bootstrap.json" in command
    assert command[-2:] == ("node", "pi-core.mjs")


@pytest.mark.asyncio
async def test_run9_provider_limits_single_configured_provider(monkeypatch):
    provider = Run9SandboxProvider()
    client = FakeClient()
    process_calls = []
    provider_config = type(
        "ProviderConfig",
        (),
        {
            "payload": {"provider": {"chat": {"models": {}}}},
            "pi_models_payload": lambda self: {},
        },
    )()

    monkeypatch.setattr(run9_sandbox, "_run9_client", lambda: client)
    monkeypatch.setattr(
        run9_sandbox.SharedAgentProviderConfig,
        "load",
        lambda: provider_config,
    )

    async def bootstrap(**kwargs):
        return None

    async def start_process(**kwargs):
        process_calls.append(kwargs)
        return object()

    monkeypatch.setattr(run9_sandbox, "bootstrap_run9_workspace", bootstrap)
    monkeypatch.setattr(run9_sandbox.Run9ProcessTransport, "start", start_process)

    await provider.spawn_agent_process(
        SandboxAgentProcessSpec(
            workspace_id="ws-1",
            owner_user_id="user-1",
            runtime_name="opencode",
            runtime_transport="acp",
            command=("opencode", "acp"),
            env={
                "OPENCODE_ACP_PROFILE": "1",
                "OPENCODE_ACP_DIRECT_STARTUP": "1",
                "OPENCODE_DISABLE_DEFAULT_PLUGINS": "1",
            },
            stdout_limit=4096,
        )
    )

    config = __import__("json").loads(
        process_calls[0]["env"]["OPENCODE_CONFIG_CONTENT"]
    )
    assert config["enabled_providers"] == ["chat"]
    assert process_calls[0]["env"]["OPENCODE_ACP_PROFILE"] == "1"
    assert process_calls[0]["env"]["OPENCODE_ACP_DIRECT_STARTUP"] == "1"
    assert process_calls[0]["pull_wait_ms"] == 50


@pytest.mark.asyncio
async def test_run9_provider_opens_owner_validated_storage(monkeypatch):
    provider = Run9SandboxProvider()
    client = FakeClient(created=False)
    monkeypatch.setattr(run9_sandbox, "_run9_client", lambda: client)

    storage = await provider.open_workspace_storage(
        workspace_id="ws-1",
        owner_user_id="user-1",
    )

    assert isinstance(storage, Run9WorkspaceStorage)
    assert client.validated["expected"]["owner_user_id"] == "user-1"
    await storage.aclose()
    assert client.closed is True