//! A minimal OpenID provider for tests: discovery, JWKS, token and userinfo endpoints with
//! RS256 tokens. The private key is a throwaway test fixture and protects nothing.
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use jsonwebtoken::{Algorithm, EncodingKey, Header};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpListener,
};

const PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY-----
MIIEuwIBADANBgkqhkiG9w0BAQEFAASCBKUwggShAgEAAoIBAQDI4xpMobGZ91uK
+uV7YFtGg0KorieswaP4ZhIqg3LkxIjpaMggyw0Fc4jliFPxn+uZAzKEF5lBrrdt
XBllxKayC6dLEHFBvrglZMwyzsj6fgssne/N1YRWsQsQ48d2KANzfGpkBB8/29g7
YhpNK6lWZo54yGFcFa2g68/yqbpQ6dSo191/CSqOnC+R2pdMe9NzI8kQBtZPkcft
hmQ05IQJoyNHgBgwUQO2vB7J23haATlvQUrZlv0ASFYgxJDKeLZOPvdLGEmezIvc
WZYdI5qEqiNDr9KRWLS/sDYxxOQkVxRP4OoWy5svru1ptJ3QEG50uUnPZkdpX/Ny
UaMFOieJAgMBAAECgf8UyqFa5XWp/ODExbcM7nE2ydLmiAuRJgvXmHfdLY72qqne
ajvq17OcvVVs3g3B3wE5PaKWPCAdhJdEtUCCB/HXTO25hOTOYAhTrdLuh6cvIUxF
wqlJGVNM7tM3ZS8mW8b7maLjMxQtep605CxYxjU5c+wD54yuIMkl58Q3bLENuM5G
Y5GeXwmAv4yrJXuFIe/h2F6A+p4IEQ64QF6VQ/rDeG8kniLXWNZjAndqGxPdLDjC
4YdXpSWeYM/dTGKUyt0SwE64vxytMNyoRWAO4i4v3JJbrST0SmAM2RwM803Ks0Ys
H9VvFoYt5kGhETOz4OKa3m6HrMwdUy7vnUYd870CgYEA59CLNpjU3Dl57RJRlHRw
2FAPM0uM9qBWiahDyhQAeHeypUFGeu1yZHZCH9Z5oOo28xRyAB27R8ozFX0JdPYl
HBuJH63boff4/DuzjmnAK92889m+cbyPcAd6BjOo545cMRiAXvu3b8cz2Mlks1bj
uWS5RIVbQvuiPqNEbHNHc5UCgYEA3diG58Ziug75V6k8U9wHyqWWRyvQdwdJQmey
FAB0LQdSsHUkPteGl2U6Ml7EL2AuqoZK5kVRK+8aGhBHyOlsjwj/sCfQfiTNDgve
MP7pWfUUjctO8ZvKHG3SG87BG2kLAW6l/vuhDfLwQe/PV8hC0yDegFfF1R4nQVYf
7XFv5yUCgYBftygEHw6k78Z8MLgmh/2gE05za5qEqqeQpz12TtPy6dmFinJDh49Y
E8wxwSjvXEXcEHERwzT063Q0RzCEbq3b8TKpNewd8HmKDEyTXO7QLfo65YlRBP5k
z3F+5g2zroA6dVdr4JtZN/Mca+t8cUAN6Jzmwflwqt3O5zL1ZJP4EQKBgQCt50Sp
f9x4mbBgRWvjGQKKmo9ZL0F1bSLCiWMRRuHLQohRhyJomK8LEvEx9xO3R6CNUER6
ii/9TquZbtW26zSMU7OLtUo8+EA+GJOOuOTX/cSk+pDh2QvwiS5BWiwKfqd/8ery
HPK9aoloRsK0v2mS85s35JZSYBbJoufYTCo6eQKBgGW1vIlLzHRMeboj82uW0yw2
Jkg6Ig45qtVhY89267E8cI7gdC92Y715nFR4v+0eiWo08k+FIaHyBsah0vklj1Q6
fj6ir9inIzkkrFmjs6u/g4h2GNGUXOJJwJzOU78C6oeB3G5OUM/pWFx0zypRGH/q
HpM1TJ5Fkyog20WBQS9+
-----END PRIVATE KEY-----";
const MODULUS: &str = "yOMaTKGxmfdbivrle2BbRoNCqK4nrMGj-GYSKoNy5MSI6WjIIMsNBXOI5YhT8Z_rmQMyhBeZQa63bVwZZcSmsgunSxBxQb64JWTMMs7I-n4LLJ3vzdWEVrELEOPHdigDc3xqZAQfP9vYO2IaTSupVmaOeMhhXBWtoOvP8qm6UOnUqNfdfwkqjpwvkdqXTHvTcyPJEAbWT5HH7YZkNOSECaMjR4AYMFEDtrweydt4WgE5b0FK2Zb9AEhWIMSQyni2Tj73SxhJnsyL3FmWHSOahKojQ6_SkVi0v7A2McTkJFcUT-DqFsubL67tabSd0BBudLlJz2ZHaV_zclGjBToniQ";
pub const CLIENT_ID: &str = "meguri-test";
pub const CLIENT_SECRET: &str = "test-client-secret";

