#!/usr/bin/env python3
"""Attach a deny policy, dry-run it, open break-glass for an outage, then put
the guardrails back.

Pass only reason + duration for break-glass. Aiden records who opened it from
your API token (``activated_by`` on the response). Prefer ``any_deny_blocks``
when several policies can disagree: one deny stops the tool.

Usage:
  export STACKGEN_URL=https://app.stackgen.com
  export STACKGEN_TOKEN=stackgen_…
  export STACKGEN_WORKSPACE_ID=<workspace-uuid>
  python examples/automation/policy-break-glass.py --agent my-oncall-agent
"""

from __future__ import annotations

import argparse
import os
import sys
import uuid
from datetime import datetime, timezone

from stackgen import StackgenClient, StackgenConfig

DENY_MARKER = "sdk-btg-emergency-marker"

WIRE_TO_RULE = {
    "deny_overrides": "any_deny_blocks",
    "permit_overrides": "any_allow_wins",
    "first_applicable": "first_matching_policy",
    "any_deny_blocks": "any_deny_blocks",
    "any_allow_wins": "any_allow_wins",
    "first_matching_policy": "first_matching_policy",
}

DENY_REGO = f"""package policy

import rego.v1

default allow := true

allow := false if {{
  endswith(input.tool.name, "_execute_parallel")
  some c in input.tool.arguments.commands
  contains(c.command, "{DENY_MARKER}")
}}

deny_reason := "sdk sample: emergency marker denied without break-glass" if {{
  not allow
}}
"""


def require_env(name: str, *fallbacks: str) -> str:
    for key in (name, *fallbacks):
        value = (os.environ.get(key) or "").strip()
        if value:
            return value
    raise SystemExit(f"missing env {name}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--agent", required=True, help="existing agent name")
    args = parser.parse_args()
    agent_name = args.agent

    client = StackgenClient(
        StackgenConfig(
            base_url=require_env("STACKGEN_URL"),
            api_token=require_env("STACKGEN_TOKEN"),
            workspace_id=require_env("STACKGEN_WORKSPACE_ID", "STACKGEN_PROJECT"),
        )
    )

    suffix = uuid.uuid4().hex[:8]
    policy_name = f"sdk-btg-deny-{suffix}"
    policy_id = ""
    shadow_attached = False
    attached = False
    glass_open = False
    previous = client.aiden.get_agent(agent_name)
    previous_rule = WIRE_TO_RULE.get(
        str(previous.get("combining_algorithm") or previous.get("combiningAlgorithm") or "")
    )

    created = client.aiden.create_policy(
        {
            "name": policy_name,
            "description": "SDK sample: deny emergency marker command",
            "type": "logic",
            "rego_source": DENY_REGO,
        }
    )
    policy_id = str(created.get("id") or "")
    if not policy_id:
        raise SystemExit(f"create_policy missing id: {created!r}")
    print(f"created policy id={policy_id}")

    try:
        client.aiden.attach_policy(
            agent_name, policy_id=policy_id, shadow=True, priority=10
        )
        shadow_attached = True
        print("attached in shadow (logged, not enforced)")

        client.aiden.when_policies_disagree(agent_name, "any_deny_blocks")
        print("conflict rule: any_deny_blocks")

        agent_info = client.aiden.get_agent(agent_name)
        agent_id = str(agent_info.get("id") or "").strip()
        if not agent_id:
            raise SystemExit(f"get_agent missing id: {agent_info!r}")
        sample_input = {
            "agent": {"id": agent_id, "name": agent_name},
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "tool": {
                "name": "shell_execute_parallel",
                "arguments": {
                    "commands": [{"command": f"printf '{DENY_MARKER}'"}],
                },
            },
        }

        eval_shadow = client.aiden.evaluate_policy(
            {"rego_source": DENY_REGO, "input": sample_input}
        )
        print(f"evaluate (shadow path) outcome={eval_shadow.get('outcome')}")

        client.aiden.detach_policy(agent_name, policy_id)
        shadow_attached = False
        client.aiden.attach_policy(
            agent_name, policy_id=policy_id, shadow=False, priority=10
        )
        attached = True
        print("reattached enforcing")

        eval_live = client.aiden.evaluate_policy(
            {"rego_source": DENY_REGO, "input": sample_input}
        )
        print(f"evaluate (enforcing) outcome={eval_live.get('outcome')}")

        # duration_minutes is 1–480. Keep the sample short; end_break_glass in finally.
        client.aiden.break_glass(
            agent_name,
            reason=f"sdk sample outage bypass {suffix}",
            duration_minutes=15,
        )
        glass_open = True

        state = client.aiden.get_break_glass(agent_name)
        print(
            "break-glass "
            f"active={state.get('active')} "
            f"activated_by={state.get('activated_by')} "
            f"reason={state.get('reason')}"
        )
        if not state.get("active"):
            raise SystemExit(f"expected break-glass active: {state!r}")
        if not state.get("activated_by"):
            raise SystemExit(
                f"expected activated_by from API token on response: {state!r}"
            )
    finally:
        cleanup_errors: list[str] = []

        def cleanup(label: str, fn) -> None:
            try:
                fn()
            except Exception as exc:  # noqa: BLE001
                cleanup_errors.append(f"{label}: {exc}")
                print(f"{label} failed: {exc}", file=sys.stderr)

        if glass_open:
            cleanup("end break-glass", lambda: client.aiden.end_break_glass(agent_name))
            if not cleanup_errors:
                print("ended break-glass")
        if (attached or shadow_attached) and policy_id:
            cleanup(
                "detach policy",
                lambda: client.aiden.detach_policy(agent_name, policy_id),
            )
            if "detach policy" not in " ".join(cleanup_errors):
                print("detached policy")
        if policy_id:
            cleanup("delete policy", lambda: client.aiden.delete_policy(policy_id))
            if not any(item.startswith("delete policy") for item in cleanup_errors):
                print(f"deleted policy id={policy_id}")
        if previous_rule and previous_rule != "any_deny_blocks":
            cleanup(
                "restore combining algorithm",
                lambda: client.aiden.when_policies_disagree(agent_name, previous_rule),
            )
            if not any(item.startswith("restore combining") for item in cleanup_errors):
                print(f"restored conflict rule: {previous_rule}")
        elif previous_rule is None:
            print(
                "combining algorithm left as any_deny_blocks "
                "(agent record had no previous rule)",
                file=sys.stderr,
            )
    if cleanup_errors:
        raise SystemExit("; ".join(cleanup_errors))


if __name__ == "__main__":
    main()