//! 司命授权：HMAC-SHA256 签名的短期凭证。
//!
//! 凭证格式：`<payload_b64url>.<sig_b64url>`
//! - payload: JSON `{exp, trace_id, risk_level, tool, character_id, run_id, body_hash, rlimit_overrides}`，
//!   字段顺序与后端 Python 客户端一致。
//! - sig: HMAC-SHA256(secret, payload_b64url) 的 base64url（无填充）。
//!
//! 安全语义：沙箱不做风险判定，只校验"调用方是否持有由司命签发的有效凭证"。
//! `body_hash` 把凭证绑死到具体执行内容（argv + stdin 摘要），防 TTL 内重放执行任意命令；
//! `rlimit_overrides` 携带司命授予的资源配额上限（授权属性，不进 body_hash），验签后信任并 clamp。
//! verify 支持多 secret 并存（主 + previous），支持无中断轮换。
//! 验签失败、过期或 body_hash 不匹配一律 401。
//! 设计依据见 docs/architecture/service-boundaries.md 与 cross-language-integration-design.md。

use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use axum::extract::{Request, State};
use axum::http::{header, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use axum::Json;
use base64::prelude::{BASE64_URL_SAFE_NO_PAD, Engine as _};
use hmac::{Hmac, Mac};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::AppState;

type HmacSha256 = Hmac<Sha256>;

/// 授权凭证 payload（验签通过后插入请求扩展，供 handler 关联 trace、body_hash 校验与资源配额）。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthPayload {
    /// 过期时间（Unix 秒）。
    pub exp: u64,
    #[serde(default)]
    pub trace_id: Option<String>,
    #[serde(default)]
    pub risk_level: Option<String>,
    #[serde(default)]
    pub tool: Option<String>,
    #[serde(default)]
    pub character_id: Option<String>,
    #[serde(default)]
    pub run_id: Option<String>,
    /// 请求内容摘要 SHA256(canonical JSON `{argv, stdin}`)，绑定凭证与具体执行内容，
    /// 防止凭证在 TTL 内被重放去执行任意命令；None 表示不绑定（仅 auth_enabled=false 本地联调）。
    #[serde(default)]
    pub body_hash: Option<String>,
    /// 司命授予该上下文的资源配额上限（key 见 `RlimitKind::as_str`：fsize/as/cpu/nproc/nofile）。
    /// 随 HMAC 凭证下发，sandbox 验签后信任并 clamp 到全局天花板；属授权属性，不进 body_hash。
    #[serde(default)]
    pub rlimit_overrides: Option<BTreeMap<String, u64>>,
    /// 文件操作摘要 SHA256(canonical JSON `{op, path}`)，绑定凭证与具体文件操作
    /// （read/write/list/stat），防 TTL 内重放执行任意文件操作。exec 凭证为 None、file 凭证
    /// 为 Some；与 body_hash 互斥使用，不参与 body_hash 计算。
    #[serde(default)]
    pub file_op_hash: Option<String>,
    /// 隔离模式授权（L0-L2 → process，L3-L4 → container）。随 HMAC 凭证下发，sandbox 验签后
    /// 据此路由执行路径；P0 期间 container 降级 process。属授权属性，不进 body_hash，与
    /// rlimit_overrides 同链路。None 表示司命未指定，按 SANDBOX_ISOLATION_DEFAULT 兜底。
    #[serde(default)]
    pub isolation_mode: Option<String>,
}

#[derive(Debug)]
pub enum AuthError {
    Missing,
    Malformed,
    BadSignature,
    Expired,
}

impl std::fmt::Display for AuthError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let s = match self {
            AuthError::Missing => "缺少 Authorization 头",
            AuthError::Malformed => "授权凭证格式不合法",
            AuthError::BadSignature => "授权凭证签名校验失败",
            AuthError::Expired => "授权凭证已过期",
        };
        f.write_str(s)
    }
}

