extern crate std;

use crate::encrypted::{Encrypted, area_of};
use crate::filesystem::{DirEntry, FileHandle, FileName, FileSystem};
use std::collections::BTreeMap;
use std::string::{String, ToString};
use std::vec;
use std::vec::Vec;

/// Files in memory, with positions: what the layer sits on in tests.
#[derive(Default)]
struct Memory {
    files: BTreeMap<String, Vec<u8>>,
    dirs: Vec<String>,
    open: BTreeMap<u16, (String, usize)>,
    next: u16,
}

fn key(path: &[u8]) -> String {
    String::from_utf8_lossy(path.strip_prefix(b"/").unwrap_or(path)).into_owned()
}

impl Memory {
    fn handle(&mut self, path: String) -> FileHandle {
        self.next = (self.next + 1) % 16;
        self.open.insert(self.next, (path, 0));
        FileHandle::new(self.next)
    }
}

impl FileSystem for Memory {
    fn open(&mut self, path: &[u8]) -> Option<FileHandle> {
        let k = key(path);
        self.files.contains_key(&k).then(|| self.handle(k))
    }
    fn create(&mut self, path: &[u8]) -> Option<FileHandle> {
        let k = key(path);
        if self.files.contains_key(&k) {
            return None;
        }
        self.files.insert(k.clone(), Vec::new());
        Some(self.handle(k))
    }
    fn read(&mut self, h: FileHandle, buf: &mut [u8]) -> Option<usize> {
        let (path, pos) = self.open.get_mut(&h.raw())?;
        let data = &self.files[path.as_str()];
        let n = buf.len().min(data.len().saturating_sub(*pos));
        buf[..n].copy_from_slice(&data[*pos..*pos + n]);
        *pos += n;
        Some(n)
    }
    fn write(&mut self, h: FileHandle, buf: &[u8]) -> Option<usize> {
        let (path, pos) = self.open.get_mut(&h.raw())?;
        let data = self.files.get_mut(path.as_str())?;
        if data.len() < *pos + buf.len() {
            data.resize(*pos + buf.len(), 0);
        }
        data[*pos..*pos + buf.len()].copy_from_slice(buf);
        *pos += buf.len();
        Some(buf.len())
    }
    fn seek(&mut self, h: FileHandle, offset: usize) -> bool {
        let Some((path, pos)) = self.open.get_mut(&h.raw()) else { return false };
        if offset > self.files[path.as_str()].len() {
            return false;
        }
        *pos = offset;
        true
    }
    fn close(&mut self, h: FileHandle) {
        self.open.remove(&h.raw());
    }
    fn file_size(&mut self, h: FileHandle) -> Option<usize> {
        let (path, _) = self.open.get(&h.raw())?;
        Some(self.files[path.as_str()].len())
    }
    /// The direct children: files, and folders made or implied by a file.
    fn list_dir(&mut self, path: &[u8], entries: &mut [DirEntry]) -> Option<usize> {
        let dir = key(path);
        let prefix = std::format!("{dir}/");
        if !self.dirs.contains(&dir) && !self.files.keys().any(|k| k.starts_with(&prefix)) {
            return None;
        }
        let mut seen: Vec<String> = Vec::new();
        let mut n = 0;
        let names = self.files.keys().map(|k| (k, true)).chain(self.dirs.iter().map(|d| (d, false)));
        for (name, file) in names.filter(|(k, _)| k.starts_with(&prefix)) {
            let rest = &name[prefix.len()..];
            let (child, deeper) = rest.split_once('/').map_or((rest, false), |(c, _)| (c, true));
            if seen.iter().any(|s| s == child) || n == entries.len() {
                continue;
            }
            seen.push(child.to_string());
            let size = if file && !deeper { self.files[name.as_str()].len() } else { 0 };
            entries[n] = DirEntry { name: FileName::from_bytes(child.as_bytes()), size, is_dir: deeper || !file };
            n += 1;
        }
        Some(n)
    }
    fn delete(&mut self, path: &[u8]) -> bool {
        self.files.remove(&key(path)).is_some()
    }
    fn mkdir(&mut self, path: &[u8]) -> bool {
        let k = key(path);
        if !self.dirs.contains(&k) {
            self.dirs.push(k);
        }
        true
    }
    /// Only an empty folder.
    fn rmdir(&mut self, path: &[u8]) -> bool {
        let k = key(path);
        let prefix = std::format!("{k}/");
        let empty = !self.files.keys().chain(self.dirs.iter()).any(|f| f.starts_with(&prefix));
        let was = self.dirs.contains(&k) || !empty;
        if empty {
            self.dirs.retain(|d| *d != k);
        }
        empty && was
    }
}

