extern crate std;

use crate::flow::*;
use keystore::{Block, KdfParams, RECORD_SIZE, Random, Record};
use std::collections::VecDeque;
use std::string::String;
use std::vec;
use std::vec::Vec;

const FAST: KdfParams = KdfParams { memory_kib: 8, passes: 1 };

/// The keystore in memory; `cut_after` writes simulate the power going.
#[derive(Default, Clone)]
struct Device {
    copies: [Vec<u8>; 2],
    writes: usize,
    cut_after: Option<usize>,
    /// The power goes in the middle of the first write that starts a
    /// wipe: only its start lands.
    tear_wipe: bool,
}

impl Device {
    fn new() -> Self {
        Self { copies: [vec![0; RECORD_SIZE], vec![0; RECORD_SIZE]], ..Self::default() }
    }

    fn record(&self) -> Option<Record> {
        keystore::Copies { copies: [&self.copies[0], &self.copies[1]] }.newest().map(|(r, _)| r)
    }
}

impl Keystore for Device {
    fn read(&mut self, copy: usize, out: &mut [u8; RECORD_SIZE]) -> bool {
        out.copy_from_slice(&self.copies[copy]);
        true
    }
    fn write(&mut self, copy: usize, data: &[u8; RECORD_SIZE]) -> bool {
        if self.cut_after.is_some_and(|n| self.writes >= n) {
            return false;
        }
        let starts_wipe = Record::decode(data).is_some_and(|r| r.state == keystore::State::Wiping);
        if self.tear_wipe && starts_wipe {
            // Into the slots: the record (and its checksum) is not whole.
            self.copies[copy][..100].copy_from_slice(&data[..100]);
            self.cut_after = Some(self.writes);
            return false;
        }
        self.writes += 1;
        self.copies[copy] = data.to_vec();
        true
    }
}

struct Counter(u8);

impl Random for Counter {
    fn fill(&mut self, buf: &mut [u8]) -> bool {
        for b in buf {
            self.0 = self.0.wrapping_mul(37).wrapping_add(11);
            *b = self.0;
        }
        true
    }
}

/// An owner typing from a script; what was shown is kept.
struct Owner {
    typing: VecDeque<&'static str>,
    /// On the screen now.
    shown: String,
    /// Ever shown, cleared or not.
    said: String,
}

