// Amtrak data decryption
//
// The encrypted response from getTrainsData uses AES-CBC with PBKDF2-SHA1 key derivation.
// Keys (salt, IV, public key) are fetched from RoutesList.v.json and RoutesList.json.
// The last 88 bytes of the response are an encrypted segment containing the private key.

use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use anyhow::{Context, Result};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};

type Aes128CbcEnc = cbc::Encryptor<aes::Aes128>;
type Aes128CbcDec = cbc::Decryptor<aes::Aes128>;

pub struct RouteInfo {
    pub zoom_level: u32,
}

pub struct DecryptionKeys {
    pub public_key: String,
    pub salt: String,
    pub iv: String,
}

/// Sum of fibonacci(0..=9) = 88. This is the length of the encrypted key segment
/// appended to the end of the response.
pub fn master_segment() -> usize {
    let mut fib = vec![0usize, 1];
    for i in 2..=9 {
        fib.push(fib[i - 2] + fib[i - 1]);
    }
    fib.iter().sum()
}

/// All salt entries in the config have length 8, so picking any element
/// and using its length as the index always yields index 8.
pub fn extract_salt(salt_array: &[String]) -> String {
    let idx = salt_array[0].len(); // always 8
    salt_array[idx].clone()
}

/// All IV entries in the config have length 32, so picking any element
/// and using its length as the index always yields index 32.
pub fn extract_iv(iv_array: &[String]) -> String {
    let idx = iv_array[0].len(); // always 32
    iv_array[idx].clone()
}

/// masterZoom = sum of all ZoomLevel values from the route list.
/// Used as the index into the arr[] to get the public key.
pub fn calculate_master_zoom(routes: &[RouteInfo]) -> u32 {
    routes.iter().map(|r| r.zoom_level).sum()
}

/// Derive an AES-128 key from a password using PBKDF2-HMAC-SHA1 (1000 iterations).
fn derive_key(password: &str, salt_hex: &str) -> [u8; 16] {
    let salt = hex::decode(salt_hex).expect("invalid salt hex");
    let mut key = [0u8; 16];
    pbkdf2::pbkdf2_hmac::<sha1::Sha1>(password.as_bytes(), &salt, 1000, &mut key);
    key
}

/// Decrypt base64-encoded AES-128-CBC ciphertext.
pub fn decrypt_aes(
    ciphertext_b64: &str,
    password: &str,
    salt_hex: &str,
    iv_hex: &str,
) -> Result<String> {
    let ciphertext = BASE64.decode(ciphertext_b64).context("invalid base64")?;
    let iv = hex::decode(iv_hex).context("invalid IV hex")?;
    let key = derive_key(password, salt_hex);

    let mut buf = ciphertext.clone();
    let plaintext = Aes128CbcDec::new_from_slices(&key, &iv)
        .context("invalid key/iv length")?
        .decrypt_padded_mut::<Pkcs7>(&mut buf)
        .map_err(|e| anyhow::anyhow!("decryption failed: {}", e))?;

    String::from_utf8(plaintext.to_vec()).context("decrypted data is not valid UTF-8")
}

/// Encrypt plaintext to base64-encoded AES-128-CBC (for tests).
pub fn encrypt_aes(plaintext: &str, password: &str, salt_hex: &str, iv_hex: &str) -> String {
    let iv = hex::decode(iv_hex).expect("invalid IV hex");
    let key = derive_key(password, salt_hex);

    // Buffer needs room for plaintext + up to 16 bytes of PKCS7 padding
    let mut buf = vec![0u8; plaintext.len() + 16];
    buf[..plaintext.len()].copy_from_slice(plaintext.as_bytes());

    let ciphertext = Aes128CbcEnc::new_from_slices(&key, &iv)
        .expect("invalid key/iv length")
        .encrypt_padded_mut::<Pkcs7>(&mut buf, plaintext.len())
        .expect("encryption failed");

    BASE64.encode(ciphertext)
}

/// Split the encrypted response into main data and the key segment.
pub fn split_encrypted_response(blob: &str) -> (&str, &str) {
    let seg_len = master_segment();
    let split_at = blob.len() - seg_len;
    (&blob[..split_at], &blob[split_at..])
}

/// Extract the private key from the decrypted segment (format: "uuid|timestamp").
pub fn extract_private_key(decrypted_segment: &str) -> &str {
    decrypted_segment
        .split('|')
        .next()
        .expect("segment should contain '|'")
}

/// Decrypt train data using a known segment length (for testing with non-88 segments).
pub fn decrypt_train_data_with_segment_len(
    encrypted_blob: &str,
    keys: &DecryptionKeys,
    segment_len: usize,
) -> Result<String> {
    let split_at = encrypted_blob.len() - segment_len;
    let main_data = &encrypted_blob[..split_at];
    let key_segment = &encrypted_blob[split_at..];

    // Step 1: Decrypt the key segment using the public key
    let decrypted_segment = decrypt_aes(key_segment, &keys.public_key, &keys.salt, &keys.iv)
        .context("failed to decrypt key segment")?;

    // Step 2: Extract the private key
    let private_key = extract_private_key(&decrypted_segment);

    // Step 3: Decrypt the main data using the private key
    decrypt_aes(main_data, private_key, &keys.salt, &keys.iv)
        .context("failed to decrypt main train data")
}

