#!/usr/bin/env python3
"""Agent Mechanics Lab 02: bind approval to one exact tool action."""
from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import json
import os
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

SECRET = b"agent-mechanics-lab-secret"
POLICY = {"read_record": "allow", "write_record": "ask", "delete_record": "deny"}
SCHEMAS: dict[str, dict[str, type]] = {
    "read_record": {"record_id": str},
    "write_record": {"record_id": str, "value": str},
    "delete_record": {"record_id": str},
}


@dataclass(frozen=True)
class Decision:
    status: str
    reason: str
    approval_token: str | None = None


def b64encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def b64decode(value: str) -> bytes:
    padding = "=" * (-len(value) % 4)
    return base64.urlsafe_b64decode(value + padding)


def canonical_action(actor: str, tool: str, arguments: dict[str, Any]) -> bytes:
    payload = {"actor": actor, "tool": tool, "arguments": arguments}
    return json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def arguments_hash(actor: str, tool: str, arguments: dict[str, Any]) -> str:
    return hashlib.sha256(canonical_action(actor, tool, arguments)).hexdigest()


def validate_schema(tool: str, arguments: dict[str, Any]) -> str | None:
    schema = SCHEMAS.get(tool)
    if schema is None:
        return "unknown_tool"
    if set(arguments) != set(schema):
        return f"fields_must_equal:{sorted(schema)}"
    for field, expected_type in schema.items():
        if not isinstance(arguments[field], expected_type):
            return f"wrong_type:{field}:{expected_type.__name__}"
    return None


def issue_approval(actor: str, tool: str, arguments: dict[str, Any], ttl_seconds: int = 60) -> str:
    payload = {
        "actor": actor,
        "tool": tool,
        "arguments_hash": arguments_hash(actor, tool, arguments),
        "expires_at": int(time.time()) + ttl_seconds,
        "nonce": b64encode(os.urandom(12)),
    }
    encoded = b64encode(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8"))
    signature = b64encode(hmac.new(SECRET, encoded.encode("ascii"), hashlib.sha256).digest())
    return f"{encoded}.{signature}"


def verify_approval(
    token: str,
    actor: str,
    tool: str,
    arguments: dict[str, Any],
    used_nonces: set[str],
    now: int | None = None,
) -> str | None:
    try:
        encoded, signature = token.split(".", 1)
        expected = b64encode(hmac.new(SECRET, encoded.encode("ascii"), hashlib.sha256).digest())
        if not hmac.compare_digest(signature, expected):
            return "bad_signature"
        payload = json.loads(b64decode(encoded))
    except (ValueError, json.JSONDecodeError):
        return "malformed_token"

    current_time = int(time.time()) if now is None else now
    if payload.get("expires_at", 0) < current_time:
        return "expired"
    if payload.get("actor") != actor or payload.get("tool") != tool:
        return "identity_or_tool_mismatch"
    if payload.get("arguments_hash") != arguments_hash(actor, tool, arguments):
        return "arguments_drift"
    nonce = str(payload.get("nonce", ""))
    if not nonce or nonce in used_nonces:
        return "replay"
    used_nonces.add(nonce)
    return None


def append_audit(path: Path, record: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps({"ts": round(time.time(), 6), **record}, sort_keys=True) + "\n")


def authorize(
    actor: str,
    tool: str,
    arguments: dict[str, Any],
    audit_path: Path,
    used_nonces: set[str],
    approval_token: str | None = None,
) -> Decision:
    schema_error = validate_schema(tool, arguments)
    if schema_error:
        decision = Decision("rejected", schema_error)
        append_audit(audit_path, {"actor": actor, "tool": tool, "arguments": arguments, **decision.__dict__})
        return decision

    policy = POLICY[tool]
    if policy == "deny":
        decision = Decision("denied", "policy_deny")
    elif policy == "allow":
        decision = Decision("executed", "policy_allow")
    elif approval_token is None:
        decision = Decision("approval_required", "policy_ask", issue_approval(actor, tool, arguments))
    else:
        error = verify_approval(approval_token, actor, tool, arguments, used_nonces)
        decision = Decision("executed", "approval_valid") if error is None else Decision("denied", error)

    append_audit(audit_path, {"actor": actor, "tool": tool, "arguments": arguments, **decision.__dict__})
    return decision


def run_demo(audit_path: Path) -> dict[str, str]:
    if audit_path.exists():
        audit_path.unlink()
    used_nonces: set[str] = set()
    actor = "user-42"
    original = {"record_id": "safe-note", "value": "approved content"}

    read_decision = authorize(actor, "read_record", {"record_id": "safe-note"}, audit_path, used_nonces)
    ask_decision = authorize(actor, "write_record", original, audit_path, used_nonces)
    assert ask_decision.approval_token
    drift_decision = authorize(
        actor,
        "write_record",
        {"record_id": "secrets.env", "value": "approved content"},
        audit_path,
        used_nonces,
        ask_decision.approval_token,
    )
    execute_decision = authorize(actor, "write_record", original, audit_path, used_nonces, ask_decision.approval_token)
    replay_decision = authorize(actor, "write_record", original, audit_path, used_nonces, ask_decision.approval_token)
    deny_decision = authorize(actor, "delete_record", {"record_id": "safe-note"}, audit_path, used_nonces)
    schema_decision = authorize(actor, "write_record", {"record_id": "missing-value"}, audit_path, used_nonces)

    return {
        "read": read_decision.status,
        "approval": ask_decision.status,
        "argument_drift": drift_decision.reason,
        "approved_write": execute_decision.status,
        "replay": replay_decision.reason,
        "delete": deny_decision.status,
        "schema": schema_decision.status,
    }


def self_test() -> None:
    with tempfile.TemporaryDirectory(prefix="agent-mechanics-gate-") as temporary:
        result = run_demo(Path(temporary) / "audit.jsonl")
        assert result == {
            "read": "executed",
            "approval": "approval_required",
            "argument_drift": "arguments_drift",
            "approved_write": "executed",
            "replay": "replay",
            "delete": "denied",
            "schema": "rejected",
        }
    print("PASS tool-gate action_bound=true replay_blocked=true audit_complete=true")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--audit", type=Path, default=Path(".agent-mechanics-lab/tool-gate/audit.jsonl"))
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        self_test()
        return 0
    result = run_demo(args.audit)
    print(json.dumps(result, indent=2, sort_keys=True))
    print(f"audit={args.audit}")
    print("PASS action_bound=true replay_blocked=true")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())