use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use p256::ecdsa::{VerifyingKey, Signature, signature::Verifier};
use sha2::{Sha256, Digest};
use serde::{Serialize, Deserialize};

#[derive(Debug, Serialize)]
pub struct RegistrationChallenge {
    pub challenge: String,
    pub rp: RpInfo,
    pub user: UserInfo,
    pub timeout: u64,
    #[serde(rename = "pubKeyCredParams")]
    pub pub_key_cred_params: Vec<PubKeyCredParam>,
    #[serde(rename = "authenticatorSelection")]
    pub authenticator_selection: AuthenticatorSelection,
    pub attestation: String,
}

#[derive(Debug, Serialize)]
pub struct AuthenticationChallenge {
    pub challenge: String,
    pub rp_id: String,
    pub timeout: u64,
    #[serde(rename = "allowCredentials")]
    pub allow_credentials: Vec<AllowCredential>,
}

#[derive(Debug, Serialize)]
pub struct RpInfo {
    pub id: String,
    pub name: String,
}

#[derive(Debug, Serialize)]
pub struct UserInfo {
    pub id: String,
    pub name: String,
    #[serde(rename = "displayName")]
    pub display_name: String,
}

#[derive(Debug, Serialize)]
pub struct PubKeyCredParam {
    #[serde(rename = "type")]
    pub cred_type: String,
    pub alg: i64,
}

#[derive(Debug, Serialize)]
pub struct AuthenticatorSelection {
    #[serde(rename = "authenticatorAttachment")]
    pub authenticator_attachment: Option<String>,
    #[serde(rename = "residentKey")]
    pub resident_key: String,
    #[serde(rename = "requireResidentKey")]
    pub require_resident_key: bool,
    #[serde(rename = "userVerification")]
    pub user_verification: String,
}

