#![no_main]

use arbitrary::Arbitrary;
use commonware_cryptography::Sha256;
use commonware_parallel::{Rayon, Sequential, Strategy};
use commonware_runtime::{
    Runner, Strategizer as _, Supervisor as _, buffer::paged::CacheRef, deterministic,
};
use commonware_storage::{
    journal::contiguous::variable::Config as VConfig,
    merkle::{Family, Location, full::Config as MerkleConfig, mmb, mmr},
    qmdb::{
        Error,
        keyless::variable::{Config, Db as Keyless},
        verify_proof,
    },
};
use commonware_utils::{NZU16, NZU64, NZUsize};
use libfuzzer_sys::fuzz_target;
use std::num::NonZeroU16;

const MAX_OPERATIONS: usize = 50;
const MAX_PROOF_OPS: u64 = 100;
type CodecConfig = (commonware_codec::RangeCfg<usize>, ());

/// Which error variant a bad-floor commit should produce.
#[derive(Debug, Clone, Copy)]
enum BadFloorExpect {
    Regression,
    BeyondSize,
}

fn assert_bad_floor_error<F: Family>(err: &Error<F>, kind: BadFloorExpect) {
    match (err, kind) {
        (Error::FloorRegressed(_, _), BadFloorExpect::Regression) => {}
        (Error::FloorBeyondSize(_, _), BadFloorExpect::BeyondSize) => {}
        _ => panic!("unexpected error for {kind:?}: {err:?}"),
    }
}

/// What floor value a fuzz-generated commit should carry. Merkleize must reject the floors that
/// the `Bad*` variants produce, and the handler asserts the expected error variant.
#[derive(Debug, Clone, Copy)]
enum FloorKind {
    /// Keep the current floor (monotonicity trivially preserved).
    Current,
    /// Advance to the commit location (the tight upper bound).
    AdvanceToCommit,
    /// Floor one below the current floor. It must be rejected as `FloorRegressed`.
    BadRegression,
    /// Floor one past the commit location. It must be rejected as `FloorBeyondSize`.
    BadBeyondCommit,
}

impl<'a> Arbitrary<'a> for FloorKind {
    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
        let choice: u8 = u.arbitrary()?;
        Ok(match choice % 4 {
            0 => FloorKind::Current,
            1 => FloorKind::AdvanceToCommit,
            2 => FloorKind::BadRegression,
            3 => FloorKind::BadBeyondCommit,
            _ => unreachable!(),
        })
    }
}

#[derive(Debug)]
enum Operation {
    Append {
        value_bytes: Vec<u8>,
    },
    Commit {
        metadata_bytes: Option<Vec<u8>>,
        floor_kind: FloorKind,
    },
    /// Build a two-level batch chain (parent -> child) with a valid parent and a child floor that
    /// regresses below the parent's floor or passes the child's own commit location. Merkleizing
    /// the child must return `FloorRegressed` or `FloorBeyondSize`.
    BadChainedCommit {
        child_kind: FloorKind,
    },
    Get {
        loc_offset: u32,
    },
    GetMetadata,
    Prune,
    Sync,
    OpCount,
    OldestRetainedLoc,
    Root,
    Proof {
        start_offset: u32,
        max_ops: u16,
    },
    HistoricalProof {
        size_offset: u32,
        start_offset: u32,
        max_ops: u16,
    },
    SimulateFailure {},
    ReopenAtMost {
        cap: u64,
    },
}

