#![no_main]

use arbitrary::Arbitrary;
use commonware_cryptography::{Sha256, sha256::Digest};
use commonware_parallel::Sequential;
use commonware_runtime::{Runner, Supervisor as _, buffer::paged::CacheRef, deterministic};
use commonware_storage::{
    index::ordered::Index,
    journal::contiguous::fixed::{Config as FConfig, Journal},
    merkle::{Family as MerkleFamily, Location, Proof, mmb, mmr},
    mmr::full::Config as MerkleConfig,
    qmdb::{
        any::{
            FixedConfig as Config,
            db::Db as AnyDb,
            ordered::{Operation, Update},
            traits::DbAny as _,
            value::FixedEncoding,
        },
        floor::Proportional,
        verify_proof,
    },
    translator::EightCap,
};
use commonware_storage_fuzz::{
    assert_ordered_neighbors,
    floor::{Plan, Recorder},
};
use commonware_utils::{NZU16, NZU64, NZUsize, sequence::FixedBytes};
use libfuzzer_sys::fuzz_target;
use std::{
    collections::{HashMap, HashSet},
    num::{NonZeroU16, NonZeroU64},
};

type Key = FixedBytes<32>;
type Value = FixedBytes<64>;
type RawKey = [u8; 32];
type RawValue = [u8; 64];
const BITMAP_CHUNK_BYTES: usize = 64;

type GenericDb<F> = AnyDb<
    F,
    deterministic::Context,
    Journal<deterministic::Context, Operation<F, Key, FixedEncoding<Value>>>,
    Index<EightCap, Location<F>>,
    Sha256,
    Update<Key, FixedEncoding<Value>>,
    BITMAP_CHUNK_BYTES,
    Sequential,
>;

const MAX_OPS: usize = 25;

#[derive(Arbitrary, Debug, Clone)]
enum QmdbOperation {
    Update {
        key: RawKey,
        value: RawValue,
    },
    Delete {
        key: RawKey,
    },
    Commit(Plan),
    OpCount,
    Root,
    Proof {
        start_loc: u64,
        max_ops: NonZeroU64,
    },
    ArbitraryProof {
        start_loc: u64,
        max_ops: NonZeroU64,
        proof_leaves: u64,
        digests: Vec<[u8; 32]>,
    },
    Get {
        key: RawKey,
    },
    GetSpan {
        key: RawKey,
    },
    GetNeighbors {
        key: RawKey,
    },
}

#[derive(Arbitrary, Debug)]
struct FuzzInput {
    operations: Vec<QmdbOperation>,
}

const PAGE_SIZE: NonZeroU16 = NZU16!(555);
const PAGE_CACHE_SIZE: usize = 100;

async fn commit_pending<F: MerkleFamily>(
    db: GenericDb<F>,
    plan: &Plan,
    pending_writes: &mut Vec<(Key, Option<Value>)>,
    committed_state: &mut HashMap<RawKey, RawValue>,
    pending_inserts: &mut HashMap<RawKey, RawValue>,
    pending_deletes: &mut HashSet<RawKey>,
) -> GenericDb<F> {
    let inherited = db.inactivity_floor_loc();
    let mut batch = db.new_batch();
    for (k, v) in pending_writes.drain(..) {
        batch = batch.write(k, v);
    }
    let mut policy = Recorder::new(plan, |seed| Value::new([seed; 64]));
    let merkleized = batch.merkleize(&db, None, &mut policy).await.unwrap();
    let db = db
        .apply_batch(merkleized)
        .await
        .expect("commit should not fail");
    let db = db.commit().await.expect("commit fsync should not fail");
    for key in pending_deletes.drain() {
        committed_state.remove(&key);
    }
    committed_state.extend(pending_inserts.drain());

    // Check the floor walk against the post-write state, then replay its decisions into it.
    *committed_state = policy.check_raw(
        committed_state.drain(),
        inherited,
        db.inactivity_floor_loc(),
        db.bounds().end - 1,
    );
    db
}