const MASTER: [u8; 32] = [9; 32];

fn unlocked() -> Encrypted<Memory> {
    let mut fs = Encrypted::new(Memory::default());
    fs.unlock(&MASTER, &[3; 32], b"u1");
    fs
}

fn put(fs: &mut impl FileSystem, path: &[u8], data: &[u8]) {
    let h = fs.create(path).expect("create");
    assert_eq!(fs.write(h, data), Some(data.len()));
    fs.close(h);
}

fn get(fs: &mut impl FileSystem, path: &[u8]) -> Option<Vec<u8>> {
    let h = fs.open(path)?;
    let size = fs.file_size(h)?;
    let mut out = vec![0u8; size + 10];
    let n = fs.read(h, &mut out);
    fs.close(h);
    out.truncate(n?);
    Some(out)
}

fn on_disk(fs: &Encrypted<Memory>, path: &str) -> Vec<u8> {
    fs.inner().files[path].clone()
}

#[test]
fn areas_are_told_by_path() {
    for (p, a) in [
        (&b"users/u1/home/readme.txt"[..], "sealed"),
        (b"/users/u1/home", "sealed"),
        (b"users/u1/permissions/kos.snake", "sealed"),
        (b"users/u1/data/kos.notes/a", "sealed"),
        (b"users/u1/apps/kos.notes/notes", "sealed"),
        (b"users/u1/data/kos.notes", "neither"),
        (b"users/u1/apps", "neither"),
        (b"users/u1/loose", "neither"),
        (b"users/u1", "neither"),
        (b"users//home", "neither"),
        (b"users", "plain"),
        (b"bin/shell", "plain"),
        (b"apps/kos.cat/cat", "plain"),
        (b"factory/home/readme.txt", "plain"),
        (b"home/x", "plain"),
        (b"usersx/u1/home/x", "plain"),
        (b"selftest.txt", "plain"),
    ] {
        assert_eq!(area_of(p), a, "{}", String::from_utf8_lossy(p));
    }
}

#[test]
fn sealed_files_read_back_and_are_noise_on_disk() {
    let mut fs = unlocked();
    let text = b"Dear diary, nobody else may read this.";
    put(&mut fs, b"users/u1/home/diary.txt", text);
    assert_eq!(get(&mut fs, b"users/u1/home/diary.txt").unwrap(), text);
    let disk = on_disk(&fs, "users/u1/home/diary.txt");
    assert!(disk.starts_with(b"KOSFILE1"));
    assert!(!disk.windows(5).any(|w| w == b"diary"), "the words are not on the disk");
    assert_eq!(disk.len(), filecrypt::disk_len(text.len()));
}

#[test]
fn plaintext_areas_stay_plaintext() {
    let mut fs = unlocked();
    put(&mut fs, b"bin/tool", b"\x7fELF");
    assert_eq!(on_disk(&fs, "bin/tool"), b"\x7fELF");
    assert_eq!(get(&mut fs, b"bin/tool").unwrap(), b"\x7fELF");
}