impl<'a> Arbitrary<'a> for Operation {
    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
        let choice: u8 = u.arbitrary()?;
        match choice % 14 {
            0 => {
                let value_len: u16 = u.arbitrary()?;
                let actual_len = ((value_len as usize) % 10000) + 1;
                let value_bytes = u.bytes(actual_len)?.to_vec();
                Ok(Operation::Append { value_bytes })
            }
            1 => {
                let has_metadata: bool = u.arbitrary()?;
                let metadata_bytes = if has_metadata {
                    let metadata_len: u16 = u.arbitrary()?;
                    let actual_len = ((metadata_len as usize) % 1000) + 1;
                    Some(u.bytes(actual_len)?.to_vec())
                } else {
                    None
                };
                let floor_kind = FloorKind::arbitrary(u)?;
                Ok(Operation::Commit {
                    metadata_bytes,
                    floor_kind,
                })
            }
            2 => {
                let loc_offset = u.arbitrary()?;
                Ok(Operation::Get { loc_offset })
            }
            3 => Ok(Operation::GetMetadata),
            4 => Ok(Operation::Prune),
            5 => Ok(Operation::Sync),
            6 => Ok(Operation::OpCount),
            7 => Ok(Operation::OldestRetainedLoc),
            8 => Ok(Operation::Root),
            9 => {
                let start_offset = u.arbitrary()?;
                let max_ops = u.arbitrary()?;
                Ok(Operation::Proof {
                    start_offset,
                    max_ops,
                })
            }
            10 => {
                let size_offset = u.arbitrary()?;
                let start_offset = u.arbitrary()?;
                let max_ops = u.arbitrary()?;
                Ok(Operation::HistoricalProof {
                    size_offset,
                    start_offset,
                    max_ops,
                })
            }
            11 => Ok(Operation::SimulateFailure {}),
            13 => Ok(Operation::ReopenAtMost {
                cap: u.arbitrary()?,
            }),
            12 => {
                // Only Bad* kinds make sense here.
                let child_kind = match u.arbitrary::<bool>()? {
                    false => FloorKind::BadRegression,
                    true => FloorKind::BadBeyondCommit,
                };
                Ok(Operation::BadChainedCommit { child_kind })
            }
            _ => unreachable!(),
        }
    }
}

#[derive(Debug)]
struct FuzzInput {
    ops: Vec<Operation>,
}

impl<'a> Arbitrary<'a> for FuzzInput {
    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
        let num_ops = u.int_in_range(1..=MAX_OPERATIONS)?;
        let ops = (0..num_ops)
            .map(|_| Operation::arbitrary(u))
            .collect::<Result<Vec<_>, _>>()?;
        Ok(FuzzInput { ops })
    }
}

const PAGE_SIZE: NonZeroU16 = NZU16!(127);
const PAGE_CACHE_SIZE: usize = 8;

type Db<F, S> = Keyless<F, deterministic::Context, Vec<u8>, Sha256, S>;

fn test_config<S: Strategy>(test_name: &str, strategy: S) -> Config<CodecConfig, S> {
    let page_cache = CacheRef::new(PAGE_SIZE, NZUsize!(PAGE_CACHE_SIZE));
    Config {
        merkle: MerkleConfig {
            journal_partition: format!("{test_name}-journal"),
            metadata_partition: format!("{test_name}-meta"),
            items_per_blob: NZU64!(3),
            write_buffer: NZUsize!(1024),
            replay_buffer: NZUsize!(1024),
            strategy,
            page_cache: page_cache.clone(),
        },
        log: VConfig {
            partition: format!("{test_name}-log"),
            write_buffer: NZUsize!(1024),
            replay_buffer: NZUsize!(1024),
            compression: None,
            codec_config: ((0..=10000).into(), ()),
            items_per_section: NZU64!(7),
            page_cache,
        },
    }
}

/// Reopen the database.
async fn reopen<F: Family, S: Strategy>(
    context: &deterministic::Context,
    suffix: &str,
    strategy: &S,
    restarts: &mut usize,
) -> Db<F, S> {
    let cfg = test_config(suffix, strategy.clone());
    let db = Db::init(
        context.child("db").with_attribute("instance", *restarts),
        cfg,
        None,
    )
    .await
    .expect("Failed to init keyless db");
    *restarts += 1;
    db
}

