"""司命授权凭证：HMAC-SHA256 短期 token，与 Rust 沙箱端 ``auth.rs`` 共用算法与密钥。

凭证格式：``<payload_b64url>.<sig_b64url>``

- payload：JSON，字段顺序固定为
  ``exp / trace_id / risk_level / tool / character_id / run_id / body_hash / rlimit_overrides / file_op_hash / isolation_mode``，
  必须与 Rust ``serde`` 字段顺序字节级一致（否则 HMAC 不匹配）。``isolation_mode`` 末尾追加，旧凭证向后兼容。
- sig：``HMAC-SHA256(secret, payload_b64url)`` 的 base64url（无填充）。
- body_hash：``compute_body_hash(argv, stdin)``，把凭证绑死到具体执行内容，防 TTL 内重放执行任意命令。

安全语义：客户端不判断风险，只把「经司命审核通过」的上下文连同请求内容摘要签成短期凭证；
沙箱端验签 + body_hash 比对通过才执行。
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import time
from dataclasses import dataclass
from typing import cast


@dataclass(frozen=True)
class AuthorizationContext:
    """一次授权的上下文：由司命审核通过后随请求下发。"""

    trace_id: str
    risk_level: str  # L0–L4，见契约字典 risk_level
    tool: str
    character_id: str | None = None
    run_id: str | None = None
    # 司命授予的资源配额上限（key: fsize/as/cpu/nproc/nofile → u64）。
    # 随凭证下发，sandbox 验签后信任并 clamp 到全局天花板；None 表示不绑定分级配额。
    # 属授权属性，不进 body_hash（body_hash 只绑 argv + stdin）。
    rlimit_overrides: dict[str, int] | None = None
    # 隔离模式授权（L0-L2 → process，L3-L4 → container）。随凭证下发，沙箱据此路由；
    # None 表示司命未指定。属授权属性，不进 body_hash，与 rlimit_overrides 同链路。
    isolation_mode: str | None = None


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


def _b64url_no_pad_decode(text: str) -> bytes:
    pad = "=" * (-len(text) % 4)
    return base64.urlsafe_b64decode(text + pad)


def compute_body_hash(argv: list[str], stdin: str | None) -> str:
    """计算请求内容摘要，与 Rust 沙箱端 ``compute_body_hash`` 字节级一致。

    口径：``SHA256(canonical JSON {"argv": [...], "stdin": null|str})``——
    JSON 紧凑（``separators=(",",":")``）、不转义非 ASCII（``ensure_ascii=False``）、
    字段顺序固定 ``argv, stdin``；返回小写十六进制（64 字符）。
    """
    payload = {"argv": argv, "stdin": stdin}
    encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def compute_file_op_hash(op: str, path: str) -> str:
    """计算文件操作摘要，与 Rust 沙箱端 ``compute_file_op_hash`` 字节级一致。

    口径：``SHA256(canonical JSON {"op": str, "path": str})``——把凭证绑死到具体文件操作
    （read/write/list/stat），防 TTL 内重放执行任意文件读写。字段顺序固定 ``op, path``。
    """
    payload = {"op": op, "path": path}
    encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def compute_body_hash_code(language: str, code: str, stdin: str | None) -> str:
    """计算 code 模式请求内容摘要，与 Rust 沙箱端 ``compute_body_hash_code`` 字节级一致。

    口径：``SHA256(canonical JSON {"language": str, "code": str, "stdin": null|str})``——
    字段顺序固定 ``language, code, stdin``、紧凑 JSON、不转义非 ASCII；返回小写十六进制。
    code 模式凭证 body_hash 用此函数（argv 模式仍用 ``compute_body_hash``），把凭证绑死到
    具体代码内容，防 TTL 内重放执行任意代码。
    """
    payload = {"language": language, "code": code, "stdin": stdin}
    encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def sign_authorization(
    secret: str,
    ctx: AuthorizationContext,
    argv: list[str] | None = None,
    stdin: str | None = None,
    ttl_seconds: int = 30,
    *,
    code: str | None = None,
    language: str | None = None,
    file_op: tuple[str, str] | None = None,
) -> str:
    """签发授权凭证。

    内容绑定（互斥，由调用方保证不同时传 argv 与 code）：
    - exec 命令模式：传 ``argv``/``stdin``，按 ``compute_body_hash(argv, stdin)`` 绑定。
    - exec 代码模式：传 ``code``/``language``/``stdin``，按 ``compute_body_hash_code`` 绑定。
    - 低风险操作（cancel）：都不传，不绑定 body_hash（file_op 单独走 file_op_hash）。

    payload 字段顺序与 Rust ``AuthPayload`` 一致（isolation_mode 末尾追加，旧凭证向后兼容）。
    """
    if not secret:
        raise ValueError("授权密钥不能为空")
    exp = int(time.time()) + ttl_seconds
    if code is not None:
        body_hash = compute_body_hash_code(language or "python", code, stdin)
    elif argv is not None:
        body_hash = compute_body_hash(argv, stdin)
    else:
        body_hash = None
    payload: dict[str, object] = {
        "exp": exp,
        "trace_id": ctx.trace_id,
        "risk_level": ctx.risk_level,
        "tool": ctx.tool,
        "character_id": ctx.character_id,
        "run_id": ctx.run_id,
        "body_hash": body_hash,
        # rlimit_overrides 内部按 key 排序，对齐 Rust serde BTreeMap 的字母序，
        # 保证多 key 时两端 payload JSON 字节一致 → HMAC 一致。
        "rlimit_overrides": (
            dict(sorted(ctx.rlimit_overrides.items())) if ctx.rlimit_overrides else None
        ),
        # file_op_hash 绑定文件操作 {op, path}（exec 路径为 None）；字段顺序对齐 Rust AuthPayload。
        "file_op_hash": compute_file_op_hash(file_op[0], file_op[1]) if file_op else None,
        # isolation_mode 末尾追加（字段顺序对齐 Rust AuthPayload；旧凭证无此字段仍可验签）。
        "isolation_mode": ctx.isolation_mode,
    }
    payload_json = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
    payload_b64 = _b64url_no_pad(payload_json.encode("utf-8"))
    digest = hmac.new(secret.encode("utf-8"), payload_b64.encode("ascii"), hashlib.sha256).digest()
    return f"{payload_b64}.{_b64url_no_pad(digest)}"


def verify_authorization(secret: str, token: str) -> dict[str, object]:
    """校验授权凭证（对称实现，供本地测试与跨端一致性校验）。"""
    payload_b64, sep, sig_b64 = token.rpartition(".")
    if not sep:
        raise ValueError("授权凭证格式不合法")
    expected = _b64url_no_pad(
        hmac.new(secret.encode("utf-8"), payload_b64.encode("ascii"), hashlib.sha256).digest()
    )
    if not hmac.compare_digest(expected, sig_b64):
        raise ValueError("授权凭证签名校验失败")
    payload = cast("dict[str, object]", json.loads(_b64url_no_pad_decode(payload_b64)))
    exp = payload.get("exp")
    if not isinstance(exp, int | float) or exp <= time.time():
        raise ValueError("授权凭证已过期")
    return payload