//! The hub as an OpenID Connect issuer for jobs (hub spec §13.5,
//! oxygenhub#103): a job that declares `identity "AUDIENCE" { mount-path … }`
//! gets a signed JWT naming the repository, bookmark, run and job, so a
//! cloud or a vault that trusts this hub can let the job in without a
//! long-lived secret. The key is RSA 2048 in `keys/oidc.pem`, made at
//! first use; `/.well-known/openid-configuration` and
//! `/.well-known/jwks.json` publish what verifiers need.

use rsa::pkcs1v15::{SigningKey, VerifyingKey};
use rsa::pkcs8::{DecodePrivateKey, EncodePrivateKey};
use rsa::signature::{SignatureEncoding, Signer, Verifier};
use rsa::traits::PublicKeyParts;
use rsa::RsaPrivateKey;
use sha2::Sha256;
use std::path::Path;

use crate::api::ApiError;
use crate::Shared;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::Json;

/// The longest a job token lives, whatever the job's limit.
pub const MAX_TTL_S: u64 = 6 * 3600;

pub struct Issuer {
    signing: SigningKey<Sha256>,
    verifying: VerifyingKey<Sha256>,
    /// The key id: the first 16 hex of BLAKE3 over the public key's
    /// modulus, as the JWKS names it.
    pub kid: String,
    n: String,
    e: String,
}

/// What a token says about its job.
pub struct JobClaims<'a> {
    pub issuer: &'a str,
    pub audience: &'a str,
    pub owner: &'a str,
    pub name: &'a str,
    pub visibility: &'a str,
    pub bookmark: &'a str,
    pub run_number: i64,
    pub job: &'a str,
    pub state: &'a str,
    pub runner: &'a str,
    pub actor: &'a str,
    /// The environment the job deploys to, only once it was approved
    /// (oxygenflow#59).
    pub environment: Option<&'a str>,
    /// Who approved the job, and the sealed approval's hash (#235).
    pub approved_by: Option<&'a str>,
    pub approval: Option<&'a str>,
    pub now_ms: i64,
    pub ttl_s: u64,
}

