use rsa::{RsaPrivateKey, RsaPublicKey, Pkcs1v15Sign};
use rsa::pkcs1::DecodeRsaPublicKey;
use rsa::pkcs8::DecodePublicKey;
use sha2::{Sha256, Digest};
use base64::Engine;
use base64::engine::general_purpose::STANDARD as B64;

pub fn sign_request(
    private_key: &RsaPrivateKey,
    key_id: &str,
    method: &str,
    path: &str,
    host: &str,
    body: Option<&[u8]>,
) -> Vec<(String, String)> {
    let date = chrono::Utc::now().format("%a, %d %b %Y %H:%M:%S GMT").to_string();
    let mut headers_to_sign = vec![
        format!("(request-target): {} {}", method.to_lowercase(), path),
        format!("host: {}", host),
        format!("date: {}", date),
    ];

    let mut result = vec![
        ("Host".into(), host.into()),
        ("Date".into(), date.clone()),
    ];

    let signed_headers;
    if let Some(body) = body {
        let digest = B64.encode(Sha256::digest(body));
        let digest_header = format!("SHA-256={}", digest);
        headers_to_sign.push(format!("digest: {}", digest_header));
        result.push(("Digest".into(), digest_header));
        signed_headers = "(request-target) host date digest";
    } else {
        signed_headers = "(request-target) host date";
    }

    let sign_string = headers_to_sign.join("\n");
    let hash = Sha256::digest(sign_string.as_bytes());
    let mut digest_info = vec![0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20];
    digest_info.extend_from_slice(&hash);
    let padding = Pkcs1v15Sign::new_unprefixed();
    let signature = private_key.sign(padding, &digest_info).expect("signing failed");
    let sig_b64 = B64.encode(&signature);

    let sig_header = format!(
        "keyId=\"{}\",algorithm=\"rsa-sha256\",headers=\"{}\",signature=\"{}\"",
        key_id, signed_headers, sig_b64
    );
    result.push(("Signature".into(), sig_header));

    result
}

pub fn verify_request(
    public_key_pem: &str,
    method: &str,
    path: &str,
    headers: &[(String, String)],
    body: Option<&[u8]>,
) -> Result<(), String> {
    let sig_header = headers.iter()
        .find(|(k, _)| k.to_lowercase() == "signature")
        .map(|(_, v)| v.as_str())
        .ok_or("missing Signature header")?;

    let mut _key_id = "";
    let mut signed_headers_str = "";
    let mut signature_b64 = "";
    let mut algorithm = "";

    for part in sig_header.split(',') {
        let part = part.trim();
        if let Some(v) = part.strip_prefix("keyId=\"").and_then(|s| s.strip_suffix('"')) {
            _key_id = v;
        } else if let Some(v) = part.strip_prefix("headers=\"").and_then(|s| s.strip_suffix('"')) {
            signed_headers_str = v;
        } else if let Some(v) = part.strip_prefix("signature=\"").and_then(|s| s.strip_suffix('"')) {
            signature_b64 = v;
        } else if let Some(v) = part.strip_prefix("algorithm=\"").and_then(|s| s.strip_suffix('"')) {
            algorithm = v;
        }
    }

    if signature_b64.is_empty() || signed_headers_str.is_empty() {
        return Err("malformed Signature header".into());
    }

    if !algorithm.is_empty() && algorithm != "rsa-sha256" {
        return Err(format!("unsupported algorithm: {}", algorithm));
    }

    let header_map: std::collections::HashMap<String, String> = headers.iter()
        .map(|(k, v)| (k.to_lowercase(), v.clone()))
        .collect();

    let mut sign_parts = Vec::new();
    for h in signed_headers_str.split(' ') {
        let value = if h == "(request-target)" {
            format!("{} {}", method.to_lowercase(), path)
        } else if let Some(v) = header_map.get(h) {
            v.clone()
        } else {
            return Err(format!("missing signed header: {}", h));
        };
        sign_parts.push(format!("{}: {}", h, value));
    }
    let sign_string = sign_parts.join("\n");

    if signed_headers_str.contains("digest") {
        let digest_val = header_map.get("digest")
            .ok_or("signed headers include digest but Digest header is missing")?;
        let body = body.ok_or("signed headers include digest but request body is empty")?;
        let expected = format!("SHA-256={}", B64.encode(Sha256::digest(body)));
        if digest_val != &expected {
            return Err("digest mismatch".into());
        }
    }

    let public_key = RsaPublicKey::from_pkcs1_pem(public_key_pem)
        .or_else(|_| RsaPublicKey::from_public_key_pem(public_key_pem))
        .map_err(|e| format!("invalid public key: {}", e))?;
    let sig_bytes = B64.decode(signature_b64)
        .map_err(|_| "invalid signature base64")?;
    let hash = Sha256::digest(sign_string.as_bytes());
    let mut digest_info = vec![0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20];
    digest_info.extend_from_slice(&hash);
    let padding = Pkcs1v15Sign::new_unprefixed();

    public_key.verify(padding, &digest_info, &sig_bytes)
        .map_err(|e| format!("signature verification failed: {}", e))?;

    Ok(())
}