/// 签发授权凭证（与后端 Python `sign_authorization` 共用同一算法与密钥）。
#[allow(dead_code)] // 沙箱端仅校验（verify）；sign 供对称测试与未来本地签发场景保留
pub fn sign(secret: &str, payload: &AuthPayload) -> anyhow::Result<String> {
    let json = serde_json::to_vec(payload)?;
    let payload_b64 = BASE64_URL_SAFE_NO_PAD.encode(&json);
    let mut mac =
        HmacSha256::new_from_slice(secret.as_bytes()).expect("HMAC 接受任意长度密钥");
    mac.update(payload_b64.as_bytes());
    let sig_b64 = BASE64_URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
    Ok(format!("{payload_b64}.{sig_b64}"))
}

/// 校验授权凭证（多 secret）：逐个尝试 secrets，任一验签通过即放行。
///
/// 返回 `(payload, 命中 secret 的索引)`：0 为主 secret（`SANDBOX_AUTH_SECRET`），
/// \>0 为旧 secret（`SANDBOX_AUTH_SECRET_PREVIOUS` 第几个）。索引供 metric 区分轮换进度。
/// 格式 / 过期错误与 secret 无关，只判一次；仅验签阶段遍历多 secret（支持无中断轮换的并存窗口）。
pub fn verify_any(secrets: &[String], token: &str) -> Result<(AuthPayload, usize), AuthError> {
    let (payload_b64, sig_b64) = token.rsplit_once('.').ok_or(AuthError::Malformed)?;
    let sig_bytes = BASE64_URL_SAFE_NO_PAD
        .decode(sig_b64.as_bytes())
        .map_err(|_| AuthError::Malformed)?;

    // 验签：任一 secret 通过即可（常量时间比较由 hmac verify_slice 提供）。
    let mut slot = None;
    for (i, secret) in secrets.iter().enumerate() {
        let mut mac = HmacSha256::new_from_slice(secret.as_bytes())
            .map_err(|_| AuthError::Malformed)?;
        mac.update(payload_b64.as_bytes());
        if mac.verify_slice(&sig_bytes).is_ok() {
            slot = Some(i);
            break;
        }
    }
    let slot = slot.ok_or(AuthError::BadSignature)?;

    let payload_bytes = BASE64_URL_SAFE_NO_PAD
        .decode(payload_b64.as_bytes())
        .map_err(|_| AuthError::Malformed)?;
    let payload: AuthPayload =
        serde_json::from_slice(&payload_bytes).map_err(|_| AuthError::Malformed)?;

    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);
    if payload.exp <= now {
        return Err(AuthError::Expired);
    }
    Ok((payload, slot))
}

/// 单 secret 校验（测试便捷包装：内部转 `verify_any`）。
/// 仅测试用——生产路径 `require_authorization` 直连 `verify_any` 以支持多 secret 轮换。
#[cfg(test)]
pub fn verify(secret: &str, token: &str) -> Result<AuthPayload, AuthError> {
    verify_any(&[secret.to_string()], token).map(|(p, _)| p)
}

/// 计算请求内容摘要：SHA256(canonical JSON of `{argv, stdin}`)。
///
/// 与后端 Python `compute_body_hash` 字节级一致：
/// - JSON 紧凑（无空格），字段顺序固定 `argv, stdin`；`serde_json` 默认紧凑输出 + Option None → null，
///   对齐 Python `json.dumps(..., separators=(",",":"), ensure_ascii=False)`
/// - 返回小写十六进制（64 字符），对齐 `hashlib.sha256().hexdigest()`
/// - 注意：`rlimit_overrides` 不参与摘要（属授权属性，非执行内容）
pub fn compute_body_hash(argv: &[String], stdin: &Option<String>) -> String {
    #[derive(Serialize)]
    struct BodyHashInput<'a> {
        argv: &'a [String],
        stdin: &'a Option<String>,
    }
    let bytes = serde_json::to_vec(&BodyHashInput { argv, stdin })
        .expect("argv/stdin 序列化不会失败");
    let digest = Sha256::digest(&bytes);
    digest.iter().map(|b| format!("{:02x}", b)).collect()
}