impl Issuer {
    /// The key from `<keys>/oidc.pem`, or a new 2048-bit one written
    /// there (mode 0600).
    pub fn load_or_create(keys: &Path) -> anyhow::Result<Issuer> {
        std::fs::create_dir_all(keys)?;
        let path = keys.join("oidc.pem");
        let key = match std::fs::read_to_string(&path) {
            Ok(pem) => RsaPrivateKey::from_pkcs8_pem(&pem).map_err(|e| anyhow::anyhow!("{}: {e}", path.display()))?,
            Err(_) => {
                let key = RsaPrivateKey::new(&mut rand::rngs::OsRng, 2048)?;
                let pem = key.to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)?;
                std::fs::write(&path, pem.as_bytes())?;
                #[cfg(unix)]
                {
                    use std::os::unix::fs::PermissionsExt;
                    std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
                }
                key
            }
        };
        Ok(Issuer::from_key(key))
    }

    fn from_key(key: RsaPrivateKey) -> Issuer {
        let public = key.to_public_key();
        let n_bytes = public.n().to_bytes_be();
        let kid = blake3::hash(&n_bytes).to_hex()[..16].to_string();
        Issuer {
            n: b64url(&n_bytes),
            e: b64url(&public.e().to_bytes_be()),
            verifying: VerifyingKey::new(public),
            signing: SigningKey::new(key),
            kid,
        }
    }

    /// `/.well-known/openid-configuration`.
    pub fn discovery(&self, issuer: &str) -> serde_json::Value {
        serde_json::json!({
            "issuer": issuer,
            "jwks_uri": format!("{issuer}/.well-known/jwks.json"),
            "response_types_supported": ["id_token"],
            "subject_types_supported": ["public"],
            "id_token_signing_alg_values_supported": ["RS256"],
            "claims_supported": ["iss", "sub", "aud", "exp", "iat", "nbf", "jti", "repository", "repository_owner",
                "repository_visibility", "bookmark", "run_number", "job", "state", "runner", "actor", "environment", "approved_by", "approval"],
            "scopes_supported": ["openid"],
        })
    }

    /// `/.well-known/jwks.json`.
    pub fn jwks(&self) -> serde_json::Value {
        serde_json::json!({ "keys": [{ "kty": "RSA", "use": "sig", "alg": "RS256", "kid": self.kid, "n": self.n, "e": self.e }] })
    }

    /// A signed token for a job: `sub` is `repo:OWNER/NAME:bookmark:B`,
    /// so a trust policy can match a repository and a bookmark.
    pub fn mint(&self, c: &JobClaims<'_>) -> String {
        let iat = c.now_ms / 1000;
        let exp = iat + c.ttl_s.clamp(60, MAX_TTL_S) as i64;
        let header = serde_json::json!({ "alg": "RS256", "typ": "JWT", "kid": self.kid });
        let mut payload = serde_json::json!({
            "iss": c.issuer,
            "sub": format!("repo:{}/{}:bookmark:{}", c.owner, c.name, c.bookmark),
            "aud": c.audience,
            "iat": iat,
            "nbf": iat - 30,
            "exp": exp,
            "jti": crate::auth::random_hex(16),
            "repository": format!("{}/{}", c.owner, c.name),
            "repository_owner": c.owner,
            "repository_visibility": c.visibility,
            "bookmark": c.bookmark,
            "run_number": c.run_number,
            "job": c.job,
            "state": c.state,
            "runner": c.runner,
            "actor": c.actor,
        });
        if let Some(e) = c.environment {
            payload["environment"] = serde_json::Value::from(e);
        }
        if let Some(a) = c.approved_by {
            payload["approved_by"] = serde_json::Value::from(a);
        }
        if let Some(a) = c.approval {
            payload["approval"] = serde_json::Value::from(a);
        }
        let signing_input = format!("{}.{}", b64url(header.to_string().as_bytes()), b64url(payload.to_string().as_bytes()));
        let sig = self.signing.sign(signing_input.as_bytes());
        format!("{signing_input}.{}", b64url(&sig.to_bytes()))
    }

    /// A signed token for the hub itself acting as `subject` (a backup's
    /// `backup:OWNER`, oxygenhub#169), for `audience`, living `ttl_s`.
    pub fn mint_for(&self, issuer: &str, subject: &str, audience: &str, now_ms: i64, ttl_s: u64) -> String {
        let iat = now_ms / 1000;
        let header = serde_json::json!({ "alg": "RS256", "typ": "JWT", "kid": self.kid });
        let payload = serde_json::json!({
            "iss": issuer,
            "sub": subject,
            "aud": audience,
            "iat": iat,
            "nbf": iat - 30,
            "exp": iat + ttl_s.clamp(60, MAX_TTL_S) as i64,
            "jti": crate::auth::random_hex(16),
        });
        let signing_input = format!("{}.{}", b64url(header.to_string().as_bytes()), b64url(payload.to_string().as_bytes()));
        let sig = self.signing.sign(signing_input.as_bytes());
        format!("{signing_input}.{}", b64url(&sig.to_bytes()))
    }

    /// The claims of a token this hub signed, when its signature holds,
    /// it is for `audience` and it has not expired at `now_ms`.
    pub fn verify(&self, token: &str, audience: &str, now_ms: i64) -> Option<serde_json::Value> {
        let mut parts = token.split('.');
        let (h, p, s) = (parts.next()?, parts.next()?, parts.next()?);
        if parts.next().is_some() {
            return None;
        }
        let header: serde_json::Value = serde_json::from_slice(&b64url_decode(h)?).ok()?;
        if header["alg"] != "RS256" || header["kid"] != self.kid.as_str() {
            return None;
        }
        let sig = rsa::pkcs1v15::Signature::try_from(b64url_decode(s)?.as_slice()).ok()?;
        self.verifying.verify(format!("{h}.{p}").as_bytes(), &sig).ok()?;
        let claims: serde_json::Value = serde_json::from_slice(&b64url_decode(p)?).ok()?;
        let now = now_ms / 1000;
        (claims["aud"] == audience && claims["exp"].as_i64()? > now && claims["nbf"].as_i64()? <= now).then_some(claims)
    }
}

/// `text` with every token this hub signed replaced by `***`: a job
/// that prints its identity does not leave it in the log.
pub fn mask(text: &str, kid: &str) -> String {
    if !text.contains("eyJ") {
        return text.to_string();
    }
    let is_tok = |c: char| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.';
    let mut out = String::with_capacity(text.len());
    let mut rest = text;
    while let Some(i) = rest.find("eyJ") {
        out.push_str(&rest[..i]);
        let end = rest[i..].find(|c: char| !is_tok(c)).map(|e| i + e).unwrap_or(rest.len());
        let cand = &rest[i..end];
        let ours = cand
            .split('.')
            .next()
            .and_then(b64url_decode)
            .and_then(|h| serde_json::from_slice::<serde_json::Value>(&h).ok())
            .is_some_and(|h| h["kid"] == kid);
        if ours && cand.matches('.').count() == 2 {
            out.push_str("***");
        } else {
            out.push_str(cand);
        }
        rest = &rest[end..];
    }
    out.push_str(rest);
    out
}

pub(crate) fn b64url(bytes: &[u8]) -> String {
    const T: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
    let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
    for chunk in bytes.chunks(3) {
        let n = chunk.iter().fold(0u32, |acc, b| (acc << 8) | *b as u32) << (8 * (3 - chunk.len()));
        for i in 0..=chunk.len() {
            out.push(T[((n >> (18 - 6 * i)) & 63) as usize] as char);
        }
    }
    out
}