/// Check strict, non-wrapping neighbors against the committed model, excluding queued writes.
async fn assert_neighbors<F: MerkleFamily>(
    db: &GenericDb<F>,
    committed_state: &HashMap<RawKey, RawValue>,
    key: RawKey,
) {
    let query = Key::new(key);
    let prev = db
        .get_prev_key(&query)
        .await
        .expect("get_prev_key should not fail");
    let next = db
        .get_next_key(&query)
        .await
        .expect("get_next_key should not fail");
    assert_ordered_neighbors(
        committed_state.keys().copied().map(Key::new),
        &query,
        prev,
        next,
    );
}

fn fuzz_family<F: MerkleFamily>(data: &FuzzInput, suffix: &str) {
    let runner = deterministic::Runner::default();

    runner.start(|context| {
        let operations = data.operations.clone();
        async move {
            let page_cache = CacheRef::new(PAGE_SIZE, NZUsize!(PAGE_CACHE_SIZE));
            let cfg = Config::<EightCap, Sequential> {
                merkle_config: MerkleConfig {
                    journal_partition: format!("test-qmdb-mmr-journal-{suffix}"),
                    metadata_partition: format!("test-qmdb-mmr-metadata-{suffix}"),
                    items_per_blob: NZU64!(500000),
                    write_buffer: NZUsize!(1024),
                    replay_buffer: NZUsize!(1024),
                    strategy: Sequential,
                    page_cache: page_cache.clone(),
                },
                journal_config: FConfig {
                    partition: format!("test-qmdb-log-journal-{suffix}"),
                    items_per_blob: NZU64!(500000),
                    write_buffer: NZUsize!(1024),
                    replay_buffer: NZUsize!(1024),
                    page_cache,
                },
                translator: EightCap,
                init_cache: Some(NZUsize!(3)),
                init_buffer: NZUsize!(1 << 21),
                init_concurrency: (),
            };

            let mut db: GenericDb<F> =
                commonware_storage::qmdb::any::init(context.child("storage"), cfg, None)
                    .await
                    .expect("init qmdb");

            // committed_state tracks state after apply_batch. pending_expected tracks
            // uncommitted mutations that haven't been applied yet.
            let mut committed_state: HashMap<RawKey, RawValue> = HashMap::new();
            let mut pending_inserts: HashMap<RawKey, RawValue> = HashMap::new();
            let mut pending_deletes: HashSet<RawKey> = HashSet::new();
            let mut all_keys: HashSet<RawKey> = HashSet::new();
            let mut pending_writes: Vec<(Key, Option<Value>)> = Vec::new();

            // Every commit walks the floor under the plan of the most recent Commit operation.
            let mut plan = Plan::Proportional;

            for op in operations.iter().take(MAX_OPS) {
                db = match op {
                    QmdbOperation::Update { key, value } => {
                        let k = Key::new(*key);
                        let v = Value::new(*value);

                        pending_writes.push((k, Some(v)));
                        pending_deletes.remove(key);
                        pending_inserts.insert(*key, *value);
                        all_keys.insert(*key);
                        db
                    }

                    QmdbOperation::Delete { key } => {
                        let k = Key::new(*key);
                        pending_writes.push((k, None));
                        pending_inserts.remove(key);
                        pending_deletes.insert(*key);
                        all_keys.insert(*key);
                        db
                    }

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

                    QmdbOperation::Commit(next) => {
                        plan = next.clone();
                        commit_pending(
                            db, &plan, &mut pending_writes, &mut committed_state,
                            &mut pending_inserts, &mut pending_deletes,
                        ).await
                    }

                    QmdbOperation::Root => {
                        let db = commit_pending(
                            db, &plan, &mut pending_writes, &mut committed_state,
                            &mut pending_inserts, &mut pending_deletes,
                        ).await;
                        db.root();
                        db
                    }

                    QmdbOperation::Proof { start_loc, max_ops } => {
                        let db = commit_pending(
                            db, &plan, &mut pending_writes, &mut committed_state,
                            &mut pending_inserts, &mut pending_deletes,
                        ).await;
                        let actual_op_count = db.bounds().end;

                        if actual_op_count > 0 {
                            let current_root = db.root();
                            let adjusted_start = Location::<F>::new(*start_loc % *actual_op_count);
                            let (proof, log) = db
                                .proof(adjusted_start, *max_ops)
                                .await
                                .expect("proof should not fail");

                            assert!(
                                verify_proof::<Sha256, _, _>(
                                    &proof,
                                    adjusted_start,
                                    &log,
                                    &current_root),
                                "Proof verification failed for start_loc={adjusted_start}, max_ops={max_ops}",
                            );
                        }
                        db
                    }

                    QmdbOperation::ArbitraryProof { start_loc, max_ops , proof_leaves, digests} => {
                        let db = commit_pending(
                            db, &plan, &mut pending_writes, &mut committed_state,
                            &mut pending_inserts, &mut pending_deletes,
                        ).await;
                        let actual_op_count = db.bounds().end;

                        let proof = Proof {
                            leaves: Location::<F>::new(*proof_leaves),
                            inactive_peaks: 0,
                            digests: digests.iter().map(|d| Digest::from(*d)).collect(),
                        };

                        if actual_op_count > 0 {
                            let current_root = db.root();
                            let adjusted_start = Location::<F>::new(*start_loc % *actual_op_count);

                            if let Ok(res) = db
                                .proof(adjusted_start, *max_ops)
                                .await {
                                    let _ = verify_proof::<Sha256, _, _>(
                                        &proof,
                                        adjusted_start,
                                        &res.1,
                                        &current_root);

                            }
                        }
                        db
                    }

                    QmdbOperation::Get { key } => {
                        let k = Key::new(*key);
                        let result = db.get(&k).await.expect("get should not fail");

                        // Verify against committed state only.
                        match committed_state.get(key) {
                            Some(expected_value) => {
                                let v = result.expect("get should not fail");
                                let v_bytes: &[u8; 64] = v.as_ref().try_into().expect("bytes");
                                assert_eq!(v_bytes, expected_value, "Value mismatch for key {key:?}");
                            }
                            None => {
                                assert!(
                                    result.is_none(),
                                    "Found unexpected value for key {key:?} that was never touched",
                                );
                            }
                        }
                        all_keys.insert(*key);
                        db
                    }

                    QmdbOperation::GetSpan { key } => {
                        let k = Key::new(*key);
                        let result = db.get_span(&k).await.expect("get should not fail");
                        assert_eq!(result.is_some(), !db.is_empty(), "span should be empty only if db is empty");
                        db
                    }

                    QmdbOperation::GetNeighbors { key } => {
                        assert_neighbors(&db, &committed_state, *key).await;
                        all_keys.insert(*key);
                        db
                    }
                };
            }

            // Final commit to ensure all operations are persisted.
            if !pending_writes.is_empty() {
                db = commit_pending(
                    db, &plan, &mut pending_writes, &mut committed_state,
                    &mut pending_inserts, &mut pending_deletes,
                ).await;
            }

            // Comprehensive final verification - check ALL keys ever touched.
            for key in &all_keys {
                let k = Key::new(*key);
                let result = db.get(&k).await.expect("final get should not fail");

                match committed_state.get(key) {
                    Some(expected_value) => {
                        let v = result.expect("get should not fail");
                        let v_bytes: &[u8; 64] = v.as_ref().try_into().expect("bytes");
                        assert_eq!(
                            v_bytes, expected_value,
                            "Final value mismatch for key {key:?}"
                        );
                    }
                    None => {
                        assert!(
                            result.is_none(),
                            "Deleted key {key:?} should remain deleted, but found value",
                        );
                    },
                }
            }

            for key in all_keys
                .iter()
                .copied()
                .chain([[0u8; 32], [u8::MAX; 32]])
            {
                assert_neighbors(&db, &committed_state, key).await;
            }

            let batch = db.new_batch().merkleize(&db, None, &mut Proportional).await.unwrap();
            let db = db
                .apply_batch(batch)
                .await
                .expect("final commit should not fail");
            let db = db
                .commit()
                .await
                .expect("final commit fsync should not fail");
            db.destroy().await.expect("destroy should not fail");
        }
    });
}

fn fuzz(input: FuzzInput) {
    fuzz_family::<mmr::Family>(&input, "fuzz-mmr");
    fuzz_family::<mmb::Family>(&input, "fuzz-mmb");
}

fuzz_target!(|input: FuzzInput| {
    fuzz(input);
});