use solana_runtime::bank::{Bank, Result};
use solana_sdk::genesis_block::GenesisBlock;
use solana_sdk::hash::hash;
use solana_sdk::native_program::ProgramError;
use solana_sdk::pubkey::Pubkey;
use solana_sdk::signature::{Keypair, KeypairUtil};
use solana_sdk::system_instruction::SystemInstruction;
use solana_sdk::transaction::{Instruction, InstructionError, Transaction, TransactionError};
use solana_vote_api::vote_instruction::{Vote, VoteInstruction};
use solana_vote_api::vote_state::VoteState;
use solana_vote_api::vote_transaction::VoteTransaction;

struct VoteBank<'a> {
    bank: &'a Bank,
}

impl<'a> VoteBank<'a> {
    fn new(bank: &'a Bank) -> Self {
        bank.add_native_program("solana_vote_program", &solana_vote_api::id());
        Self { bank }
    }

    fn create_vote_account(
        &self,
        from_keypair: &Keypair,
        vote_id: &Pubkey,
        lamports: u64,
    ) -> Result<()> {
        let blockhash = self.bank.last_blockhash();
        let tx = VoteTransaction::new_account(from_keypair, vote_id, blockhash, lamports, 0);
        self.bank.process_transaction(&tx)
    }

    fn create_vote_account_with_delegate(
        &self,
        from_keypair: &Keypair,
        vote_keypair: &Keypair,
        delegate_id: &Pubkey,
        lamports: u64,
    ) -> Result<()> {
        let blockhash = self.bank.last_blockhash();
        let tx = VoteTransaction::new_account_with_delegate(
            from_keypair,
            vote_keypair,
            delegate_id,
            blockhash,
            lamports,
            0,
        );
        self.bank.process_transaction(&tx)
    }

    fn submit_vote(
        &self,
        staking_account: &Pubkey,
        vote_keypair: &Keypair,
        tick_height: u64,
    ) -> Result<VoteState> {
        let blockhash = self.bank.last_blockhash();
        let tx =
            VoteTransaction::new_vote(staking_account, vote_keypair, tick_height, blockhash, 0);
        self.bank.process_transaction(&tx)?;
        self.bank.register_tick(&hash(blockhash.as_ref()));

        let vote_account = self.bank.get_account(&vote_keypair.pubkey()).unwrap();
        Ok(VoteState::deserialize(&vote_account.data).unwrap())
    }
}

#[test]
fn test_vote_bank_basic() {
    let (genesis_block, from_keypair) = GenesisBlock::new(10_000);
    let bank = Bank::new(&genesis_block);
    let vote_bank = VoteBank::new(&bank);

    let vote_keypair = Keypair::new();
    let vote_id = vote_keypair.pubkey();
    vote_bank
        .create_vote_account(&from_keypair, &vote_id, 100)
        .unwrap();

    let vote_state = vote_bank.submit_vote(&vote_id, &vote_keypair, 0).unwrap();
    assert_eq!(vote_state.votes.len(), 1);
}

#[test]
fn test_vote_bank_delegate() {
    let (genesis_block, from_keypair) = GenesisBlock::new(10_000);
    let bank = Bank::new(&genesis_block);
    let vote_bank = VoteBank::new(&bank);
    let vote_keypair = Keypair::new();
    let delegate_keypair = Keypair::new();
    let delegate_id = delegate_keypair.pubkey();
    vote_bank
        .create_vote_account_with_delegate(&from_keypair, &vote_keypair, &delegate_id, 100)
        .unwrap();
}

#[test]
fn test_vote_via_bank_with_no_signature() {
    let (genesis_block, mallory_keypair) = GenesisBlock::new(10_000);
    let bank = Bank::new(&genesis_block);
    let vote_bank = VoteBank::new(&bank);

    let vote_keypair = Keypair::new();
    let vote_id = vote_keypair.pubkey();
    vote_bank
        .create_vote_account(&mallory_keypair, &vote_id, 100)
        .unwrap();

    let mallory_id = mallory_keypair.pubkey();
    let blockhash = bank.last_blockhash();
    let vote_ix = Instruction::new(
        solana_vote_api::id(),
        &VoteInstruction::Vote(Vote::new(0)),
        vec![(vote_id, false)], // <--- attack!! No signature.
    );

    // Sneak in an instruction so that the transaction is signed but
    // the 0th account in the second instruction is not! The program
    // needs to check that it's signed.
    let move_ix = SystemInstruction::new_move(&mallory_id, &vote_id, 1);
    let mut tx = Transaction::new(vec![move_ix, vote_ix]);
    tx.sign(&[&mallory_keypair], blockhash);

    let result = bank.process_transaction(&tx);

    // And ensure there's no vote.
    let vote_account = bank.get_account(&vote_id).unwrap();
    let vote_state = VoteState::deserialize(&vote_account.data).unwrap();
    assert_eq!(vote_state.votes.len(), 0);

    assert_eq!(
        result,
        Err(TransactionError::InstructionError(
            1,
            InstructionError::ProgramError(ProgramError::InvalidArgument)
        ))
    );
}