use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
use risc0_zkvm::sha::{Impl, Sha256};

const PRIZE: u128 = 150;

type Instruction = u128;

struct Challenge {
    difficulty: u8,
    seed: [u8; 32],
}

impl Challenge {
    fn new(bytes: &[u8]) -> Self {
        assert_eq!(bytes.len(), 33);
        let difficulty = bytes[0];
        assert!(difficulty <= 32);

        let mut seed = [0; 32];
        seed.copy_from_slice(&bytes[1..]);
        Self { difficulty, seed }
    }

    fn is_nonce_valid(&self, nonce: Instruction) -> bool {
        let mut bytes = [0; 32 + 16];
        bytes.copy_from_slice(&self.seed);
        bytes[32..].copy_from_slice(&nonce.to_le_bytes());
        let digest: [u8; 32] = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
        let difficulty = self.difficulty as usize;
        digest[..difficulty].iter().all(|&b| b == 0)
    }

    fn next_data(self) -> [u8; 33] {
        let mut result = [0; 33];
        result[0] = self.difficulty;
        result.copy_from_slice(Impl::hash_bytes(&self.seed).as_bytes());
        result
    }
}

/// A pinata program
fn main() {
    // Read input accounts.
    // It is expected to receive only two accounts: [pinata_account, winner_account]
    let ProgramInput {
        pre_states,
        instruction: nonce,
    } = read_nssa_inputs::<Instruction>();

    let [pinata, winner] = match pre_states.try_into() {
        Ok(array) => array,
        Err(_) => return,
    };

    let data = Challenge::new(&pinata.account.data);

    if !data.is_nonce_valid(nonce) {
        return;
    }

    let mut pinata_post = pinata.account.clone();
    let mut winner_post = winner.account.clone();
    pinata_post.balance -= PRIZE;
    pinata_post.data = data.next_data().to_vec();
    winner_post.balance += PRIZE;

    write_nssa_outputs(vec![pinata, winner], vec![pinata_post, winner_post]);
}