use aes::Aes256;
use aes::cipher::{BlockDecrypt, KeyInit};
use aes::cipher::generic_array::GenericArray;
use base64::{engine::general_purpose::{STANDARD, URL_SAFE}, Engine};
use urlencoding;
use axum::{extract::Query, routing::get, Router};
use serde::Deserialize;
use sha1_smol::Sha1;

#[derive(Debug, Deserialize)]
pub struct WecomVerify {
    msg_signature: String,
    timestamp: String,
    nonce: String,
    echostr: String,
}

const WECOM_TOKEN: &str = "mytoken123456";
const CORP_ID: &str = "wwa7bb7aec981103b4";
const ENCODING_AES_KEY: &str = "PXP7FjoinIPc9WscGymDlf1VwMyBLh1cKJJSJFx2SO8";

fn url_decode(s: &str) -> String {
    match urlencoding::decode(s) {
        Ok(decoded) => decoded.into_owned(),
        Err(_) => s.to_string(),
    }
}

fn decrypt_aes_cbc(key: &[u8; 32], iv: &[u8; 16], ciphertext: &[u8]) -> Vec<u8> {
    let cipher = Aes256::new(key.into());
    let mut plaintext = Vec::with_capacity(ciphertext.len());
    let mut prev_iv = *iv;

    for block in ciphertext.chunks_exact(16) {
        let mut block_arr = GenericArray::clone_from_slice(block);
        cipher.decrypt_block(&mut block_arr);

        // CBC 解密：异或 IV/前一个密文块
        for i in 0..16 {
            block_arr[i] ^= prev_iv[i];
        }

        plaintext.extend_from_slice(&block_arr);
        prev_iv.copy_from_slice(block);
    }

    // 企业微信 echostr 特殊处理：强制去除 PKCS7 填充
    // 因为 echostr 是随机串，解密后填充可能无效，但仍需去除
    let len = plaintext.len();
    if len == 0 { return plaintext; }
    
    let pad = plaintext[len - 1] as usize;
    if pad >= 1 && pad <= 16 && len >= pad {
        // 强制去除填充，不验证有效性
        plaintext.truncate(len - pad);
    }
    
    plaintext
}

fn decrypt_echo_str(encrypted: &str) -> Result<String, String> {
    // 企业微信官方要求：EncodingAESKey 需要在末尾添加 "=" 再解码
    let key = URL_SAFE.decode(format!("{}=", ENCODING_AES_KEY)).map_err(|e| e.to_string())?;
    // 企业微信 URL 验证（echostr）使用标准 Base64 解码
    let ciphertext = STANDARD.decode(encrypted).map_err(|e| e.to_string())?;
    
    if key.len() != 32 {
        return Err(format!("Invalid key length: {}", key.len()));
    }
    
    let mut key_array = [0u8; 32];
    key_array.copy_from_slice(&key);
    
    let iv = &key_array[0..16];
    let mut iv_array = [0u8; 16];
    iv_array.copy_from_slice(iv);
    
    let plaintext = decrypt_aes_cbc(&key_array, &iv_array, &ciphertext);
    
    // 企业微信官方格式：16字节随机串 + 4字节长度(网络序) + 消息内容 + CorpID
    if plaintext.len() < 20 {
        return Err("Decrypted data too short".into());
    }

    // 解析 4 字节长度（网络序，大端）
    let msg_len = u32::from_be_bytes([plaintext[16], plaintext[17], plaintext[18], plaintext[19]]);
    let msg_start = 20;
    let msg_end = msg_start + msg_len as usize;
    
    if msg_end > plaintext.len() {
        return Err("Invalid message length".into());
    }

    // 截取正确的消息内容
    let msg = &plaintext[msg_start..msg_end];
    let result = String::from_utf8_lossy(msg).to_string();

    eprintln!("✅ 最终解密结果: {}", result);
    Ok(result)
}

fn verify_signature(msg_signature: &str, timestamp: &str, nonce: &str, echostr: &str) -> bool {
    let mut arr = vec![WECOM_TOKEN, timestamp, nonce, echostr];
    arr.sort();

    let mut hasher = Sha1::new();
    hasher.update(arr.join("").as_bytes());
    let signature = hasher.digest().to_string();

    signature == msg_signature
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/wecom", get(verify));

    println!("✅ 服务已启动：127.0.0.1:8000");
    println!("Token: {}", WECOM_TOKEN);
    println!("CorpID: {}", CORP_ID);

    axum::serve(
        tokio::net::TcpListener::bind("127.0.0.1:8000").await.unwrap(),
        app
    ).await.unwrap();
}

async fn verify(Query(q): Query<WecomVerify>) -> String {
    let timestamp = url_decode(&q.timestamp);
    let nonce = url_decode(&q.nonce);
    let echostr = url_decode(&q.echostr);
    
    eprintln!("=== 收到企业微信验证请求 ===");
    eprintln!("msg_signature: {}", q.msg_signature);
    eprintln!("timestamp: {}", timestamp);
    eprintln!("nonce: {}", nonce);
    eprintln!("echostr: {}", echostr);
    
    if !verify_signature(&q.msg_signature, &timestamp, &nonce, &echostr) {
        eprintln!("签名验证失败");
        return "invalid signature".to_string();
    }
    
    eprintln!("签名验证成功");

    match decrypt_echo_str(&echostr) {
        Ok(msg) => {
            eprintln!("解密成功，msg: {}", msg);
            msg
        },
        Err(e) => {
            eprintln!("解密失败: {}", e);
            "decrypt error".to_string()
        }
    }
}