/// What a single sign-in attempt looks like to the relying party.
#[derive(Clone)]
pub struct Scenario {
    pub claims: Value,
    pub userinfo: Option<Value>,
    pub issuer: Option<String>,
    pub audience: Option<Value>,
    pub expires_in: i64,
    pub nonce: Option<String>,
    pub corrupt_signature: bool,
    pub token_status: u16,
}
impl Scenario {
    pub fn new(claims: Value) -> Self {
        Self {
            claims,
            userinfo: None,
            issuer: None,
            audience: None,
            expires_in: 300,
            nonce: None,
            corrupt_signature: false,
            token_status: 200,
        }
    }
}

struct Grant {
    scenario: Scenario,
    nonce: String,
    challenge: String,
}
#[derive(Default)]
pub struct State {
    codes: HashMap<String, Grant>,
    tokens: HashMap<String, Value>,
    pub token_requests: usize,
    pub userinfo_requests: usize,
}
pub struct Server {
    pub base: String,
    pub state: Arc<Mutex<State>>,
    task: tokio::task::JoinHandle<()>,
}
impl Drop for Server {
    fn drop(&mut self) {
        self.task.abort();
    }
}
impl Server {
    pub fn issuer(&self) -> String {
        format!("{}/realms/corp", self.base)
    }
    /// Registers an authorization the way a browser visit to the provider would.
    pub fn authorize(&self, authorization_url: &str, scenario: Scenario) -> String {
        let url = reqwest::Url::parse(authorization_url).unwrap();
        let query: HashMap<_, _> = url.query_pairs().into_owned().collect();
        assert_eq!(query["response_type"], "code");
        assert_eq!(query["client_id"], CLIENT_ID);
        assert_eq!(query["code_challenge_method"], "S256");
        assert!(query["scope"].split(' ').any(|s| s == "openid"));
        let code = format!("code-{}", uuid::Uuid::new_v4().simple());
        self.state.lock().unwrap().codes.insert(
            code.clone(),
            Grant {
                scenario,
                nonce: query["nonce"].clone(),
                challenge: query["code_challenge"].clone(),
            },
        );
        code
    }
    pub async fn start() -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let base = format!("http://{}", listener.local_addr().unwrap());
        let state = Arc::new(Mutex::new(State::default()));
        let (shared, root) = (state.clone(), base.clone());
        let task = tokio::spawn(async move {
            loop {
                let (mut socket, _) = listener.accept().await.unwrap();
                let (shared, root) = (shared.clone(), root.clone());
                tokio::spawn(async move {
                    let mut buffer = vec![];
                    let (head_end, length) = loop {
                        let mut chunk = [0; 4096];
                        let n = socket.read(&mut chunk).await.unwrap();
                        if n == 0 {
                            return;
                        }
                        buffer.extend_from_slice(&chunk[..n]);
                        if let Some(end) = buffer.windows(4).position(|w| w == b"\r\n\r\n") {
                            let head = String::from_utf8_lossy(&buffer[..end]);
                            let length = head
                                .lines()
                                .find_map(|l| {
                                    l.to_lowercase()
                                        .strip_prefix("content-length: ")
                                        .and_then(|v| v.parse::<usize>().ok())
                                })
                                .unwrap_or(0);
                            break (end + 4, length);
                        }
                    };
                    while buffer.len() < head_end + length {
                        let mut chunk = [0; 4096];
                        let n = socket.read(&mut chunk).await.unwrap();
                        if n == 0 {
                            return;
                        }
                        buffer.extend_from_slice(&chunk[..n]);
                    }
                    let head = String::from_utf8_lossy(&buffer[..head_end]).to_string();
                    let body =
                        String::from_utf8_lossy(&buffer[head_end..head_end + length]).to_string();
                    let mut first = head.lines().next().unwrap().split_whitespace();
                    let method = first.next().unwrap().to_owned();
                    let path = first.next().unwrap().to_owned();
                    let (status, result) = handle(
                        &mut shared.lock().unwrap(),
                        &root,
                        &method,
                        &path,
                        &head,
                        &body,
                    );
                    let body = result.to_string();
                    let response = format!(
                        "HTTP/1.1 {status} Result\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
                        body.len()
                    );
                    let _ = socket.write_all(response.as_bytes()).await;
                });
            }
        });
        Self { base, state, task }
    }
}