/// 计算 code 模式请求内容摘要：SHA256(canonical JSON `{language, code, stdin}`)。
///
/// 与后端 Python `compute_body_hash_code` 字节级一致：
/// - JSON 紧凑、字段顺序固定 `language, code, stdin`；Option None → null
/// - 返回小写十六进制（64 字符）
/// - code 模式凭证 body_hash 用此函数（argv 模式仍用 compute_body_hash），把凭证绑死到
///   具体代码内容，防 TTL 内重放执行任意代码。
pub fn compute_body_hash_code(language: &str, code: &str, stdin: &Option<String>) -> String {
    #[derive(Serialize)]
    struct CodeHashInput<'a> {
        language: &'a str,
        code: &'a str,
        stdin: &'a Option<String>,
    }
    let bytes = serde_json::to_vec(&CodeHashInput {
        language,
        code,
        stdin,
    })
    .expect("language/code/stdin 序列化不会失败");
    let digest = Sha256::digest(&bytes);
    digest.iter().map(|b| format!("{:02x}", b)).collect()
}

/// 计算文件操作摘要：SHA256(canonical JSON `{op, path}`)。
///
/// 与后端 Python `compute_file_op_hash` 字节级一致：紧凑 JSON、字段顺序固定 `op, path`、
/// 小写十六进制。把凭证绑死到具体文件操作，防 TTL 内重放执行任意文件读写。
pub fn compute_file_op_hash(op: &str, path: &str) -> String {
    #[derive(Serialize)]
    struct FileOpHashInput<'a> {
        op: &'a str,
        path: &'a str,
    }
    let bytes = serde_json::to_vec(&FileOpHashInput { op, path })
        .expect("op/path 序列化不会失败");
    let digest = Sha256::digest(&bytes);
    digest.iter().map(|b| format!("{:02x}", b)).collect()
}

/// axum 中间件：校验 `Authorization: Bearer <token>`；`auth_enabled=false` 时直接放行。
pub async fn require_authorization(
    State(state): State<Arc<AppState>>,
    mut req: Request,
    next: Next,
) -> Response {
    if !state.config.auth_enabled {
        // 本地联调：跳过验签，但仍插入占位 payload，使 handler 能统一从扩展取 AuthPayload。
        req.extensions_mut().insert(AuthPayload {
            exp: 0,
            trace_id: None,
            risk_level: None,
            tool: None,
            character_id: None,
            run_id: None,
            body_hash: None,
            rlimit_overrides: None,
            file_op_hash: None,
            isolation_mode: None,
        });
        return next.run(req).await;
    }

    let token = req
        .headers()
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.strip_prefix("Bearer ").map(str::trim));

    match token {
        Some(t) => {
            // 主 secret + previous secrets 一并尝试（无中断轮换的并存窗口）。
            let mut secrets = vec![state.config.auth_secret.clone()];
            secrets.extend(state.config.auth_secret_previous.iter().cloned());
            match verify_any(&secrets, t) {
                Ok((payload, slot)) => {
                    state.executor.metrics().record_auth_secret(slot);
                    req.extensions_mut().insert(payload);
                    next.run(req).await
                }
                Err(e) => unauthorized(e),
            }
        }
        None => unauthorized(AuthError::Missing),
    }
}

fn unauthorized(err: AuthError) -> Response {
    let body = Json(serde_json::json!({
        "error": { "code": "sandbox_unauthorized", "message": err.to_string() }
    }));
    (StatusCode::UNAUTHORIZED, body).into_response()
}

#[cfg(test)]
mod tests {
    use super::*;

    const SECRET: &str = "test-shared-secret";