/// Decrypt train data from the raw API response.
pub fn decrypt_train_data(encrypted_blob: &str, keys: &DecryptionKeys) -> Result<String> {
    decrypt_train_data_with_segment_len(encrypted_blob, keys, master_segment())
}

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

    // --- Fibonacci / master segment ---

    #[test]
    fn test_master_segment_is_88() {
        // The master segment length is the sum of fibonacci(0..=9)
        // which equals 0+1+1+2+3+5+8+13+21+34 = 88
        assert_eq!(master_segment(), 88);
    }

    // --- Key extraction ---

    #[test]
    fn test_extract_keys_from_config() {
        // All salt entries have length 8, so any random pick -> index 8
        let salt_array: Vec<String> = (0..20).map(|i| format!("{:08x}", i * 12345678)).collect();
        assert_eq!(salt_array[0].len(), 8);
        let salt = extract_salt(&salt_array);
        assert_eq!(salt, salt_array[8]);

        // All IV entries have length 32, so any random pick -> index 32
        let iv_array: Vec<String> = (0..50)
            .map(|i| format!("{:032x}", i * 12345678u128))
            .collect();
        assert_eq!(iv_array[0].len(), 32);
        let iv = extract_iv(&iv_array);
        assert_eq!(iv, iv_array[32]);
    }

    #[test]
    fn test_master_zoom_calculation() {
        // masterZoom is the sum of all ZoomLevel values from the route list
        let routes = vec![
            RouteInfo { zoom_level: 6 },
            RouteInfo { zoom_level: 5 },
            RouteInfo { zoom_level: 4 },
        ];
        assert_eq!(calculate_master_zoom(&routes), 15);
    }

    // --- AES decryption ---

    #[test]
    fn test_aes_decrypt_roundtrip() {
        // Encrypt something with known params, then decrypt it
        let plaintext = "hello amtrak";
        let password = "test-password-uuid";
        let salt_hex = "9a3686ac";
        let iv_hex = "c6eb2f7f5c4740c1a2f708fefd947d39";

        let encrypted = encrypt_aes(plaintext, password, salt_hex, iv_hex);
        let decrypted = decrypt_aes(&encrypted, password, salt_hex, iv_hex).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_decrypt_aes_bad_data_returns_error() {
        let result = decrypt_aes(
            "not-valid-base64!!!",
            "key",
            "9a3686ac",
            "c6eb2f7f5c4740c1a2f708fefd947d39",
        );
        assert!(result.is_err());
    }

    // --- Full pipeline ---

    #[test]
    fn test_split_encrypted_response() {
        // Given a blob of length 200, the last 88 chars are the encrypted key segment
        let blob = "A".repeat(112) + &"B".repeat(88);
        let (main_data, key_segment) = split_encrypted_response(&blob);
        assert_eq!(main_data.len(), 112);
        assert_eq!(key_segment.len(), 88);
        assert_eq!(key_segment, "B".repeat(88));
    }

    #[test]
    fn test_extract_private_key_from_decrypted_segment() {
        // The decrypted segment is "private-key-uuid|2026-03-15T11:05:05.000Z"
        let segment = "37b1b306-903f-4ae7-ae93-ceb9e267215d|2026-03-15T11:05:05.000Z";
        let key = extract_private_key(segment);
        assert_eq!(key, "37b1b306-903f-4ae7-ae93-ceb9e267215d");
    }

    #[test]
    fn test_full_decrypt_pipeline_with_known_data() {
        // Create a mini encrypted blob using known keys, then decrypt it
        let salt_hex = "9a3686ac";
        let iv_hex = "c6eb2f7f5c4740c1a2f708fefd947d39";
        let public_key = "69af143c-e8cf-47f8-bf09-fc1f61e5cc33";
        let private_key = "37b1b306-903f-4ae7-ae93-ceb9e267215d";

        // The actual train JSON (simplified)
        let train_json = r#"{"type":"FeatureCollection","features":[]}"#;

        // Encrypt the train data with the private key
        let encrypted_main = encrypt_aes(train_json, private_key, salt_hex, iv_hex);

        // Create the key segment: "private_key|timestamp"
        let key_segment_plain = format!("{}|2026-03-15T11:05:05.000Z", private_key);
        let encrypted_segment = encrypt_aes(&key_segment_plain, public_key, salt_hex, iv_hex);

        // Combine: encrypted_main + encrypted_segment (segment must be exactly 88 chars)
        // Note: in reality the segment is always 88 chars. For this test we verify the pipeline
        // works with our split/decrypt logic.
        let combined = format!("{}{}", encrypted_main, encrypted_segment);

        let keys = DecryptionKeys {
            public_key: public_key.to_string(),
            salt: salt_hex.to_string(),
            iv: iv_hex.to_string(),
        };

        let segment_len = encrypted_segment.len();
        let result = decrypt_train_data_with_segment_len(&combined, &keys, segment_len);
        assert!(result.is_ok(), "decrypt failed: {:?}", result.err());
        assert_eq!(result.unwrap(), train_json);
    }
}