fn fuzz_family<F: Family, S: Strategy>(
    input: &FuzzInput,
    suffix: &str,
    strategy: impl FnOnce(&deterministic::Context) -> S,
) {
    let runner = deterministic::Runner::default();

    runner.start(|context| async move {
        let strategy = strategy(&context);
        let cfg = test_config(suffix, strategy.clone());
        let mut db: Db<F, S> = Db::init(context.child("storage"), cfg, None)
            .await
            .expect("Failed to init keyless db");
        let mut restarts = 0usize;

        let mut pending_appends: Vec<Vec<u8>> = Vec::new();
        let mut commits = std::collections::BTreeMap::new();

        for op in &input.ops {
            commits.insert(db.bounds().end, (db.root(), db.inactivity_floor_loc()));
            db = match op {
                Operation::Append { value_bytes } => {
                    pending_appends.push(value_bytes.clone());
                    db
                }

                Operation::Commit { metadata_bytes, floor_kind } => {
                    let pending_count = pending_appends.len() as u64;
                    let end = db.bounds().end;
                    let commit_loc = end.as_u64() + pending_count;
                    let current_floor = db.inactivity_floor_loc();

                    // Pick the floor for this commit. `Bad*` kinds are guaranteed to trigger
                    // the expected error; Valid kinds (Current/AdvanceToCommit) apply cleanly.
                    let (floor, expect_err) = match floor_kind {
                        FloorKind::Current => (current_floor, None),
                        FloorKind::AdvanceToCommit => (Location::<F>::new(commit_loc), None),
                        FloorKind::BadRegression => {
                            // Only meaningful when current floor > 0; otherwise fall back to Current.
                            if current_floor.as_u64() == 0 {
                                (current_floor, None)
                            } else {
                                let bad = Location::<F>::new(current_floor.as_u64() - 1);
                                (bad, Some(BadFloorExpect::Regression))
                            }
                        }
                        FloorKind::BadBeyondCommit => {
                            let bad = Location::<F>::new(commit_loc + 1);
                            (bad, Some(BadFloorExpect::BeyondSize))
                        }
                    };

                    let mut batch = db.new_batch();
                    for v in pending_appends.drain(..) {
                        batch = batch.append(v);
                    }
                    let merkleized = batch.merkleize(&db, metadata_bytes.clone(), floor).await;

                    match expect_err {
                        None => {
                            let merkleized = merkleized.expect("valid floor must merkleize");
                            let db = db
                                .apply_batch(merkleized)
                                .await
                                .expect("batch with a valid floor must apply");
                            db.commit().await.expect("Commit should not fail")
                        }
                        Some(kind) => {
                            // Merkleize validates the floor, so the batch never reaches apply.
                            let Err(err) = merkleized else {
                                panic!("bad floor must be rejected");
                            };
                            assert_bad_floor_error(&err, kind);
                            db
                        }
                    }
                }

                Operation::BadChainedCommit { child_kind } => {
                    let end = db.bounds().end;

                    // Parent: one append at `end`, commit at `end + 1`, floor at its commit.
                    // Child: one append at `end + 2`, commit at `end + 3`.
                    let parent_commit_loc = end.as_u64() + 1;
                    let child_commit_loc = parent_commit_loc + 2;
                    let (child_floor, kind) = match child_kind {
                        FloorKind::BadRegression => (
                            Location::<F>::new(parent_commit_loc - 1),
                            BadFloorExpect::Regression,
                        ),
                        FloorKind::BadBeyondCommit => (
                            Location::<F>::new(child_commit_loc + 1),
                            BadFloorExpect::BeyondSize,
                        ),
                        _ => continue, // only bad kinds are meaningful here
                    };

                    // Keep pending_appends for later operations.
                    let parent = db
                        .new_batch()
                        .append(vec![0u8; 1])
                        .merkleize(&db, None, Location::<F>::new(parent_commit_loc))
                        .await
                        .expect("valid parent must merkleize");
                    let Err(err) = parent
                        .new_batch::<Sha256>()
                        .append(vec![1u8; 1])
                        .merkleize(&db, None, child_floor)
                        .await
                    else {
                        panic!("bad child floor must be rejected");
                    };
                    assert_bad_floor_error(&err, kind);
                    db
                }

                Operation::Get { loc_offset } => {
                    let op_count = db.bounds().end;
                    if op_count > 0 {
                        let loc = (*loc_offset as u64) % op_count.as_u64();
                        let _ = db.get(loc.into()).await;
                    }
                    db
                }

                Operation::GetMetadata => {
                    let _ = db.get_metadata().await;
                    db
                }

                Operation::Prune => {
                    let pending_count = pending_appends.len() as u64;
                    let mut batch = db.new_batch();
                    for v in pending_appends.drain(..) {
                        batch = batch.append(v);
                    }
                    // Advance the floor to the new commit location so the subsequent prune
                    // actually removes data. This exercises more of the code path than pruning
                    // at a stale floor would.
                    let end = db.bounds().end;
                    let floor = Location::<F>::new(end.as_u64() + pending_count);
                    let merkleized = batch.merkleize(&db, None, floor).await.unwrap();
                    let db = db.apply_batch(merkleized).await.expect("Commit should not fail");
                    let db = db.commit().await.expect("Commit should not fail");
                    let floor = db.inactivity_floor_loc();
                    db.prune(floor).await.expect("Prune should not fail")
                }

                Operation::Sync => {
                    let mut batch = db.new_batch();
                    for v in pending_appends.drain(..) {
                        batch = batch.append(v);
                    }
                    let merkleized = batch.merkleize(&db, None, db.inactivity_floor_loc()).await.unwrap();
                    let db = db.apply_batch(merkleized).await.expect("Commit should not fail");
                    db.sync().await.expect("Sync should not fail")
                }

                Operation::OpCount => {
                    let _ = db.bounds().end;
                    db
                }

                Operation::OldestRetainedLoc => {
                    let _ = db.bounds().start;
                    db
                }

                Operation::Root => {
                    let mut batch = db.new_batch();
                    for v in pending_appends.drain(..) {
                        batch = batch.append(v);
                    }
                    let merkleized = batch.merkleize(&db, None, db.inactivity_floor_loc()).await.unwrap();
                    let db = db.apply_batch(merkleized).await.expect("Commit should not fail");
                    let db = db.commit().await.expect("Commit should not fail");
                    let _ = db.root();
                    db
                }

                Operation::Proof {
                    start_offset,
                    max_ops,
                } => {
                    let op_count = db.bounds().end;
                    if op_count == 0 {
                        continue;
                    }
                    let mut batch = db.new_batch();
                    for v in pending_appends.drain(..) {
                        batch = batch.append(v);
                    }
                    let merkleized = batch.merkleize(&db, None, db.inactivity_floor_loc()).await.unwrap();
                    let db = db.apply_batch(merkleized).await.expect("Commit should not fail");
                    let db = db.commit().await.expect("Commit should not fail");
                    let start_loc = (*start_offset as u64) % op_count.as_u64();
                    let max_ops_value = ((*max_ops as u64) % MAX_PROOF_OPS) + 1;
                    let start_loc: Location<F> = Location::new(start_loc);
                    let root = db.root();
                    if let Ok((proof, ops)) = db.proof(start_loc, NZU64!(max_ops_value)).await {
                            assert!(
                                verify_proof::<Sha256, _, _>(
                                    &proof,
                                    start_loc,
                                    &ops,
                                    &root),
                                "Failed to verify proof for start loc{start_loc} with ops {max_ops} ops",
                            );
                    }
                    db
                }

                Operation::HistoricalProof {
                    size_offset,
                    start_offset,
                    max_ops,
                } => {
                    let op_count = db.bounds().end;
                    if op_count == 0 {
                        continue;
                    }
                    let mut batch = db.new_batch();
                    for v in pending_appends.drain(..) {
                        batch = batch.append(v);
                    }
                    let merkleized = batch.merkleize(&db, None, db.inactivity_floor_loc()).await.unwrap();
                    let db = db.apply_batch(merkleized).await.expect("Commit should not fail");
                    let db = db.commit().await.expect("Commit should not fail");
                    // Use post-commit op_count so it's consistent with the root.
                    let op_count = db.bounds().end;
                    let size = ((*size_offset as u64) % op_count.as_u64()) + 1;
                    let size: Location<F> = Location::new(size);
                    let start_loc = (*start_offset as u64) % *size;
                    let start_loc: Location<F> = Location::new(start_loc);
                    let max_ops_value = ((*max_ops as u64) % MAX_PROOF_OPS) + 1;
                    let root = db.root();
                    if let Ok((proof, ops)) = db
                        .historical_proof(op_count, start_loc, NZU64!(max_ops_value))
                            .await {
                            assert!(
                                verify_proof::<Sha256, _, _>(
                                    &proof,
                                    start_loc,
                                    &ops,
                                    &root),
                                "Failed to verify historical proof for start loc{start_loc} with max ops {max_ops}",
                            );
                        }
                    db
                }

                Operation::ReopenAtMost { cap } => {
                    pending_appends.clear();
                    let before_bounds = db.bounds();
                    let before_root = db.root();
                    let cap = Location::<F>::new(*cap % (*before_bounds.end + 3));
                    let selected = commits.range(..=cap).next_back().map(|(end, state)| (*end, *state));
                    _ = db.sync().await.expect("sync before cap");
                    let opened = Db::<F, S>::init(
                        context.child("capped").with_attribute("instance", restarts),
                        test_config(suffix, strategy.clone()), Some(cap),
                    ).await;
                    restarts += 1;
                    match opened {
                        Ok(opened) => {
                            let (end, (root, floor)) = selected.expect("cap must select a modeled commit");
                            assert_eq!(opened.bounds().end, end);
                            assert_eq!(opened.root(), root);
                            assert_eq!(opened.inactivity_floor_loc(), floor);
                            drop(opened);
                            let reopened = reopen(&context, suffix, &strategy, &mut restarts).await;
                            assert_eq!(reopened.bounds().end, end);
                            assert_eq!(reopened.root(), root);
                            commits.retain(|candidate, _| *candidate <= end);
                            reopened
                        }
                        Err(error @ (Error::InvalidInitializationBound
                            | Error::Journal(commonware_storage::journal::Error::ItemPruned(_)))) => {
                            match error {
                                Error::InvalidInitializationBound => assert_eq!(cap, 0),
                                Error::Journal(_) => {
                                    assert_ne!(cap, 0);
                                    assert!(selected.is_none_or(|(end, _)| end <= before_bounds.start));
                                }
                                _ => unreachable!(),
                            }
                            let reopened = reopen(&context, suffix, &strategy, &mut restarts).await;
                            assert_eq!(reopened.root(), before_root);
                            reopened
                        }
                        Err(err) => panic!("unexpected bounded initialization error. {err:?}"),
                    }
                }

                Operation::SimulateFailure{} => {
                    pending_appends.clear();
                    drop(db);

                    reopen(&context, suffix, &strategy, &mut restarts).await
                }
            };
        }

        let mut batch = db.new_batch();
        for v in pending_appends.drain(..) {
            batch = batch.append(v);
        }
        let merkleized = batch.merkleize(&db, None, db.inactivity_floor_loc()).await.unwrap();
        let db = db
            .apply_batch(merkleized)
            .await
            .expect("Commit should not fail");
        db.destroy().await.expect("Destroy should not fail");
    });
}

fuzz_target!(|input: FuzzInput| {
    fuzz_family::<mmr::Family, Sequential>(&input, "fuzz-mmr-sequential", |_| Sequential);
    fuzz_family::<mmb::Family, Sequential>(&input, "fuzz-mmb-sequential", |_| Sequential);
    fuzz_family::<mmr::Family, Rayon>(&input, "fuzz-mmr-rayon", |context| {
        context.strategy(NZUsize!(2))
    });
    fuzz_family::<mmb::Family, Rayon>(&input, "fuzz-mmb-rayon", |context| {
        context.strategy(NZUsize!(2))
    });
});