#[test]
fn big_files_and_writes_across_chunks_round_trip() {
    let mut fs = unlocked();
    let big: Vec<u8> = (0..3 * filecrypt::PLAIN_CHUNK + 123).map(|i| (i % 251) as u8).collect();
    put(&mut fs, b"users/u1/data/kos.notes/big", &big);
    assert_eq!(get(&mut fs, b"users/u1/data/kos.notes/big").unwrap(), big);
    // Overwrite a stretch spanning two chunks in the middle.
    let h = fs.open(b"users/u1/data/kos.notes/big").unwrap();
    let at = filecrypt::PLAIN_CHUNK - 10;
    assert!(fs.seek(h, at));
    assert_eq!(fs.write(h, &[0xAA; 30]), Some(30));
    fs.close(h);
    let mut expect = big.clone();
    expect[at..at + 30].fill(0xAA);
    assert_eq!(get(&mut fs, b"users/u1/data/kos.notes/big").unwrap(), expect);
    // Grow exactly from a chunk boundary: the old last chunk stops being last.
    let mut fs = unlocked();
    let full = vec![7u8; filecrypt::PLAIN_CHUNK];
    put(&mut fs, b"users/u1/home/f", &full);
    let h = fs.open(b"users/u1/home/f").unwrap();
    assert!(fs.seek(h, filecrypt::PLAIN_CHUNK));
    assert_eq!(fs.write(h, b"more"), Some(4));
    fs.close(h);
    let mut expect = full.clone();
    expect.extend_from_slice(b"more");
    assert_eq!(get(&mut fs, b"users/u1/home/f").unwrap(), expect);
}

#[test]
fn one_app_cannot_read_anothers_file_even_moved_into_its_folder() {
    let mut fs = unlocked();
    put(&mut fs, b"users/u1/data/kos.notes/secret", b"notes only");
    let stolen = on_disk(&fs, "users/u1/data/kos.notes/secret");
    fs.inner_mut_for_test().files.insert("users/u1/data/kos.chat/secret".to_string(), stolen);
    let h = fs.open(b"users/u1/data/kos.chat/secret").expect("it looks like a sealed file");
    let mut out = [0u8; 64];
    assert_eq!(fs.read(h, &mut out), None, "but it opens only under kos.notes' key");
}

#[test]
fn locked_the_areas_refuse_everything() {
    let mut fs = unlocked();
    put(&mut fs, b"users/u1/home/x", b"y");
    let mut locked = Encrypted::new(fs.into_inner());
    assert!(locked.open(b"users/u1/home/x").is_none());
    assert!(locked.create(b"users/u1/home/z").is_none());
    assert!(!locked.delete(b"users/u1/home/x"));
    assert!(!locked.mkdir(b"users/u1/data/kos.notes"));
    assert!(locked.list_dir(b"users/u1/home", &mut [DirEntry::empty(); 4]).is_none());
    assert!(locked.open(b"bin/anything").is_none(), "(absent) but plaintext areas are not refused");
    // Another master key (a wiped and re-set-up device) reads noise.
    let mut other = Encrypted::new(locked.into_inner());
    other.unlock(&[1; 32], &[2; 32], b"u1");
    let h = other.open(b"users/u1/home/x").unwrap();
    assert_eq!(other.read(h, &mut [0u8; 8]), None);
}

#[test]
fn plaintext_planted_in_an_encrypted_area_is_refused() {
    let mut fs = unlocked();
    fs.inner_mut_for_test().files.insert("users/u1/permissions/kos.evil".to_string(), b"[decision]\ngranted = [\"network\"]\n".to_vec());
    assert!(fs.open(b"users/u1/permissions/kos.evil").is_none());
}

#[test]
fn tampering_and_truncation_are_caught() {
    let mut fs = unlocked();
    let data = vec![5u8; 2 * filecrypt::PLAIN_CHUNK + 50];
    put(&mut fs, b"users/u1/home/t", &data);
    let good = on_disk(&fs, "users/u1/home/t");
    // One flipped byte.
    let mut bent = good.clone();
    bent[100] ^= 1;
    fs.inner_mut_for_test().files.insert("users/u1/home/t".to_string(), bent);
    assert_eq!(get(&mut fs, b"users/u1/home/t"), None);
    // Cut at a chunk boundary: the remaining last chunk was not sealed as last.
    let cut = good[..filecrypt::chunk_offset(2)].to_vec();
    fs.inner_mut_for_test().files.insert("users/u1/home/t".to_string(), cut);
    assert_eq!(get(&mut fs, b"users/u1/home/t"), None);
    // Chunks swapped.
    let (a, b) = (filecrypt::chunk_offset(0), filecrypt::chunk_offset(1));
    let mut swapped = good.clone();
    swapped[a..b].copy_from_slice(&good[b..b + filecrypt::CHUNK]);
    swapped[b..b + filecrypt::CHUNK].copy_from_slice(&good[a..b]);
    fs.inner_mut_for_test().files.insert("users/u1/home/t".to_string(), swapped);
    assert_eq!(get(&mut fs, b"users/u1/home/t"), None);
}