fn b64url_decode(s: &str) -> Option<Vec<u8>> {
    let val = |c: u8| -> Option<u32> {
        Some(match c {
            b'A'..=b'Z' => (c - b'A') as u32,
            b'a'..=b'z' => (c - b'a') as u32 + 26,
            b'0'..=b'9' => (c - b'0') as u32 + 52,
            b'-' => 62,
            b'_' => 63,
            _ => return None,
        })
    };
    let bytes = s.trim_end_matches('=').as_bytes();
    let mut out = Vec::with_capacity(bytes.len() * 3 / 4);
    for chunk in bytes.chunks(4) {
        if chunk.len() == 1 {
            return None;
        }
        let mut n = 0u32;
        for (i, c) in chunk.iter().enumerate() {
            n |= val(*c)? << (18 - 6 * i);
        }
        out.push((n >> 16) as u8);
        if chunk.len() > 2 {
            out.push((n >> 8) as u8);
        }
        if chunk.len() > 3 {
            out.push(n as u8);
        }
    }
    Some(out)
}

/// `GET /.well-known/openid-configuration`.
pub async fn discovery(State(app): State<Shared>) -> Result<Response, ApiError> {
    let issuer = app.oidc().await.map_err(|e| ApiError(StatusCode::INTERNAL_SERVER_ERROR, format!("oidc: {e}")))?;
    Ok(([(axum::http::header::CACHE_CONTROL, "public, max-age=3600")], Json(issuer.discovery(app.config.base_url()))).into_response())
}

/// `GET /.well-known/jwks.json`.
pub async fn jwks(State(app): State<Shared>) -> Result<Response, ApiError> {
    let issuer = app.oidc().await.map_err(|e| ApiError(StatusCode::INTERNAL_SERVER_ERROR, format!("oidc: {e}")))?;
    Ok(([(axum::http::header::CACHE_CONTROL, "public, max-age=3600")], Json(issuer.jwks())).into_response())
}

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

    #[test]
    fn tokens_verify_against_the_published_key_and_are_masked() {
        let dir = tempfile::tempdir().unwrap();
        let issuer = Issuer::load_or_create(dir.path()).unwrap();
        let again = Issuer::load_or_create(dir.path()).unwrap();
        assert_eq!(issuer.kid, again.kid, "the key is kept");
        let claims = JobClaims {
            issuer: "https://hub.example",
            audience: "sts.amazonaws.com",
            owner: "andy",
            name: "demo",
            visibility: "public",
            bookmark: "main",
            run_number: 7,
            job: "deploy",
            state: "ab".repeat(32).leak(),
            runner: "earth-1",
            actor: "andy",
            environment: Some("prod"),
            approved_by: Some("andy"),
            approval: Some("ab12"),
            now_ms: 1_800_000_000_000,
            ttl_s: 900,
        };
        let token = issuer.mint(&claims);
        assert_eq!(token.matches('.').count(), 2);
        let got = again.verify(&token, "sts.amazonaws.com", 1_800_000_500_000).expect("verifies");
        assert_eq!(got["sub"], "repo:andy/demo:bookmark:main");
        assert_eq!(got["run_number"], 7);
        assert_eq!(got["environment"], "prod");
        assert_eq!((got["approved_by"].as_str(), got["approval"].as_str()), (Some("andy"), Some("ab12")));
        let plain = issuer.mint(&JobClaims { environment: None, approved_by: None, approval: None, ..claims });
        assert!(again.verify(&plain, "sts.amazonaws.com", 1_800_000_500_000).unwrap().get("environment").is_none(), "only when approved");
        assert_eq!(got["exp"].as_i64().unwrap() - got["iat"].as_i64().unwrap(), 900);
        assert!(again.verify(&token, "other", 1_800_000_500_000).is_none(), "another audience");
        assert!(again.verify(&token, "sts.amazonaws.com", 1_800_001_000_000).is_none(), "expired");
        let other = Issuer::load_or_create(tempfile::tempdir().unwrap().path()).unwrap();
        assert!(other.verify(&token, "sts.amazonaws.com", 1_800_000_500_000).is_none(), "another key");
        let tampered = format!("{}x", token);
        assert!(again.verify(&tampered, "sts.amazonaws.com", 1_800_000_500_000).is_none());
        let jwks = issuer.jwks();
        assert_eq!(jwks["keys"][0]["kid"], issuer.kid.as_str());
        assert_eq!(b64url_decode(jwks["keys"][0]["e"].as_str().unwrap()).unwrap(), vec![1, 0, 1]);
        assert_eq!(mask(&format!("token={token} done eyJnot.a.token"), &issuer.kid), "token=*** done eyJnot.a.token");
        assert_eq!(b64url_decode(&b64url(b"hello?")).unwrap(), b"hello?");
    }
}