"""Swarm-side agent launcher tests (agent_swarm#16: Swarm launches agents + sets the limit).

Covers the pure/limit/env/resolve/backend-selection logic of orchestrator/agent_launcher:
  * launch_count + plan_launch_specs cap the pool at MAX_AGENTS_PER_USER (never push the user over);
  * each launch spec composes the agent env the runtime reads (ORCHESTRATOR_URL / AGENT_ID /
    AGENT_CAPABILITIES / OPENAI_API_BASE / OPENAI_API_KEY / HEICODE_USER_ID);
  * the model key is resolved server-side (override → azkv secret_ref dev-map → OPENAI_API_KEY),
    never fabricated;
  * `command` backend builds argv from the template; `none` backend is a no-op (no spawn).

Hermetic: no Redis / model / subprocess (backend forced to none/command-build only).

Run from agent_swarm_v6:
    python scripts/test-agent-launcher.py
"""
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from orchestrator import agent_launcher as al

failures = []


def check(name, cond):
    print(("PASS" if cond else "FAIL"), "-", name)
    if not cond:
        failures.append(name)


class FakeRun:
    def __init__(self, swarm_id="swarm-abc"):
        self.swarm_id = swarm_id


def test_launch_count():
    # Default hard pod-count window is [MIN, MAX] = [3, 16]. count = max(3, min(16, min(pool, cap-connected))).
    check("count = max(MIN, min(pool, limit-connected)) — pool 3, no connected -> 3",
          al.launch_count(pool_size=3, limit=10, connected_user_agents=0) == 3)
    # MIN floor takes PRIORITY over cap headroom: pool 5, cap 10, 8 connected -> headroom 2, floored to MIN 3.
    check("MIN floor (3) wins over cap headroom (headroom 2 -> 3)",
          al.launch_count(pool_size=5, limit=10, connected_user_agents=8) == 3)
    # Even when the user is already AT/over the cap, the MIN floor still applies (intentional [3,16]
    # hard-floor-first rule; WS register-time cap fail-closes any pods that exceed the cap).
    check("MIN floor still applies when already at cap (headroom 0 -> 3)",
          al.launch_count(pool_size=3, limit=10, connected_user_agents=10) == 3)
    check("count clamps to pool when cap is high (pool 3 -> 3)",
          al.launch_count(pool_size=3, limit=100, connected_user_agents=0) == 3)
    # MAX ceiling (16): a big pool with plenty of headroom is capped at 16.
    check("MAX ceiling (16) caps a large pool",
          al.launch_count(pool_size=50, limit=100, connected_user_agents=0) == 16)
    # A high pool but limited cap headroom: pool 50, cap 20, 12 connected -> headroom 8, within [3,16] -> 8.
    check("cap headroom (8) honored when within [MIN,MAX]",
          al.launch_count(pool_size=50, limit=20, connected_user_agents=12) == 8)
    # env-overridable window: MIN=1, MAX=4 -> headroom 2 honored (not floored), big pool capped at 4.
    os.environ["AGENT_LAUNCH_MIN_POOL"] = "1"
    os.environ["AGENT_LAUNCH_MAX_POOL"] = "4"
    check("env MIN=1 -> small headroom not floored",
          al.launch_count(pool_size=5, limit=10, connected_user_agents=8) == 2)
    check("env MAX=4 -> large pool capped at 4",
          al.launch_count(pool_size=50, limit=100, connected_user_agents=0) == 4)
    os.environ.pop("AGENT_LAUNCH_MIN_POOL")
    os.environ.pop("AGENT_LAUNCH_MAX_POOL")


def test_pool_window():
    # desired_pool_size() is clamped to [MIN, MAX] = [3, 16] by default.
    for k in ("AGENT_LAUNCH_POOL_SIZE", "AGENT_LAUNCH_MIN_POOL", "AGENT_LAUNCH_MAX_POOL"):
        os.environ.pop(k, None)
    check("desired_pool_size default = 3 (== MIN)", al.desired_pool_size() == 3)
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "1"
    check("desired_pool_size floored to MIN 3 (requested 1)", al.desired_pool_size() == 3)
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "99"
    check("desired_pool_size capped to MAX 16 (requested 99)", al.desired_pool_size() == 16)
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "8"
    check("desired_pool_size passes through within window (8)", al.desired_pool_size() == 8)
    # env-overridable window.
    os.environ["AGENT_LAUNCH_MIN_POOL"] = "5"
    os.environ["AGENT_LAUNCH_MAX_POOL"] = "6"
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "1"
    check("env MIN=5 floors desired_pool_size", al.desired_pool_size() == 5)
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "20"
    check("env MAX=6 caps desired_pool_size", al.desired_pool_size() == 6)
    check("min_pool_size()/max_pool_size() read env", al.min_pool_size() == 5 and al.max_pool_size() == 6)
    for k in ("AGENT_LAUNCH_POOL_SIZE", "AGENT_LAUNCH_MIN_POOL", "AGENT_LAUNCH_MAX_POOL"):
        os.environ.pop(k, None)


