"""A10 minimal authorization lab; Python 3.10+, stdlib only.

Run:
    python3 A10-Agent-Permissions-and-Constraints-v1.1-demo.py

This is a deterministic policy lab, not a production sandbox.
"""

from dataclasses import dataclass, field
from enum import Enum
from hashlib import sha256
from typing import Any
import json


class Verdict(str, Enum):
    ALLOW = "ALLOW"
    DENY = "DENY"
    APPROVAL = "APPROVAL_REQUIRED"


@dataclass(frozen=True)
class Principal:
    agent_id: str
    tenant: str
    roles: frozenset[str]


@dataclass(frozen=True)
class Request:
    principal: Principal
    action: str
    resource: str
    args: dict[str, Any]
    run_id: str
    call_id: str


@dataclass(frozen=True)
class Decision:
    verdict: Verdict
    reason: str
    policy_version: str


@dataclass
class Ledger:
    records: list[dict[str, Any]] = field(default_factory=list)
    committed: dict[str, dict[str, Any]] = field(default_factory=dict)
    used: dict[str, int] = field(default_factory=dict)

    def log(self, req: Request, dec: Decision):
        self.records.append({
            "call_id": req.call_id,
            "agent": req.principal.agent_id,
            "action": req.action,
            "resource": req.resource,
            "decision": dec.verdict.value,
            "reason": dec.reason,
            "policy_version": dec.policy_version,
        })


class Policy:
    version = "a10-lab-v1"

    def evaluate(self, req: Request, ledger: Ledger) -> Decision:

        def result(verdict: Verdict, reason: str):
            return Decision(verdict, reason, self.version)

        # Explicit scope: no cross-tenant path,
        # and no unregistered action.

        if not req.resource.startswith(
            f"tenant/{req.principal.tenant}/"
        ):
            return result(Verdict.DENY, "TENANT_SCOPE")

        if req.action not in {"read", "draft", "publish"}:
            return result(Verdict.DENY, "UNKNOWN_ACTION")

        if (
            req.action == "read"
            and "reader" not in req.principal.roles
        ):
            return result(Verdict.DENY, "ROLE_REQUIRED")

        if (
            req.action in {"draft", "publish"}
            and "author" not in req.principal.roles
        ):
            return result(Verdict.DENY, "ROLE_REQUIRED")

        if ledger.used.get(req.run_id, 0) >= 3:
            return result(Verdict.DENY, "RUN_BUDGET")

        if req.action == "publish":
            return result(Verdict.APPROVAL, "HUMAN_APPROVAL")

        return result(Verdict.ALLOW, "POLICY_MATCH")


class Gateway:

    def __init__(self):
        self.ledger = Ledger()
        self.policy = Policy()
        self.storage: dict[str, str] = {}

    def invoke(
        self,
        req: Request,
        approval_token: str | None = None,
    ) -> dict[str, Any]:

        # Repeated call_id must be identical.
        # Production needs atomic durable storage.

        digest = sha256(
            json.dumps({
                "agent": req.principal.agent_id,
                "tenant": req.principal.tenant,
                "roles": sorted(req.principal.roles),
                "action": req.action,
                "resource": req.resource,
                "args": req.args,
                "run_id": req.run_id,
            }, sort_keys=True).encode()
        ).hexdigest()

        if req.call_id in self.ledger.committed:

            previous = self.ledger.committed[req.call_id]

            if previous["digest"] != digest:
                return {
                    "ok": False,
                    "error": "CALL_ID_CONFLICT",
                }

            return previous["result"]

        dec = self.policy.evaluate(req, self.ledger)

        self.ledger.log(req, dec)

        if dec.verdict == Verdict.DENY:
            return {
                "ok": False,
                "error": dec.reason,
            }

        if (
            dec.verdict == Verdict.APPROVAL
            and approval_token != f"approved:{req.call_id}"
        ):
            return {
                "ok": False,
                "error": "APPROVAL_REQUIRED",
            }

        # All side effects are executed
        # after policy and approval checks.

        if req.action == "read":

            result = {
                "ok": True,
                "value": self.storage.get(req.resource),
            }

        else:

            content = req.args.get("content")

            if (
                not isinstance(content, str)
                or len(content) > 1000
            ):
                return {
                    "ok": False,
                    "error": "INVALID_CONTENT",
                }

            self.storage[req.resource] = content

            result = {
                "ok": True,
                "written": req.resource,
                "action": req.action,
            }

        self.ledger.used[req.run_id] = (
            self.ledger.used.get(req.run_id, 0) + 1
        )

        self.ledger.committed[req.call_id] = {
            "digest": digest,
            "result": result,
        }

        return result


def run_checks():

    g = Gateway()

    p = Principal(
        "research-agent",
        "acme",
        frozenset({"reader", "author"}),
    )

    def req(
        action,
        resource,
        call_id,
        run="run-1",
        args=None,
        principal=p,
    ):
        return Request(
            principal,
            action,
            resource,
            args or {},
            run,
            call_id,
        )

    # 1. Authorized write

    good = req(
        "draft",
        "tenant/acme/report.md",
        "c1",
        args={"content": "draft"},
    )

    assert g.invoke(good)["ok"]

    assert (
        g.storage["tenant/acme/report.md"]
        == "draft"
    )

    # 2. Idempotency

    assert (
        g.invoke(good)["ok"]
        and g.ledger.used["run-1"] == 1
    )

    # 3. Call ID conflict

    assert g.invoke(
        req(
            "draft",
            "tenant/acme/report.md",
            "c1",
            args={"content": "tamper"},
        )
    )["error"] == "CALL_ID_CONFLICT"

    # 4. Cross-tenant access denied

    assert g.invoke(
        req(
            "read",
            "tenant/other/secret",
            "c2",
        )
    )["error"] == "TENANT_SCOPE"

    # 5. Unknown action denied

    assert g.invoke(
        req(
            "delete",
            "tenant/acme/report.md",
            "c3",
        )
    )["error"] == "UNKNOWN_ACTION"

    # 6. Human approval required

    pub = req(
        "publish",
        "tenant/acme/report.md",
        "c4",
        args={"content": "public"},
    )

    assert (
        g.invoke(pub)["error"]
        == "APPROVAL_REQUIRED"
    )

    assert (
        g.storage["tenant/acme/report.md"]
        == "draft"
    )

    # 7. Approved execution

    assert g.invoke(
        pub,
        approval_token="approved:c4",
    )["ok"]

    assert (
        g.storage["tenant/acme/report.md"]
        == "public"
    )

    # 8. Role restriction

    reader = Principal(
        "reader-agent",
        "acme",
        frozenset({"reader"}),
    )

    assert g.invoke(
        req(
            "draft",
            "tenant/acme/other.md",
            "c5",
            principal=reader,
        )
    )["error"] == "ROLE_REQUIRED"

    # 9. Budget enforcement

    assert g.invoke(
        req(
            "read",
            "tenant/acme/report.md",
            "c6",
        )
    )["ok"]

    assert g.invoke(
        req(
            "read",
            "tenant/acme/report.md",
            "c7",
        )
    )["error"] == "RUN_BUDGET"

    # 10. Input validation

    assert g.invoke(
        req(
            "draft",
            "tenant/acme/huge",
            "c8",
            run="run-2",
            args={"content": "x" * 1001},
        )
    )["error"] == "INVALID_CONTENT"

    print("A10 policy lab: 10 checks passed")

    print(
        "audit_records=",
        len(g.ledger.records),
        "committed=",
        len(g.ledger.committed),
    )


if __name__ == "__main__":
    run_checks()