use openpgp::parse::stream::*;
use openpgp::parse::Parse;
use openpgp::*;
use std::fs;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::path::Path;

pub struct KnownKeys {
    keys: Vec<openpgp::TPK>,
}
impl VerificationHelper for &KnownKeys {
    fn get_public_keys(&mut self, _ids: &[KeyID]) -> Result<Vec<TPK>> {
        Ok(self.keys.clone())
    }
    fn check(&mut self, structure: &MessageStructure) -> Result<()> {
        Ok(()) // Implement your verification policy here.
    }
}

impl KnownKeys {
    pub fn load_dir(dir: impl AsRef<Path>) -> io::Result<KnownKeys> {
        let mut keys: Vec<openpgp::TPK> = Vec::with_capacity(3);
        for f in fs::read_dir(dir)? {
            let f = f?;
            if f.metadata()?.is_dir() {
                continue;
            }
            let tpk = openpgp::TPK::from_file(f.path()).unwrap();
            println!("Fingerprint: {}", tpk.fingerprint());
            keys.push(tpk);
        }
        Ok(KnownKeys { keys: keys })
    }

    pub fn verify(&mut self, r: impl Read) -> io::Result<Vec<u8>> {
        let mut content = Vec::with_capacity(2048);
        let helper = &*self;
        let mut v = Verifier::<&KnownKeys>::from_reader(r, helper, None).map_err(|e| {
            io::Error::new(io::ErrorKind::InvalidData, "Failed to verify signature")
        })?;
        let mut buf = [0u8; 512];
        let bp = "brownpaper".as_bytes();
        loop {
            match v.read(&mut buf)? {
                0 => break,
                read => {
                    // first buffer read
                    if content.len() == 0 {
                        if !(buf.len() > bp.len() && bp == &buf[0..bp.len()]) {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                "Failed to verify signature(prefix)",
                            ));
                        } else {
                            // remove prefix
                            content.extend_from_slice(&buf[bp.len()..read])
                        }
                    } else {
                        content.extend_from_slice(&buf[0..read]);
                    }
                }
            }
        }
        if v.read_to_end(&mut content).is_err() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "Signature Mismatch",
            ));
        }
        Ok(content)
    }
}