def test_plan_specs():
    run = FakeRun()
    body = {"orchestration_plan": {"objective": "x"}, "billing_context": {"default_model_id": "gpt-x"}}
    # 8 connected, cap 10 -> headroom 2, but MIN floor 3 wins (intentional [3,16] hard-floor-first).
    specs = al.plan_launch_specs(run, body, connected_user_agents=8, limit=10, pool_size=3,
                                 model_key="sk-test", orchestrator_url="ws://orch:8000", user_id="u-1")
    check("plan applies MIN floor over cap headroom (8 connected, cap 10 -> 3 not 2)", len(specs) == 3)
    s = specs[0]
    check("spec env has ORCHESTRATOR_URL", s.env.get("ORCHESTRATOR_URL") == "ws://orch:8000")
    check("spec env has model key (server-side injected)", s.env.get("OPENAI_API_KEY") == "sk-test")
    check("spec env has model id from billing_context", s.env.get("OPENAI_MODEL") == "gpt-x")
    check("spec env has HEICODE_USER_ID for per-user cap", s.env.get("HEICODE_USER_ID") == "u-1")
    check("spec has AGENT_ID + capabilities", bool(s.agent_id) and bool(s.env.get("AGENT_CAPABILITIES")))
    check("agent ids unique", len({sp.agent_id for sp in specs}) == len(specs))
    # No key -> OPENAI_API_KEY omitted (not fabricated), no user -> HEICODE_USER_ID omitted.
    specs2 = al.plan_launch_specs(run, body, connected_user_agents=0, limit=10, pool_size=1,
                                  model_key=None, orchestrator_url="ws://orch", user_id=None)
    check("no model key -> OPENAI_API_KEY omitted", "OPENAI_API_KEY" not in specs2[0].env)
    check("no user -> HEICODE_USER_ID omitted", "HEICODE_USER_ID" not in specs2[0].env)


def test_resolve_model_key():
    for k in ("AGENT_LAUNCH_MODEL_KEY", "OPENAI_API_KEY", "HEICODE_SECRET_res_model_1"):
        os.environ.pop(k, None)
    # override wins
    os.environ["AGENT_LAUNCH_MODEL_KEY"] = "sk-override"
    check("override key wins", al.resolve_model_key({"billing_context": {"secret_ref": "azkv://kv/secrets/res_model_1"}}) == "sk-override")
    os.environ.pop("AGENT_LAUNCH_MODEL_KEY")
    # azkv secret_ref -> dev env map
    os.environ["HEICODE_SECRET_res_model_1"] = "sk-from-kv"
    check("azkv secret_ref resolved via dev map", al.resolve_model_key({"billing_context": {"secret_ref": "azkv://kv/secrets/res_model_1"}}) == "sk-from-kv")
    os.environ.pop("HEICODE_SECRET_res_model_1")
    # fallback to orchestrator OPENAI_API_KEY
    os.environ["OPENAI_API_KEY"] = "sk-orch"
    check("fallback to orchestrator OPENAI_API_KEY", al.resolve_model_key({}) == "sk-orch")
    os.environ.pop("OPENAI_API_KEY")
    check("unresolved -> None (never fabricated)", al.resolve_model_key({"billing_context": {"secret_ref": "azkv://kv/secrets/missing"}}) is None)


