use super::*;
use alloc::collections::BTreeMap;

#[derive(Default, Clone)]
struct Image {
    len: u64,
    pages: BTreeMap<u64, Vec<u8>>,
}
impl Storage for Image {
    fn len(&mut self) -> Result<u64> {
        Ok(self.len)
    }
    fn resize(&mut self, n: u64) -> Result<()> {
        self.len = n;
        Ok(())
    }
    fn read(&mut self, p: u64, b: &mut [u8]) -> Result<()> {
        if p + b.len() as u64 > self.len {
            return Err(Error::Io);
        }
        if let Some(bytes) = self.pages.get(&p) {
            b.copy_from_slice(&bytes[..b.len()]);
        } else {
            b.fill(0);
        }
        Ok(())
    }
    fn write(&mut self, p: u64, b: &[u8]) -> Result<()> {
        if p + b.len() as u64 > self.len {
            return Err(Error::Io);
        }
        self.pages.insert(p, b.to_vec());
        Ok(())
    }
    fn sync(&mut self) -> Result<()> {
        Ok(())
    }
}

#[derive(Clone, Copy, Debug)]
enum Tear {
    Prefix(usize),
    Scatter,
}
impl Tear {
    fn keep(self, i: usize) -> bool {
        match self {
            Self::Prefix(n) => i < n,
            Self::Scatter => i % 3 != 1,
        }
    }
}
struct Fault {
    working: Image,
    durable: Image,
    calls: usize,
    fail: usize,
    tear: Tear,
}
impl Fault {
    fn new(base: &Image, fail: usize, tear: Tear) -> Self {
        Self {
            working: base.clone(),
            durable: base.clone(),
            calls: 0,
            fail,
            tear,
        }
    }
    fn tick(&mut self) -> bool {
        self.calls += 1;
        self.calls == self.fail
    }
    fn spill(&mut self, offset: u64) {
        let source = &self.working.pages[&offset];
        let target = self
            .durable
            .pages
            .entry(offset)
            .or_insert_with(|| vec![0; source.len()]);
        for (i, byte) in source.iter().enumerate() {
            if self.tear.keep(i) {
                target[i] = *byte;
            }
        }
    }
}
impl Storage for Fault {
    fn len(&mut self) -> Result<u64> {
        self.working.len()
    }
    fn read(&mut self, p: u64, b: &mut [u8]) -> Result<()> {
        self.working.read(p, b)
    }
    fn resize(&mut self, n: u64) -> Result<()> {
        self.working.resize(n)?;
        // Expose length immediately, as a filesystem can on failure.
        self.durable.resize(n)?;
        if self.tick() {
            return Err(Error::Io);
        }
        Ok(())
    }
    fn write(&mut self, p: u64, b: &[u8]) -> Result<()> {
        self.working.write(p, b)?;
        self.spill(p); // Early persistence can include any modeled tear, even before a sync.
        if self.tick() {
            return Err(Error::Io);
        }
        Ok(())
    }
    fn sync(&mut self) -> Result<()> {
        if self.tick() {
            for p in self.working.pages.keys().copied().collect::<Vec<_>>() {
                self.spill(p);
            }
            return Err(Error::Io);
        }
        self.durable = self.working.clone();
        Ok(())
    }
}

const OLD: [Change; 3] = [
    Change {
        key: 1,
        value: Some([1; 32]),
    },
    Change {
        key: 250,
        value: Some([2; 32]),
    },
    Change {
        key: CAPACITY - 1,
        value: Some([9; 32]),
    },
];
const NEW: [Change; 3] = [
    Change {
        key: 1,
        value: Some([3; 32]),
    },
    Change {
        key: 2,
        value: Some([4; 32]),
    },
    Change {
        key: 250,
        value: None,
    },
];
fn before() -> Image {
    let mut db = PageIndex::create(Image::default()).unwrap();
    db.apply(&OLD).unwrap();
    db.into_storage()
}
fn model<S: Storage>(db: &mut PageIndex<S>) -> Vec<(u32, Value)> {
    let mut rows = Vec::new();
    db.scan(|k, v| {
        rows.push((k, v));
        true
    })
    .unwrap();
    rows
}