    fn payload(exp_offset_secs: i64) -> AuthPayload {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0);
        AuthPayload {
            exp: (now as i64 + exp_offset_secs) as u64,
            trace_id: Some("trace_abc".into()),
            risk_level: Some("L1".into()),
            tool: Some("shell".into()),
            character_id: Some("character_01H".into()),
            run_id: Some("run_01H".into()),
            body_hash: Some("deadbeef".into()),
            rlimit_overrides: None,
            file_op_hash: None,
            isolation_mode: None,
        }
    }

    #[test]
    fn sign_then_verify_roundtrip() {
        let token = sign(SECRET, &payload(60)).unwrap();
        let claims = verify(SECRET, &token).expect("合法凭证应校验通过");
        assert_eq!(claims.trace_id.as_deref(), Some("trace_abc"));
        assert_eq!(claims.risk_level.as_deref(), Some("L1"));
        assert_eq!(claims.body_hash.as_deref(), Some("deadbeef"));
    }

    #[test]
    fn rejects_wrong_secret() {
        let token = sign(SECRET, &payload(60)).unwrap();
        assert!(matches!(
            verify("wrong-secret", &token),
            Err(AuthError::BadSignature)
        ));
    }

    #[test]
    fn rejects_tampered_token() {
        let mut token = sign(SECRET, &payload(60)).unwrap();
        // 篡改签名末位
        let last = token.len() - 1;
        token.replace_range(last.., if &token[last..] == "A" { "B" } else { "A" });
        assert!(verify(SECRET, &token).is_err());
    }

    #[test]
    fn rejects_expired() {
        let token = sign(SECRET, &payload(-10)).unwrap();
        assert!(matches!(verify(SECRET, &token), Err(AuthError::Expired)));
    }

    #[test]
    fn rejects_malformed() {
        assert!(matches!(
            verify(SECRET, "not-a-valid-token"),
            Err(AuthError::Malformed)
        ));
    }

    #[test]
    fn body_hash_known_vector_matches_python() {
        // 与后端 Python compute_body_hash 共享同一固定向量，保证跨端字节级一致。
        // 向量：{"argv":["sh","-c","echo hi"],"stdin":null}
        let argv = vec!["sh".to_string(), "-c".to_string(), "echo hi".to_string()];
        let hash = compute_body_hash(&argv, &None);
        assert_eq!(
            hash,
            "f5fa97ea2ee92c505174672da52c59c952e49c32fa56a1db5bd983c10a829540"
        );
    }

    #[test]
    fn body_hash_distinguishes_none_and_empty_stdin() {
        let argv = vec!["sh".to_string()];
        // None → "stdin":null；Some("") → "stdin":""，两者字节不同
        assert_ne!(
            compute_body_hash(&argv, &None),
            compute_body_hash(&argv, &Some(String::new()))
        );
    }

    #[test]
    fn file_op_hash_known_vector_matches_python() {
        // 与后端 Python compute_file_op_hash 共享同一固定向量，保证跨端字节级一致。
        // 向量：{"op":"read","path":"data/note.txt"}
        let hash = compute_file_op_hash("read", "data/note.txt");
        assert_eq!(
            hash,
            "717f67252c61cf1681cb3c6018d3801b79a511564664567b57d12d775bb8fa53"
        );
    }

    #[test]
    fn file_op_hash_differs_per_op_and_path() {
        // 凭证绑死具体文件操作：op 或 path 不同则摘要不同，杜绝重放任意文件读写。
        let read_a = compute_file_op_hash("read", "a.txt");
        let write_a = compute_file_op_hash("write", "a.txt");
        let read_b = compute_file_op_hash("read", "b.txt");
        assert_ne!(read_a, write_a); // 同 path 不同 op
        assert_ne!(read_a, read_b); // 同 op 不同 path
    }

    #[test]
    fn verify_any_accepts_previous_secret() {
        // 轮换并存窗口：主 secret + 旧 secret 都应被接受，且命中 slot 不同。
        let token_old = sign("old-secret", &payload(60)).unwrap();
        let token_new = sign("new-secret", &payload(60)).unwrap();
        let secrets = vec!["new-secret".to_string(), "old-secret".to_string()];
        // 新 secret 签发 → 命中 slot 0（primary）
        let (p, slot) = verify_any(&secrets, &token_new).unwrap();
        assert_eq!(slot, 0);
        assert_eq!(p.trace_id.as_deref(), Some("trace_abc"));
        // 旧 secret 签发 → 命中 slot 1（previous）
        let (_, slot) = verify_any(&secrets, &token_old).unwrap();
        assert_eq!(slot, 1);
    }

    #[test]
    fn verify_any_rejects_when_no_secret_matches() {
        let token = sign("actual-secret", &payload(60)).unwrap();
        let secrets = vec!["wrong1".to_string(), "wrong2".to_string()];
        assert!(matches!(
            verify_any(&secrets, &token),
            Err(AuthError::BadSignature)
        ));
    }

    #[test]
    fn rlimit_overrides_roundtrip_through_payload() {
        // rlimit_overrides 随凭证序列化/反序列化，且不进 body_hash
        let mut p = payload(60);
        let mut overrides = BTreeMap::new();
        overrides.insert("fsize".to_string(), 200_000_000u64);
        p.rlimit_overrides = Some(overrides);
        let token = sign(SECRET, &p).unwrap();
        let claims = verify(SECRET, &token).unwrap();
        assert_eq!(
            claims.rlimit_overrides.as_ref().unwrap().get("fsize"),
            Some(&200_000_000u64)
        );
        // body_hash 不受 rlimit_overrides 影响（测试用占位值，未绑定真实 argv/stdin）
        assert_eq!(claims.body_hash.as_deref(), Some("deadbeef"));
    }

    #[test]
    fn body_hash_code_known_vector_matches_python() {
        // 与后端 Python compute_body_hash_code 共享固定向量，保证跨端字节级一致。
        // 向量：{"language":"python","code":"print(1+1)","stdin":null}
        let hash = compute_body_hash_code("python", "print(1+1)", &None);
        assert_eq!(
            hash,
            "e7d633605676c3b19c642ee2437d91f9fb22562aaffad40f217d84c587322b69"
        );
    }

    #[test]
    fn body_hash_code_distinguishes_language_and_stdin() {
        let base = compute_body_hash_code("python", "print(1+1)", &None);
        // 不同 language → 不同摘要
        assert_ne!(base, compute_body_hash_code("node", "print(1+1)", &None));
        // 不同 stdin（None vs Some）→ 不同摘要
        assert_ne!(
            base,
            compute_body_hash_code("python", "print(1+1)", &Some("42".into()))
        );
    }

    #[test]
    fn isolation_mode_roundtrip_and_backward_compat() {
        // 新字段随凭证序列化/反序列化
        let mut p = payload(60);
        p.isolation_mode = Some("container".into());
        let token = sign(SECRET, &p).unwrap();
        let claims = verify(SECRET, &token).unwrap();
        assert_eq!(claims.isolation_mode.as_deref(), Some("container"));

        // 向后兼容：手工构造不含 isolation_mode 的旧 payload JSON，仍可验签且字段为 None。
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0);
        let legacy_json = format!(
            "{{\"exp\":{},\"trace_id\":\"trace_abc\",\"risk_level\":\"L1\",\"tool\":\"shell\",\"character_id\":\"character_01H\",\"run_id\":\"run_01H\",\"body_hash\":\"deadbeef\",\"rlimit_overrides\":null,\"file_op_hash\":null}}",
            now as i64 + 60
        );
        let payload_b64 = BASE64_URL_SAFE_NO_PAD.encode(legacy_json.as_bytes());
        let mut mac = HmacSha256::new_from_slice(SECRET.as_bytes()).unwrap();
        mac.update(payload_b64.as_bytes());
        let sig_b64 = BASE64_URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
        let token = format!("{payload_b64}.{sig_b64}");
        let claims = verify(SECRET, &token).expect("旧凭证（无 isolation_mode）应可验签");
        assert!(claims.isolation_mode.is_none());
    }
}