impl Owner {
    fn types(lines: &[&'static str]) -> Self {
        Self { typing: lines.iter().copied().collect(), shown: String::new(), said: String::new() }
    }
}

impl Prompt for Owner {
    fn say(&mut self, text: &str) {
        self.shown.push_str(text);
        self.said.push_str(text);
    }
    /// What a fresh screen shows: nothing that came before.
    fn clear(&mut self) {
        self.shown.clear();
    }
    fn secret(&mut self, question: &str, out: &mut Secret) -> bool {
        self.shown.push_str(question);
        let Some(line) = self.typing.pop_front() else { return false };
        out.set(line.as_bytes());
        true
    }
}

fn memory() -> Vec<Block> {
    vec![Block::default(); FAST.memory_kib as usize]
}

fn run_with(dev: &mut Device, owner: &mut Owner) -> Result<keystore::Key, Failure> {
    run(dev, &mut Counter(5), owner, FAST, &mut memory())
}

/// A device set up with passcode 1234 and duress 0000; its master key.
fn set_up() -> (Device, [u8; 32]) {
    let mut dev = Device::new();
    let key = run_with(&mut dev, &mut Owner::types(&["1234", "1234", "0000", "0000"])).unwrap();
    (dev, key.0)
}

#[test]
fn a_fresh_device_is_set_up_and_then_unlocks_with_its_passcode() {
    let (mut dev, key) = set_up();
    let mut owner = Owner::types(&["1234"]);
    assert_eq!(run_with(&mut dev, &mut owner).unwrap().0, key, "the same master key every boot");
    assert_eq!(dev.record().unwrap().failures, 0);
}

#[test]
fn setup_insists_on_a_confirmed_passcode_of_four_or_more() {
    let mut dev = Device::new();
    let mut owner = Owner::types(&["1234", "1243", "12", "12", "123456", "123456", "", ""]);
    run_with(&mut dev, &mut owner).unwrap();
    assert!(owner.shown.contains("They differ"));
    assert!(owner.shown.contains("At least 4"));
    let mut unlock = Owner::types(&["123456"]);
    assert!(run_with(&mut dev, &mut unlock).is_ok(), "no duress chosen: the passcode is 123456");
}

#[test]
fn the_duress_passcode_must_differ() {
    let mut dev = Device::new();
    let mut owner = Owner::types(&["1234", "1234", "1234", "1234", "", ""]);
    run_with(&mut dev, &mut owner).unwrap();
    assert!(owner.shown.contains("must differ"));
}

#[test]
fn every_guess_is_counted_before_it_is_checked() {
    let (mut dev, _) = set_up();
    // The power goes after the guess was counted, before the check could
    // clear it: even the right passcode stays counted.
    dev.cut_after = Some(dev.writes + 1);
    let mut owner = Owner::types(&["1234"]);
    assert_eq!(run_with(&mut dev, &mut owner).err(), Some(Failure::Device));
    assert_eq!(dev.record().unwrap().failures, 1, "the guess stays counted");
}

#[test]
fn wrong_passcodes_are_counted_and_a_right_one_resets_the_count() {
    let (mut dev, key) = set_up();
    let mut owner = Owner::types(&["1111", "2222", "1234"]);
    assert_eq!(run_with(&mut dev, &mut owner).unwrap().0, key);
    assert_eq!(owner.shown.matches("Wrong passcode").count(), 2);
    assert_eq!(dev.record().unwrap().failures, 0);
}

#[test]
fn too_many_wrong_passcodes_wipe_the_keys_and_the_device_starts_fresh() {
    let (mut dev, key) = set_up();
    let mut guesses = vec!["9999"; THRESHOLD as usize];
    // Then it looks fresh: the owner sets it up again.
    guesses.extend(["5678", "5678", "", ""]);
    let mut owner = Owner::types(&guesses);
    let new = run_with(&mut dev, &mut owner).unwrap();
    assert!(owner.said.contains("One try left"));
    assert!(owner.shown.ends_with("This device has no passcode yet. Choose one; it protects everything on it.\nNew passcode: Again: A duress passcode wipes the device when it is entered instead. Leave it empty for none.\nDuress passcode: Again: "));
    assert_ne!(new.0, key, "the old master key is gone");
    let mut old = Owner::types(&["1234"]);
    let mut dev2 = dev.clone();
    dev2.cut_after = None;
    assert!(run_with(&mut dev2, &mut old).is_err(), "the old passcode opens nothing");
}

#[test]
fn the_duress_passcode_wipes_silently_and_the_device_looks_fresh() {
    let (mut dev, key) = set_up();
    let mut owner = Owner::types(&["0000", "4321", "4321", "", ""]);
    let new = run_with(&mut dev, &mut owner).unwrap();
    assert!(!owner.shown.contains("Wrong"), "nothing tells it was a duress passcode");
    assert!(owner.shown.starts_with("This device has no passcode yet"), "nor is the passcode prompt left showing");
    assert!(owner.shown.contains("no passcode yet"));
    assert_ne!(new.0, key);
}

#[test]
fn a_wipe_cut_short_is_finished_on_the_next_boot() {
    let (mut dev, key) = set_up();
    // Duress, and the power goes during the wipe (after its first write).
    let marked = dev.writes + 2;
    dev.cut_after = Some(marked + 1);
    assert_eq!(run_with(&mut dev, &mut Owner::types(&["0000"])).err(), Some(Failure::Device));
    dev.cut_after = None;
    // Next boot: the wipe completes before anything is asked; then setup.
    let mut owner = Owner::types(&["4321", "4321", "", ""]);
    let new = run_with(&mut dev, &mut owner).unwrap();
    assert!(owner.shown.starts_with("This device has no passcode yet"), "no passcode is asked of a half-wiped device");
    assert_ne!(new.0, key);
}

#[test]
fn a_duress_cut_short_before_the_wipe_still_wipes_next_boot() {
    let (mut dev, _) = set_up();
    // The power goes right after the duress was recorded (before any wipe write).
    dev.cut_after = Some(dev.writes + 2);
    assert_eq!(run_with(&mut dev, &mut Owner::types(&["0000"])).err(), Some(Failure::Device));
    assert!(dev.record().unwrap().must_wipe());
}

#[test]
fn a_duress_whose_wipe_is_torn_still_wipes_next_boot() {
    let (mut dev, _) = set_up();
    // Counted, recorded, then the wipe's first write is torn.
    dev.tear_wipe = true;
    assert_eq!(run_with(&mut dev, &mut Owner::types(&["0000"])).err(), Some(Failure::Device));
    dev.cut_after = None;
    assert!(dev.record().unwrap().must_wipe(), "the duress was recorded before the wipe began");
}