#[test]
fn every_commit_boundary_and_torn_locator_prefix_is_old_or_new() {
    let base = before();
    let mut db = PageIndex::open(base.clone()).unwrap();
    let old = model(&mut db);
    db.apply(&NEW).unwrap();
    let new = model(&mut db);
    for tear in (0..=32).map(Tear::Prefix).chain([
        Tear::Prefix(1024),
        Tear::Prefix(4095),
        Tear::Prefix(PAGE),
        Tear::Scatter,
    ]) {
        for fail in 1.. {
            let mut candidate = PageIndex::open(Fault::new(&base, fail, tear)).unwrap();
            let ack = candidate.apply(&NEW).is_ok();
            if !ack {
                assert_eq!(candidate.get(1), Err(Error::Poisoned));
            }
            let device = candidate.into_storage();
            let mut recovered = PageIndex::open(device.durable.clone()).unwrap();
            let rows = model(&mut recovered);
            assert!(
                rows == new || (!ack && rows == old),
                "{tear:?}, call {fail}"
            );
            if fail <= 4 {
                assert_eq!(rows, old, "before publication: {tear:?}, call {fail}");
            }
            if fail >= 7 {
                assert_eq!(
                    rows, new,
                    "after first durable record: {tear:?}, call {fail}"
                );
            }
            assert_eq!(recovered.generation(), if rows == new { 3 } else { 2 });
            if ack {
                assert_eq!(fail, device.calls + 1);
                break;
            }
        }
    }
}

#[test]
fn repeated_interruptions_repair_the_only_valid_witness_first() {
    let mut db = PageIndex::open(before()).unwrap();
    db.apply(&NEW).unwrap();
    let base = db.into_storage();
    for damaged in [0, PAGE as u64] {
        let mut partial = base.clone();
        partial.pages.get_mut(&damaged).unwrap().fill(0);
        let mut db = PageIndex::open(partial.clone()).unwrap();
        let old = model(&mut db);
        let next = [Change {
            key: 1,
            value: Some([8; 32]),
        }];
        db.apply(&next).unwrap();
        let new = model(&mut db);
        for fail in 1.. {
            let mut db = PageIndex::open(Fault::new(&partial, fail, Tear::Scatter)).unwrap();
            let ack = db.apply(&next).is_ok();
            let image = db.into_storage().durable;
            let mut db = PageIndex::open(image.clone()).unwrap();
            let first = model(&mut db);
            assert!(first == new || (!ack && first == old));
            // Another complete round of interrupted writes begins at each recovered boundary.
            for cut in 1.. {
                let mut db = PageIndex::open(Fault::new(&image, cut, Tear::Prefix(15))).unwrap();
                let finished = db.apply(&next).is_ok();
                let mut db = PageIndex::open(db.into_storage().durable).unwrap();
                let second = model(&mut db);
                assert!(second == new || (!finished && second == first));
                if finished {
                    break;
                }
            }
            if ack {
                break;
            }
        }
    }
}

#[test]
fn initialization_and_corrupt_or_unrelated_media() {
    assert_eq!(crc(b"123456789"), 0xcbf43926);
    let empty = Image::default();
    // Initializer owns storage, so retain durable state using a borrowed adapter.
    struct Borrowed<'a>(&'a mut Fault);
    impl Storage for Borrowed<'_> {
        fn len(&mut self) -> Result<u64> {
            self.0.len()
        }
        fn resize(&mut self, n: u64) -> Result<()> {
            self.0.resize(n)
        }
        fn read(&mut self, p: u64, b: &mut [u8]) -> Result<()> {
            self.0.read(p, b)
        }
        fn write(&mut self, p: u64, b: &[u8]) -> Result<()> {
            self.0.write(p, b)
        }
        fn sync(&mut self) -> Result<()> {
            self.0.sync()
        }
    }
    for tear in [
        Tear::Prefix(0),
        Tear::Prefix(16),
        Tear::Prefix(PAGE),
        Tear::Scatter,
    ] {
        for fail in 1.. {
            let mut fault = Fault::new(&empty, fail, tear);
            let ack = PageIndex::create(Borrowed(&mut fault)).is_ok();
            match PageIndex::open(fault.durable) {
                Ok(mut db) => assert!(model(&mut db).is_empty()),
                Err(Error::Ambiguous) if !ack => (),
                _ => panic!("bad initial image after {fail}"),
            }
            if ack {
                break;
            }
        }
    }

    let base = before();
    for p in [root_offset(2), leaf_offset(0, 1)] {
        let mut broken = base.clone();
        broken.pages.get_mut(&p).unwrap()[30] ^= 0x80;
        assert!(matches!(PageIndex::open(broken), Err(Error::Corrupt)));
    }
    let mut erased = base.clone();
    erased.pages.remove(&0);
    erased.pages.remove(&(PAGE as u64));
    assert!(matches!(PageIndex::open(erased), Err(Error::Ambiguous)));
    let mut short = base.clone();
    short.len -= 1;
    assert!(matches!(PageIndex::open(short), Err(Error::Invalid)));
    assert!(matches!(PageIndex::create(base), Err(Error::Invalid)));

    let legacy = ratdb_core::Database::open(
        ratdb_core::MemoryStorage::default(),
        ratdb_core::Config {
            bank_size: FILE_BYTES / 2,
            ..Default::default()
        },
    )
    .unwrap()
    .into_storage();
    assert!(PageIndex::open(legacy.clone()).is_err());
    assert!(PageIndex::create(legacy).is_err());
}