fn header<'a>(head: &'a str, name: &str) -> Option<&'a str> {
    head.lines().find_map(|l| {
        let (key, value) = l.split_once(':')?;
        key.eq_ignore_ascii_case(name).then(|| value.trim())
    })
}

fn sign(claims: &Value, corrupt: bool) -> String {
    let mut header = Header::new(Algorithm::RS256);
    header.kid = Some("test-key".into());
    let mut token = jsonwebtoken::encode(
        &header,
        claims,
        &EncodingKey::from_rsa_pem(PRIVATE_KEY.as_bytes()).unwrap(),
    )
    .unwrap();
    if corrupt {
        let last = token.pop().unwrap();
        token.push(if last == 'A' { 'B' } else { 'A' });
    }
    token
}

fn handle(
    s: &mut State,
    root: &str,
    method: &str,
    path: &str,
    head: &str,
    body: &str,
) -> (u16, Value) {
    let issuer = format!("{root}/realms/corp");
    let path = path.split('?').next().unwrap();
    match (method, path) {
        ("GET", "/realms/corp/.well-known/openid-configuration") => (
            200,
            json!({
                "issuer":issuer,
                "authorization_endpoint":format!("{issuer}/protocol/openid-connect/auth"),
                "token_endpoint":format!("{issuer}/protocol/openid-connect/token"),
                "jwks_uri":format!("{issuer}/protocol/openid-connect/certs"),
                "userinfo_endpoint":format!("{issuer}/protocol/openid-connect/userinfo"),
            }),
        ),
        ("GET", "/realms/corp/protocol/openid-connect/certs") => (
            200,
            json!({"keys":[{"kty":"RSA","use":"sig","alg":"RS256","kid":"test-key","n":MODULUS,"e":"AQAB"}]}),
        ),
        ("POST", "/realms/corp/protocol/openid-connect/token") => {
            s.token_requests += 1;
            let expected = format!(
                "Basic {}",
                base64::engine::general_purpose::STANDARD
                    .encode(format!("{CLIENT_ID}:{CLIENT_SECRET}"))
            );
            if header(head, "authorization") != Some(expected.as_str()) {
                return (401, json!({"error":"invalid_client"}));
            }
            let form: HashMap<_, _> = reqwest::Url::parse(&format!("http://x/?{body}"))
                .unwrap()
                .query_pairs()
                .into_owned()
                .collect();
            let Some(grant) = form.get("code").and_then(|c| s.codes.remove(c)) else {
                return (400, json!({"error":"invalid_grant"}));
            };
            let verifier = form.get("code_verifier").cloned().unwrap_or_default();
            if form.get("grant_type").map(String::as_str) != Some("authorization_code")
                || URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes())) != grant.challenge
                || !form
                    .get("redirect_uri")
                    .is_some_and(|u| u.ends_with("/api/oidc/callback"))
            {
                return (400, json!({"error":"invalid_grant"}));
            }
            let scenario = grant.scenario;
            if scenario.token_status != 200 {
                return (scenario.token_status, json!({"error":"server_error"}));
            }
            let now = chrono::Utc::now().timestamp();
            let mut claims = scenario.claims.clone();
            claims["iss"] = json!(scenario.issuer.unwrap_or(issuer));
            claims["aud"] = scenario.audience.unwrap_or(json!(CLIENT_ID));
            claims["iat"] = json!(now);
            claims["exp"] = json!(now + scenario.expires_in);
            claims["nonce"] = json!(scenario.nonce.unwrap_or(grant.nonce));
            let access = format!("access-{}", uuid::Uuid::new_v4().simple());
            if let Some(info) = scenario.userinfo {
                s.tokens.insert(access.clone(), info);
            }
            (
                200,
                json!({"id_token":sign(&claims, scenario.corrupt_signature),"access_token":access,"token_type":"Bearer"}),
            )
        }
        ("GET", "/realms/corp/protocol/openid-connect/userinfo") => {
            s.userinfo_requests += 1;
            let token = header(head, "authorization")
                .and_then(|h| h.strip_prefix("Bearer "))
                .unwrap_or_default();
            match s.tokens.get(token) {
                Some(info) => (200, info.clone()),
                None => (401, json!({"error":"invalid_token"})),
            }
        }
        _ => (404, json!({})),
    }
}