use sha2::{Digest, Sha256};

use crate::errors::{AuthError, AuthResult};
use crate::handlers::challenge::{ChallengeResponse, PowProofRequest};
use crate::share::crypto;

pub fn pow_hash(challenge: &ChallengeResponse, nonce: u64, payload: &[u8]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(b"pow.v1\n");
    hasher.update(challenge.context.as_bytes());
    hasher.update([0u8]);

    hasher.update(challenge.difficulty.to_be_bytes());
    hasher.update([0u8]);

    hasher.update(challenge.issued_at.to_be_bytes());
    hasher.update([0u8]);

    hasher.update(challenge.expires_at.to_be_bytes());
    hasher.update([0u8]);

    hasher.update((payload.len() as u64).to_be_bytes());
    hasher.update(payload);
    hasher.update([0u8]);

    hasher.update(nonce.to_be_bytes());

    hasher.finalize().into()
}

pub fn meets_difficulty(hash: &[u8], difficulty: u32) -> bool {
    let hash_bits = (hash.len() * 8) as u32;

    if difficulty == 0 {
        return true;
    }

    if difficulty > hash_bits {
        return false;
    }

    let full_bytes = (difficulty / 8) as usize;
    let remaining_bits = (difficulty % 8) as u8;
    if !hash[..full_bytes].iter().all(|byte| *byte == 0) {
        return false;
    }
    if remaining_bits == 0 {
        return true;
    }
    if hash.len() <= full_bytes {
        return false;
    }
    let mask = 0xffu8 << (8 - remaining_bits);
    (hash[full_bytes] & mask) == 0
}

pub fn verify_pow(
    challenge: &ChallengeResponse,
    proof: &PowProofRequest,
    expected_payload: &[u8],
) -> AuthResult<()> {
    if crypto::now_secs() > challenge.expires_at {
        return Err(AuthError::Expired);
    }

    if proof.payload != expected_payload {
        return Err(AuthError::InvalidPayload);
    }

    let hash = pow_hash(challenge, proof.nonce, &proof.payload);

    if !meets_difficulty(&hash, challenge.difficulty) {
        return Err(AuthError::InvalidProof);
    }

    Ok(())
}