#[test]
fn limits_and_ordered_model_remain_valid_on_rejection() {
    let mut db = PageIndex::open(before()).unwrap();
    let old = model(&mut db);
    for changes in [
        vec![OLD[0]; 1025],
        vec![OLD[1], OLD[0]],
        vec![OLD[0]; 2],
        vec![Change {
            key: CAPACITY,
            value: Some([0; 32]),
        }],
    ] {
        assert_eq!(db.apply(&changes), Err(Error::Invalid));
        assert_eq!(db.generation(), 2);
        assert_eq!(model(&mut db), old);
    }
    assert_eq!(db.get(CAPACITY), Err(Error::Invalid));
    let mut rng = 7u32;
    let mut expected: BTreeMap<_, _> = old.into_iter().collect();
    for i in 0..1200 {
        rng = rng.wrapping_mul(1664525).wrapping_add(1013904223);
        let key = rng % 901;
        let value = (!rng.is_multiple_of(7)).then_some([(i % 251) as u8; 32]);
        db.apply(&[Change { key, value }]).unwrap();
        if let Some(v) = value {
            expected.insert(key, v);
        } else {
            expected.remove(&key);
        }
        if i % 100 == 0 {
            db = PageIndex::open(db.into_storage()).unwrap();
        }
        assert_eq!(db.get(key).unwrap(), value);
    }
    assert_eq!(model(&mut db), expected.into_iter().collect::<Vec<_>>());
}

#[test]
fn valid_checksums_cannot_substitute_a_stale_root_or_leaf_or_conflicting_locator() {
    let first = before();
    let mut db = PageIndex::open(first.clone()).unwrap();
    db.apply(&NEW).unwrap();
    let current = db.into_storage();
    let mut wrong_root = current.clone();
    wrong_root
        .pages
        .insert(root_offset(3), first.pages[&root_offset(2)].clone());
    assert!(matches!(PageIndex::open(wrong_root), Err(Error::Corrupt)));
    let mut wrong_leaf = current.clone();
    wrong_leaf
        .pages
        .insert(leaf_offset(0, 2), first.pages[&leaf_offset(0, 1)].clone());
    assert!(matches!(PageIndex::open(wrong_leaf), Err(Error::Corrupt)));
    let mut conflict = current.clone();
    conflict.pages.insert(0, first.pages[&0].clone()); // reversed pair, publication order impossible
    assert!(matches!(PageIndex::open(conflict), Err(Error::Ambiguous)));
    let mut conflict = current;
    let primary = conflict.pages.get_mut(&0).unwrap();
    primary[16] ^= 1;
    seal(primary); // same generation, incompatible roots
    assert!(matches!(PageIndex::open(conflict), Err(Error::Ambiguous)));
}

#[test]
fn adapter_read_errors_never_become_empty_pages_or_stale_fallback() {
    struct ReadFail(Image, Option<u64>);
    impl Storage for ReadFail {
        fn len(&mut self) -> Result<u64> {
            self.0.len()
        }
        fn resize(&mut self, n: u64) -> Result<()> {
            self.0.resize(n)
        }
        fn read(&mut self, p: u64, b: &mut [u8]) -> Result<()> {
            if self.1 == Some(p) {
                return Err(Error::Io);
            }
            self.0.read(p, b)
        }
        fn write(&mut self, p: u64, b: &[u8]) -> Result<()> {
            self.0.write(p, b)
        }
        fn sync(&mut self) -> Result<()> {
            self.0.sync()
        }
    }
    let first = before();
    for off in [0, PAGE as u64, root_offset(2), leaf_offset(2, 1)] {
        assert!(matches!(
            PageIndex::open(ReadFail(first.clone(), Some(off))),
            Err(Error::Io)
        ));
    }
    let mut db = PageIndex::open(ReadFail(first, None)).unwrap();
    let old = model(&mut db);
    db.storage.1 = Some(leaf_offset(2, 1)); // previous group already wrote a candidate leaf
    assert_eq!(db.apply(&NEW), Err(Error::Io));
    assert_eq!(db.get(2), Err(Error::Poisoned));
    let mut reopened = PageIndex::open(db.into_storage().0).unwrap();
    assert_eq!(model(&mut reopened), old);
}