#[derive(Debug, Serialize)]
pub struct AllowCredential {
    #[serde(rename = "type")]
    pub cred_type: String,
    pub id: String,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct StoredCredential {
    pub credential_id: Vec<u8>,
    pub public_key_x: Vec<u8>,
    pub public_key_y: Vec<u8>,
    pub counter: u32,
}

#[derive(Debug, Deserialize)]
pub struct RegistrationResponse {
    #[allow(dead_code)]
    pub id: String,
    #[serde(rename = "clientDataJSON")]
    pub client_data_json: String,
    #[serde(rename = "attestationObject")]
    pub attestation_object: String,
}

#[derive(Debug, Deserialize)]
pub struct AuthenticationResponse {
    pub id: String,
    #[serde(rename = "clientDataJSON")]
    pub client_data_json: String,
    #[serde(rename = "authenticatorData")]
    pub authenticator_data: String,
    pub signature: String,
}

#[derive(Debug, Deserialize)]
struct ClientData {
    #[serde(rename = "type")]
    ctype: String,
    challenge: String,
    origin: String,
}

pub fn generate_challenge() -> String {
    use rand::Rng;
    let mut rng = rand::rng();
    let bytes: Vec<u8> = (0..32).map(|_| rng.random::<u8>()).collect();
    URL_SAFE_NO_PAD.encode(&bytes)
}

pub fn create_registration_challenge(
    challenge: &str,
    rp_id: &str,
    rp_name: &str,
    user_id: &str,
    user_name: &str,
) -> RegistrationChallenge {
    RegistrationChallenge {
        challenge: challenge.to_string(),
        rp: RpInfo { id: rp_id.to_string(), name: rp_name.to_string() },
        user: UserInfo {
            id: URL_SAFE_NO_PAD.encode(user_id.as_bytes()),
            name: user_name.to_string(),
            display_name: user_name.to_string(),
        },
        timeout: 60000,
        pub_key_cred_params: vec![
            PubKeyCredParam { cred_type: "public-key".into(), alg: -7 }, // ES256
        ],
        authenticator_selection: AuthenticatorSelection {
            authenticator_attachment: None,
            resident_key: "preferred".into(),
            require_resident_key: false,
            user_verification: "preferred".into(),
        },
        attestation: "none".into(),
    }
}

pub fn create_authentication_challenge(
    challenge: &str,
    rp_id: &str,
    credentials: &[StoredCredential],
) -> AuthenticationChallenge {
    AuthenticationChallenge {
        challenge: challenge.to_string(),
        rp_id: rp_id.to_string(),
        timeout: 60000,
        allow_credentials: credentials.iter().map(|c| AllowCredential {
            cred_type: "public-key".into(),
            id: URL_SAFE_NO_PAD.encode(&c.credential_id),
        }).collect(),
    }
}

pub fn verify_registration(
    response: &RegistrationResponse,
    expected_challenge: &str,
    expected_origin: &str,
    expected_rp_id: &str,
) -> Result<StoredCredential, String> {
    // decode and verify clientDataJSON
    let client_data_bytes = URL_SAFE_NO_PAD.decode(&response.client_data_json)
        .map_err(|_| "invalid clientDataJSON base64")?;
    let client_data: ClientData = serde_json::from_slice(&client_data_bytes)
        .map_err(|_| "invalid clientDataJSON")?;

    if client_data.ctype != "webauthn.create" {
        return Err("wrong type in clientDataJSON".into());
    }
    if client_data.challenge != expected_challenge {
        return Err("challenge mismatch".into());
    }
    if client_data.origin != expected_origin {
        return Err(format!("origin mismatch: got {}, expected {}", client_data.origin, expected_origin));
    }

    // decode attestationObject (CBOR)
    let att_bytes = URL_SAFE_NO_PAD.decode(&response.attestation_object)
        .map_err(|_| "invalid attestationObject base64")?;
    let att: ciborium::Value = ciborium::from_reader(&att_bytes[..])
        .map_err(|e| format!("invalid CBOR: {}", e))?;

    let att_map = att.as_map().ok_or("attestationObject is not a map")?;
    let auth_data_val = att_map.iter()
        .find(|(k, _)| k.as_text() == Some("authData"))
        .map(|(_, v)| v)
        .ok_or("missing authData")?;
    let auth_data = auth_data_val.as_bytes().ok_or("authData is not bytes")?;

    if auth_data.len() < 37 {
        return Err("authData too short".into());
    }

    // verify rpIdHash
    let rp_id_hash = &auth_data[..32];
    let expected_hash = Sha256::digest(expected_rp_id.as_bytes());
    if rp_id_hash != &expected_hash[..] {
        return Err("rpIdHash mismatch".into());
    }

    // check flags: bit 0 = UP (user present), bit 6 = AT (attested credential data)
    let flags = auth_data[32];
    if flags & 0x01 == 0 {
        return Err("user not present".into());
    }
    if flags & 0x40 == 0 {
        return Err("no attested credential data".into());
    }

    let counter = u32::from_be_bytes([auth_data[33], auth_data[34], auth_data[35], auth_data[36]]);

    // parse attested credential data (starts at byte 37)
    // aaguid (16 bytes) + credentialIdLength (2 bytes) + credentialId + COSE key
    if auth_data.len() < 55 {
        return Err("authData too short for credential data".into());
    }
    let cred_id_len = u16::from_be_bytes([auth_data[53], auth_data[54]]) as usize;
    let cred_id_end = 55 + cred_id_len;
    if auth_data.len() < cred_id_end {
        return Err("authData too short for credential ID".into());
    }
    let credential_id = auth_data[55..cred_id_end].to_vec();

    // parse COSE key (remaining bytes)
    let cose_key_bytes = &auth_data[cred_id_end..];
    let cose_key: ciborium::Value = ciborium::from_reader(cose_key_bytes)
        .map_err(|e| format!("invalid COSE key CBOR: {}", e))?;

    let key_map = cose_key.as_map().ok_or("COSE key is not a map")?;
    let get_int = |label: i64| -> Option<&ciborium::Value> {
        key_map.iter().find(|(k, _)| {
            k.as_integer().and_then(|i| i64::try_from(i).ok()) == Some(label)
        }).map(|(_, v)| v)
    };

    // verify alg is ES256 (-7)
    let alg = get_int(3)
        .and_then(|v| v.as_integer())
        .and_then(|i| i64::try_from(i).ok())
        .ok_or("missing alg")?;
    if alg != -7 {
        return Err(format!("unsupported algorithm: {}", alg));
    }

    // extract x (-2) and y (-3) coordinates
    let x = get_int(-2).and_then(|v| v.as_bytes()).ok_or("missing x coordinate")?;
    let y = get_int(-3).and_then(|v| v.as_bytes()).ok_or("missing y coordinate")?;

    if x.len() != 32 || y.len() != 32 {
        return Err("invalid key coordinates".into());
    }

    Ok(StoredCredential {
        credential_id,
        public_key_x: x.to_vec(),
        public_key_y: y.to_vec(),
        counter,
    })
}

pub fn verify_authentication(
    response: &AuthenticationResponse,
    credential: &StoredCredential,
    expected_challenge: &str,
    expected_origin: &str,
    expected_rp_id: &str,
) -> Result<u32, String> {
    // decode and verify clientDataJSON
    let client_data_bytes = URL_SAFE_NO_PAD.decode(&response.client_data_json)
        .map_err(|_| "invalid clientDataJSON base64")?;
    let client_data: ClientData = serde_json::from_slice(&client_data_bytes)
        .map_err(|_| "invalid clientDataJSON")?;

    if client_data.ctype != "webauthn.get" {
        return Err("wrong type".into());
    }
    if client_data.challenge != expected_challenge {
        return Err("challenge mismatch".into());
    }
    if client_data.origin != expected_origin {
        return Err("origin mismatch".into());
    }

    // decode authenticatorData
    let auth_data = URL_SAFE_NO_PAD.decode(&response.authenticator_data)
        .map_err(|_| "invalid authenticatorData base64")?;

    if auth_data.len() < 37 {
        return Err("authenticatorData too short".into());
    }

    // verify rpIdHash
    let rp_id_hash = &auth_data[..32];
    let expected_hash = Sha256::digest(expected_rp_id.as_bytes());
    if rp_id_hash != &expected_hash[..] {
        return Err("rpIdHash mismatch".into());
    }

    // check user present flag
    if auth_data[32] & 0x01 == 0 {
        return Err("user not present".into());
    }

    let counter = u32::from_be_bytes([auth_data[33], auth_data[34], auth_data[35], auth_data[36]]);

    // verify signature: sign(authenticatorData || SHA-256(clientDataJSON))
    let client_data_hash = Sha256::digest(&client_data_bytes);
    let mut signed_data = auth_data.to_vec();
    signed_data.extend_from_slice(&client_data_hash);

    // reconstruct P-256 public key
    let mut uncompressed = vec![0x04u8];
    uncompressed.extend_from_slice(&credential.public_key_x);
    uncompressed.extend_from_slice(&credential.public_key_y);

    let verifying_key = VerifyingKey::from_sec1_bytes(&uncompressed)
        .map_err(|e| format!("invalid public key: {}", e))?;

    let sig_bytes = URL_SAFE_NO_PAD.decode(&response.signature)
        .map_err(|_| "invalid signature base64")?;

    // WebAuthn signatures are DER-encoded
    let signature = Signature::from_der(&sig_bytes)
        .map_err(|e| format!("invalid signature DER: {}", e))?;

    verifying_key.verify(&signed_data, &signature)
        .map_err(|e| format!("signature verification failed: {}", e))?;

    Ok(counter)
}