def test_azkv_resolver():
    # ── parse azkv:// refs (pure) ──
    check("azkv parse: bare vault name -> https URL",
          al._parse_azkv_ref("azkv://heicode-vault/secrets/swarm-model-key-u1")
          == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", None))
    check("azkv parse: full host preserved",
          al._parse_azkv_ref("azkv://heicode-vault.vault.azure.net/secrets/res_git_1")
          == ("https://heicode-vault.vault.azure.net", "res_git_1", None))
    check("azkv parse: version captured",
          al._parse_azkv_ref("azkv://heicode-vault/secrets/swarm-model-key-u1/abc")
          == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", "abc"))
    check("azkv parse: short form azkv://<vault>/<name>",
          al._parse_azkv_ref("azkv://heicode-vault/swarm-model-key-u1")
          == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", None))
    check("azkv parse: bad ref -> None", al._parse_azkv_ref("azkv://heicode-vault") is None)

    # ── gating: hermetic unless workload identity injected / opt-in ──
    for k in ("AZURE_FEDERATED_TOKEN_FILE", "SECRET_RESOLVER"):
        os.environ.pop(k, None)
    check("azkv disabled without workload identity / opt-in", al._azkv_enabled() is False)
    check("disabled -> _resolve_from_keyvault returns None (no SDK/network touched)",
          al._resolve_from_keyvault("azkv://heicode-vault/secrets/swarm-model-key-u1") is None)
    os.environ["SECRET_RESOLVER"] = "azkv"
    check("azkv enabled via SECRET_RESOLVER=azkv", al._azkv_enabled() is True)
    os.environ.pop("SECRET_RESOLVER")
    os.environ["AZURE_FEDERATED_TOKEN_FILE"] = "/var/run/secrets/azure/tokens/azure-identity-token"
    check("azkv enabled via AZURE_FEDERATED_TOKEN_FILE", al._azkv_enabled() is True)
    os.environ.pop("AZURE_FEDERATED_TOKEN_FILE")


def test_command_backend_build():
    os.environ["AGENT_LAUNCH_CMD"] = "launch-agent.sh --id {agent_id} --caps {capabilities}"
    spec = al.AgentLaunchSpec(agent_id="swarm-abc-agent-1", capabilities="python,general",
                              env={"OPENAI_API_KEY": "sk-x"})
    argv = al.build_launch_command(spec)
    check("command template substitutes agent_id/capabilities",
          argv == ["launch-agent.sh", "--id", "swarm-abc-agent-1", "--caps", "python,general"])
    check("secret not on argv (passed via env)", all("sk-x" not in a for a in argv))
    os.environ.pop("AGENT_LAUNCH_CMD")


async def test_backend_none_noop():
    import asyncio  # noqa
    os.environ["AGENT_LAUNCH_BACKEND"] = "none"
    specs = [al.AgentLaunchSpec(agent_id="a1", capabilities="general", env={})]
    launched = await al.launch(specs, swarm_id="swarm-abc")
    check("backend=none launches nothing (external)", launched == [])
    os.environ.pop("AGENT_LAUNCH_BACKEND")


def test_k8s_manifests():
    import json
    spec = al.AgentLaunchSpec(
        agent_id="swarm-abc-agent-1", capabilities="python,general",
        env={"ORCHESTRATOR_URL": "ws://orch.svc:8000", "AGENT_ID": "swarm-abc-agent-1",
             "AGENT_CAPABILITIES": "python,general", "OPENAI_API_KEY": "sk-secret-xyz",
             "OPENAI_API_BASE": "https://hm/v1", "HEICODE_USER_ID": "u-1"},
    )
    # Secret manifest carries the key (applied via stdin, not argv).
    sec = al.build_secret_manifest(key_value="sk-secret-xyz", namespace="heicode-swarm",
                                   name="swarm-agent-key-swarm-abc", labels={"app": "heicode-swarm-agent"})
    check("secret kind/type", sec["kind"] == "Secret" and sec["type"] == "Opaque")
    check("secret stringData has the key", sec["stringData"]["OPENAI_API_KEY"] == "sk-secret-xyz")

    # Pod manifest references the key via secretKeyRef — NEVER inline.
    pod = al.build_pod_manifest(spec, namespace="heicode-swarm", swarm_id="swarm-abc",
                                image="img:1", secret_name="swarm-agent-key-swarm-abc",
                                service_account="swarm-agent-sa")
    blob = json.dumps(pod)
    check("pod kind/name", pod["kind"] == "Pod" and pod["metadata"]["name"] == "swarm-abc-agent-1")
    check("pod labels include swarm-id + user-id",
          pod["metadata"]["labels"].get("heicode-swarm-id") == "swarm-abc"
          and pod["metadata"]["labels"].get("heicode-user-id") == "u-1")
    env = {e["name"]: e for e in pod["spec"]["containers"][0]["env"]}
    check("OPENAI_API_KEY via secretKeyRef (not inline)",
          "value" not in env["OPENAI_API_KEY"] and env["OPENAI_API_KEY"]["valueFrom"]["secretKeyRef"]["name"] == "swarm-agent-key-swarm-abc")
    check("raw key value NOT in pod manifest (only in Secret)", "sk-secret-xyz" not in blob)
    check("non-secret env inline", env["ORCHESTRATOR_URL"]["value"] == "ws://orch.svc:8000")
    check("pod has resource limits + restartPolicy + serviceAccount",
          pod["spec"]["containers"][0]["resources"]["limits"]
          and pod["spec"]["restartPolicy"] == "OnFailure"
          and pod["spec"]["serviceAccountName"] == "swarm-agent-sa")
    # No secret_name -> key omitted entirely (never inlined).
    pod2 = al.build_pod_manifest(spec, namespace="heicode-swarm", swarm_id="swarm-abc", secret_name=None)
    env2 = {e["name"] for e in pod2["spec"]["containers"][0]["env"]}
    check("no secret -> OPENAI_API_KEY omitted from pod env", "OPENAI_API_KEY" not in env2)


def main():
    import asyncio
    test_launch_count()
    test_pool_window()
    test_plan_specs()
    test_resolve_model_key()
    test_azkv_resolver()
    test_command_backend_build()
    test_k8s_manifests()
    asyncio.run(test_backend_none_noop())
    print()
    if failures:
        print(f"{len(failures)} agent-launcher check(s) FAILED: {failures}")
        sys.exit(1)
    print("all agent-launcher checks passed")


if __name__ == "__main__":
    main()