#[test]
fn listings_show_what_files_hold() {
    let mut fs = unlocked();
    put(&mut fs, b"users/u1/home/a", b"12345");
    let mut entries = [DirEntry::empty(); 4];
    let n = fs.list_dir(b"users/u1/home", &mut entries).unwrap();
    assert_eq!((n, entries[0].size), (1, 5));
}

#[test]
fn files_cannot_live_loose_in_a_users_folder() {
    let mut fs = unlocked();
    assert!(fs.create(b"users/u1/data/loose").is_none());
    assert!(fs.create(b"users/u1/loose").is_none());
    assert!(fs.create(b"users/loose").is_none());
    assert!(fs.mkdir(b"users/u1/data/kos.notes"), "an app's folder is fine");
    let mut entries = [DirEntry::empty(); 2];
    put(&mut fs, b"users/u1/data/kos.notes/a", b"xyz");
    let n = fs.list_dir(b"users/u1/data/kos.notes", &mut entries).unwrap();
    assert_eq!((n, entries[0].size), (1, 3), "listing the folder shows what its files hold");
}

#[test]
fn only_the_signed_in_users_folder_opens() {
    let mut fs = unlocked();
    put(&mut fs, b"users/u1/home/mine", b"u1's");
    let mut inner = fs.into_inner();
    let theirs = inner.files["users/u1/home/mine"].clone();
    let mut u2 = Encrypted::new(std::mem::take(&mut inner));
    u2.unlock(&MASTER, &[3; 32], b"u2");
    assert!(u2.open(b"users/u1/home/mine").is_none(), "not someone else's, even under the same key");
    assert!(u2.create(b"users/u1/home/planted").is_none());
    assert!(!u2.mkdir(b"users/u1/home/sub"));
    assert!(u2.list_dir(b"users/u1/home", &mut [DirEntry::empty(); 4]).is_none());
    assert!(!u2.delete(b"users/u1/home/mine"));
    // Moved into u2's own home, it is u1's area key that sealed it.
    u2.inner_mut_for_test().files.insert("users/u2/home/mine".to_string(), theirs);
    let h = u2.open(b"users/u2/home/mine").unwrap();
    assert_eq!(u2.read(h, &mut [0u8; 8]), None);
    assert!(u2.list_dir(b"users", &mut [DirEntry::empty(); 4]).is_some(), "the folders themselves are no secret");
}

#[test]
fn locking_forgets_the_key_and_closes_sealed_files() {
    let mut fs = unlocked();
    put(&mut fs, b"users/u1/home/a", b"abc");
    let h = fs.open(b"users/u1/home/a").unwrap();
    fs.lock();
    assert_eq!(fs.read(h, &mut [0u8; 4]), None, "the open file went with the key");
    assert!(fs.open(b"users/u1/home/a").is_none());
    fs.unlock(&MASTER, &[3; 32], b"u1");
    assert_eq!(get(&mut fs, b"users/u1/home/a").unwrap(), b"abc", "and comes back with it");
}

#[test]
fn purging_removes_a_tree_without_its_key() {
    let mut fs = unlocked();
    put(&mut fs, b"users/u1/home/a", b"1");
    put(&mut fs, b"users/u1/data/kos.notes/b", b"2");
    assert!(fs.mkdir(b"users/u1/home/empty"));
    fs.lock();
    fs.inner_mut_for_test().files.insert("users/u2x/home/c".to_string(), b"3".to_vec());
    assert!(fs.purge(b"users/u1"));
    let left: Vec<&String> = fs.inner().files.keys().collect();
    assert_eq!(left, ["users/u2x/home/c"]);
    assert!(!fs.inner().dirs.iter().any(|d| d.starts_with("users/u1")), "folders too");
    assert!(fs.purge(b"users/nobody"), "nothing there is nothing to remove");
}