pub mod util;
pub mod pcg;
pub mod sbox;
pub mod gift256;
pub mod hash;
pub mod solver;
pub mod feistel;

use base64::{engine::general_purpose::STANDARD, Engine};
use serde::Deserialize;

#[derive(Deserialize)]
#[allow(dead_code)]
struct JwtPayload {
    #[serde(default)]
    f: u32,
    #[serde(default)]
    s: u32,
    #[serde(default)]
    t: String,
    d: String,
    #[serde(default)]
    l: String,
    #[serde(default)]
    i: String,
    #[serde(default)]
    e: u64,
    #[serde(default)]
    n: String,
    #[serde(default)]
    c: u32,
}

/// Takes a raw JWT string, runs the full PoW pipeline, returns the base64 proof.
pub fn solve_jwt(jwt_req: &str) -> Result<String, String> {
    // 1. Decode JWT payload (base64url, no signature verification)
    let parts: Vec<&str> = jwt_req.split('.').collect();
    if parts.len() < 2 {
        return Err("Invalid JWT format".into());
    }
    let payload_b64 = parts[1];
    let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(payload_b64)
        .map_err(|e| format!("Invalid base64 in JWT payload: {e}"))?;
    let payload: JwtPayload = serde_json::from_slice(&payload_bytes)
        .map_err(|e| format!("Invalid JSON in JWT payload: {e}"))?;

    // 2. Decode challenge data from `d` field
    let challenge_data = if let Some(eq_pos) = payload.d.find('=') {
        let split_pos = eq_pos + payload.d[eq_pos..].chars().take_while(|&c| c == '=').count();
        let part1 = &payload.d[..split_pos];
        let part2 = &payload.d[split_pos..];

        let mut data = STANDARD.decode(part1)
            .map_err(|e| format!("Invalid base64 in challenge part1: {e}"))?;
        if !part2.is_empty() {
            let p2 = STANDARD.decode(part2).unwrap_or_else(|_| {
                let padded = format!("{}{}", part2, "=".repeat((4 - part2.len() % 4) % 4));
                STANDARD.decode(&padded).expect("Invalid base64 in challenge part2")
            });
            data.extend_from_slice(&p2);
        }
        data
    } else {
        STANDARD.decode(&payload.d)
            .map_err(|e| format!("Invalid base64 in challenge data: {e}"))?
    };

    // 3. Split into encrypted_payload and seed_bytes
    let (encrypted_payload, seed_bytes) = if let Some(eq_pos) = payload.d.find('=') {
        let split_pos = eq_pos + payload.d[eq_pos..].chars().take_while(|&c| c == '=').count();
        let part1_len = STANDARD.decode(&payload.d[..split_pos])
            .map_err(|e| format!("Invalid base64 in part1: {e}"))?.len();
        (&challenge_data[..part1_len], &challenge_data[part1_len..])
    } else {
        (challenge_data.as_slice(), &[] as &[u8])
    };

    // 4. Feistel decrypt
    if encrypted_payload.len() < 33 {
        return Err(format!(
            "Encrypted payload too short: {} bytes (need >= 33)",
            encrypted_payload.len()
        ));
    }

    let (key_material, _control) = feistel::feistel_decrypt(encrypted_payload, 0);

    let difficulty = if payload.c > 0 { payload.c } else { 1000 };

    let challenge = solver::Challenge {
        key_material,
        difficulty,
    };

    // 5. Derive PCG seed from seed_bytes
    let seed = if seed_bytes.len() >= 8 {
        u64::from_le_bytes([
            seed_bytes[0], seed_bytes[1], seed_bytes[2], seed_bytes[3],
            seed_bytes[4], seed_bytes[5], seed_bytes[6], seed_bytes[7],
        ])
    } else {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos() as u64
    };

    // 6. Solve and encode
    let proof = solver::solve(&challenge, seed);
    Ok(STANDARD.encode(&proof))
}
