use super::{Config, Error};
use crate::{Context, rmap::RMap};
use bytes::Bytes;
use commonware_codec::{Buf, CodecFixed, FixedSize, Read, ReadExt, Write as CodecWrite};
use commonware_cryptography::{Crc32, crc32};
use commonware_formatting::hex;
use commonware_runtime::{
    Blob, BufMut, Error as RError, IoBuf, WriteOptions,
    buffer::{Read as ReadBuffer, Write},
    telemetry::metrics::{Counter, MetricsExt as _},
};
use commonware_utils::{Widen, bitmap::BitMap, futures::try_join_all};
use std::{
    collections::{BTreeMap, BTreeSet, btree_map::Entry},
    marker::PhantomData,
    sync::Arc,
};
use tracing::{debug, warn};

/// Value stored in the index file.
#[derive(Debug, Clone)]
struct Record<V: CodecFixed<Cfg = ()>> {
    value: V,
    crc: u32,
}

impl<V: CodecFixed<Cfg = ()>> Record<V> {
    /// Serialize `value` followed by the CRC of its serialized bytes.
    fn encode(value: &V) -> Vec<u8> {
        let mut buf = Vec::with_capacity(Self::SIZE);
        value.write(&mut buf);
        assert_eq!(buf.len(), V::SIZE, "write() did not write expected bytes");
        let crc = Crc32::checksum(&buf);
        crc.write(&mut buf);
        buf
    }

    /// Deserialize a record, returning the value only if the stored CRC matches the raw
    /// value bytes.
    fn decode_valid(mut buf: IoBuf) -> Option<V> {
        let crc = Crc32::checksum(buf.as_ref().get(..V::SIZE)?);
        let record = Self::read(&mut buf).ok()?;
        (record.crc == crc).then_some(record.value)
    }
}

impl<V: CodecFixed<Cfg = ()>> FixedSize for Record<V> {
    const SIZE: usize = V::SIZE + crc32::Digest::SIZE;
}

impl<V: CodecFixed<Cfg = ()>> CodecWrite for Record<V> {
    fn write(&self, buf: &mut impl BufMut) {
        self.value.write(buf);
        self.crc.write(buf);
    }
}

impl<V: CodecFixed<Cfg = ()>> Read for Record<V> {
    type Cfg = ();

    fn read_cfg(buf: &mut impl Buf, _: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
        let value = V::read(buf)?;
        let crc = u32::read(buf)?;

        Ok(Self { value, crc })
    }
}

#[cfg(feature = "arbitrary")]
impl<V: CodecFixed<Cfg = ()>> arbitrary::Arbitrary<'_> for Record<V>
where
    V: for<'a> arbitrary::Arbitrary<'a>,
{
    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
        let value = V::arbitrary(u)?;
        let mut buf = Vec::with_capacity(V::SIZE);
        value.write(&mut buf);
        let crc = Crc32::checksum(&buf);
        Ok(Self { value, crc })
    }
}

/// The store's state, boxed so the public [Ordinal] handle stays pointer-sized.
struct Inner<E: Context, V: CodecFixed<Cfg = ()>> {
    // Configuration and context
    context: E,
    config: Config,

    // Index blobs for storing key records
    blobs: BTreeMap<u64, Write<E::Blob>>,

    // RMap for interval tracking
    intervals: RMap,

    // Pending sections to be synced.
    pending: BTreeSet<u64>,

    // Metrics
    puts: Counter,
    gets: Counter,
    has: Counter,
    syncs: Counter,
    pruned: Counter,

    _phantom: PhantomData<V>,
}

impl<E: Context, V: CodecFixed<Cfg = ()>> Inner<E, V> {
    /// See [Ordinal::init].
    async fn init(
        context: E,
        config: Config,
        bits: Option<BTreeMap<u64, &Option<BitMap>>>,
    ) -> Result<Self, Error> {
        // Reset the store unless committed bits are provided to recover from the stored blobs
        let record_size = Record::<V>::SIZE as u64;
        let items_per_blob = config.items_per_blob.get();
        let mut blobs = BTreeMap::new();
        let stored_blobs = if bits.is_none() {
            match context.remove(&config.partition, None).await {
                Ok(()) | Err(RError::PartitionMissing(_)) => Vec::new(),
                Err(err) => return Err(Error::Runtime(err)),
            }
        } else {
            match context.scan(&config.partition).await {
                Ok(blobs) => blobs,
                Err(RError::PartitionMissing(_)) => Vec::new(),
                Err(err) => return Err(Error::Runtime(err)),
            }
        };

        // Remove uncovered blobs and open the rest, checking for partial records
        for name in stored_blobs {
            let index = match name.try_into() {
                Ok(index) => u64::from_be_bytes(index),
                Err(nm) => Err(Error::InvalidBlobName(hex(&nm)))?,
            };

            // Remove sections the bits omit or mark no record in without opening them, so a
            // discarded section is never repaired
            let keep = match bits.as_ref().and_then(|bits| bits.get(&index)) {
                Some(Some(bits)) => bits.count_ones() != 0,
                Some(None) => true,
                None => false,
            };
            if !keep {
                context
                    .remove(&config.partition, Some(&index.to_be_bytes()))
                    .await?;
                continue;
            }

            // Open a section the bits keep
            let (blob, mut len) = context
                .open(&config.partition, &index.to_be_bytes())
                .await?;

            // Check if blob size is aligned to record size
            if len % record_size != 0 {
                warn!(
                    blob = index,
                    invalid_size = len,
                    record_size,
                    "blob size is not a multiple of record size, truncating"
                );
                len -= len % record_size;
                blob.resize(len).await?;
                blob.sync().await?;
            }

            debug!(blob = index, len, "found index blob");
            blobs.insert(index, (Arc::new(blob), len));
        }

        // Initialize intervals by scanning committed records
        debug!(
            blobs = blobs.len(),
            "rebuilding intervals from existing index"
        );
        let start = context.current();
        let mut items = 0;
        let mut intervals = RMap::new();
        if let Some(bits) = &bits {
            // Replay ignores records outside the committed bits, but recovery clears them so
            // stored blobs match the checkpointed view. Each run of consecutive unmarked slots is
            // cleared with writes of at most `write_buffer` bytes, or one record if that is larger.
            let zeros = Bytes::from(vec![0u8; config.write_buffer.get().max(Record::<V>::SIZE)]);
            for (section, (blob, size)) in &blobs {
                // A section with no bitmap requires every record, so nothing is cleared
                let Some(Some(bits)) = bits.get(section) else {
                    continue;
                };
                let slots = *size / record_size;
                let mut modified = false;
                let mut run = 0u64;
                for bit_index in 0..=slots {
                    // Extend the run through unmarked slots and clear it at the next marked slot
                    // or the end of the blob
                    if bit_index < slots && (bit_index >= bits.len() || !bits.get(bit_index)) {
                        run += 1;
                        continue;
                    }
                    if run == 0 {
                        continue;
                    }
                    let mut offset = (bit_index - run) * record_size;
                    let end = bit_index * record_size;
                    while offset < end {
                        let len = usize::try_from(end - offset)
                            .map_or(zeros.len(), |left| left.min(zeros.len()));
                        blob.write_at(offset, zeros.slice(..len), WriteOptions::default())
                            .await?;
                        offset += Widen::widen(len);
                    }
                    modified = true;
                    run = 0;
                }
                if modified {
                    blob.sync().await?;
                }
            }

            // Rebuild intervals from the committed records
            for (section, bits) in bits {
                if let Some(bits) = bits
                    && bits.count_ones() == 0
                {
                    continue;
                }

                let Some((blob, size)) = blobs.get(section) else {
                    return Err(Error::MissingRecord(section * items_per_blob));
                };

                // A section kept whole holds exactly `items_per_blob` records. A longer blob was
                // written under a larger `items_per_blob`, and replaying only its first records
                // would hide the rest.
                if bits.is_none()
                    && items_per_blob
                        .checked_mul(record_size)
                        .is_some_and(|full| *size > full)
                {
                    return Err(Error::InvalidRecord(
                        section * items_per_blob + items_per_blob,
                    ));
                }

                // A section replays every record unless a bitmap restricts replay
                // to the records it marks
                let mut set_indices = bits.as_ref().map(|bits| bits.ones_iter());
                let mut all_indices = 0..items_per_blob;

                // A committed bitmap already proves membership, so marked records are not
                // re-read and damage surfaces at get. Membership of an unmarked section
                // comes from record validity, so its records must be read.
                let mut replay_blob = bits
                    .is_none()
                    .then(|| ReadBuffer::new(Arc::clone(blob), *size, config.replay_buffer));
                while let Some(bit_index) = set_indices
                    .as_mut()
                    .map_or_else(|| all_indices.next(), |indices| indices.next())
                {
                    let index = section * items_per_blob + bit_index;
                    if bit_index >= items_per_blob {
                        return Err(Error::MissingRecord(index));
                    }
                    let offset = bit_index * record_size;
                    if offset + record_size > *size {
                        return Err(Error::MissingRecord(index));
                    }

                    // A committed record that is missing or invalid cannot be recovered
                    if let Some(replay_blob) = replay_blob.as_mut() {
                        replay_blob.seek_to(offset)?;
                        let record_buf = replay_blob.read(Record::<V>::SIZE).await?.coalesce();
                        if Record::<V>::decode_valid(record_buf).is_none() {
                            return Err(Error::MissingRecord(index));
                        }
                    }
                    items += 1;
                    intervals.insert(index);
                }
            }
        }
        debug!(
            items,
            elapsed = ?context.current().duration_since(start).unwrap_or_default(),
            "rebuilt intervals"
        );

        // Wrap blobs in write buffers
        let blobs = blobs
            .into_iter()
            .map(|(index, (blob, len))| {
                let blob = Arc::into_inner(blob).expect("replay readers must be closed");
                (index, Write::new(blob, len, config.write_buffer))
            })
            .collect();

        // Initialize metrics
        let puts = context.counter("puts", "Number of put calls");
        let gets = context.counter("gets", "Number of get calls");
        let has = context.counter("has", "Number of has calls");
        let syncs = context.counter("syncs", "Number of sync calls");
        let pruned = context.counter("pruned", "Number of pruned blobs");

        Ok(Self {
            context,
            config,
            blobs,
            intervals,
            pending: BTreeSet::new(),
            puts,
            gets,
            has,
            syncs,
            pruned,
            _phantom: PhantomData,
        })
    }

    /// See [Ordinal::put].
    async fn put(mut self: Box<Self>, index: u64, value: V) -> Result<Box<Self>, Error> {
        self.puts.inc();

        // Merge the record into the section's buffer when it fits. Otherwise take the section's
        // writer, or open one when the section has none.
        let items_per_blob = self.config.items_per_blob.get();
        let section = index / items_per_blob;
        let offset = (index % items_per_blob) * Record::<V>::SIZE as u64;
        let record = Record::encode(&value);
        let blob = match self.blobs.entry(section) {
            Entry::Occupied(mut entry) => {
                if entry.get_mut().try_write_at(offset, &record) {
                    None
                } else {
                    Some(entry.remove())
                }
            }
            Entry::Vacant(_) => {
                let (blob, len) = self
                    .context
                    .open(&self.config.partition, &section.to_be_bytes())
                    .await?;
                debug!(section, "created blob");
                Some(Write::new(blob, len, self.config.write_buffer))
            }
        };

        // Write through a writer that was just opened or whose buffer could not take the record.
        if let Some(blob) = blob {
            let blob = blob.write_at(offset, record).await?;
            self.blobs.insert(section, blob);
        }

        // Track the accepted record for syncing and lookup.
        self.pending.insert(section);
        self.intervals.insert(index);

        Ok(self)
    }

    /// See [Ordinal::get].
    async fn get(&self, index: u64) -> Result<Option<V>, Error> {
        self.gets.inc();

        // If get isn't in an interval, it doesn't exist and we don't need to access disk
        if self.intervals.get(index).is_none() {
            return Ok(None);
        }

        // Read from disk
        let items_per_blob = self.config.items_per_blob.get();
        let section = index / items_per_blob;
        let blob = self.blobs.get(&section).unwrap();
        let offset = (index % items_per_blob) * Record::<V>::SIZE as u64;
        let read_buf = blob.read_at(offset, Record::<V>::SIZE).await?.coalesce();

        // If record is valid, return it
        let value = Record::<V>::decode_valid(read_buf).ok_or(Error::InvalidRecord(index))?;
        Ok(Some(value))
    }

    /// See [Ordinal::has].
    fn has(&self, index: u64) -> bool {
        self.has.inc();

        self.intervals.get(index).is_some()
    }

    /// See [Ordinal::next_gap].
    fn next_gap(&self, index: u64) -> (Option<u64>, Option<u64>) {
        self.intervals.next_gap(index)
    }

    /// See [Ordinal::ranges].
    fn ranges(&self) -> impl Iterator<Item = (u64, u64)> + '_ {
        self.intervals.iter().map(|(&s, &e)| (s, e))
    }

    /// See [Ordinal::ranges_from].
    fn ranges_from(&self, from: u64) -> impl Iterator<Item = (u64, u64)> + '_ {
        self.intervals.iter_from(from).map(|(&s, &e)| (s, e))
    }

    /// See [Ordinal::first_index].
    fn first_index(&self) -> Option<u64> {
        self.intervals.first_index()
    }

    /// See [Ordinal::last_index].
    fn last_index(&self) -> Option<u64> {
        self.intervals.last_index()
    }

    /// See [Ordinal::missing_items].
    fn missing_items(&self, start: u64, max: usize) -> Vec<u64> {
        self.intervals.missing_items(start, max)
    }

    /// See [Ordinal::prune].
    async fn prune(mut self: Box<Self>, min: u64) -> Result<Box<Self>, Error> {
        // Collect sections to remove
        let items_per_blob = self.config.items_per_blob.get();
        let min_section = min / items_per_blob;
        let sections_to_remove: Vec<u64> = self
            .blobs
            .keys()
            .filter(|&&section| section < min_section)
            .copied()
            .collect();

        // Remove the collected sections
        for section in sections_to_remove {
            if let Some(blob) = self.blobs.remove(&section) {
                drop(blob);
                self.context
                    .remove(&self.config.partition, Some(&section.to_be_bytes()))
                    .await?;

                // Remove the corresponding index range from intervals
                let start_index = section * items_per_blob;
                let end_index = (section + 1) * items_per_blob - 1;
                self.intervals.remove(start_index, end_index);
                debug!(section, start_index, end_index, "pruned blob");
            }

            // Update metrics
            self.pruned.inc();
        }

        // Clean pending entries that fall into pruned sections.
        self.pending.retain(|&section| section >= min_section);

        Ok(self)
    }

    /// See [Ordinal::sync].
    async fn sync(mut self: Box<Self>) -> Result<Box<Self>, Error> {
        self.syncs.inc();

        if self.pending.is_empty() {
            return Ok(self);
        }

        // Own pending writers across their syncs and restore them only after every sync succeeds.
        let futures: Vec<_> = self
            .blobs
            .extract_if(.., |section, _| self.pending.contains(section))
            .map(|(section, blob)| async move { blob.sync().await.map(|blob| (section, blob)) })
            .collect();
        let blobs = try_join_all(futures).await?;
        self.blobs.extend(blobs);

        // Clear pending sections.
        self.pending.clear();

        Ok(self)
    }

    /// See [Ordinal::destroy].
    async fn destroy(self) -> Result<(), Error> {
        for (i, blob) in self.blobs.into_iter() {
            drop(blob);
            self.context
                .remove(&self.config.partition, Some(&i.to_be_bytes()))
                .await?;
            debug!(section = i, "destroyed blob");
        }
        match self.context.remove(&self.config.partition, None).await {
            Ok(()) => {}
            Err(RError::PartitionMissing(_)) => {
                // Partition already removed or never existed.
            }
            Err(err) => return Err(Error::Runtime(err)),
        }
        Ok(())
    }
}

/// Implementation of [Ordinal].
///
/// Mutating functions consume the store and return it only on success: an error (or a dropped
/// future) destroys the handle.
pub struct Ordinal<E: Context, V: CodecFixed<Cfg = ()>>(Box<Inner<E, V>>);

impl<E: Context, V: CodecFixed<Cfg = ()>> std::fmt::Debug for Ordinal<E, V> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Ordinal")
            .field("first_index", &self.0.intervals.first_index())
            .field("last_index", &self.0.intervals.last_index())
            .finish_non_exhaustive()
    }
}

impl<E: Context, V: CodecFixed<Cfg = ()>> Ordinal<E, V> {
    /// Initialize a new [Ordinal] instance with a collection of [BitMap]s (indicating which
    /// records should be considered available).
    ///
    /// If a section is not provided in the [BTreeMap], all records in that section are considered
    /// unavailable. If a [BitMap] is provided for a section, all records in that section are
    /// considered available if and only if the [BitMap] is set for the record. If a section is provided
    /// but no [BitMap] is populated, all records in that section are considered available, and the
    /// section must hold exactly `items_per_blob` records.
    ///
    /// Passing `Some(BTreeMap::new())` or `None` removes all stored sections and starts empty.
    pub async fn init(
        context: E,
        config: Config,
        bits: Option<BTreeMap<u64, &Option<BitMap>>>,
    ) -> Result<Self, Error> {
        Ok(Self(Box::new(Inner::init(context, config, bits).await?)))
    }

    /// Add a value at the specified index (pending until sync).
    pub async fn put(mut self, index: u64, value: V) -> Result<Self, Error> {
        self.0 = self.0.put(index, value).await?;
        Ok(self)
    }

    /// Get the value for a given index.
    pub async fn get(&self, index: u64) -> Result<Option<V>, Error> {
        self.0.get(index).await
    }

    /// Check if an index exists.
    pub fn has(&self, index: u64) -> bool {
        self.0.has(index)
    }

    /// Get the next gap information for backfill operations.
    pub fn next_gap(&self, index: u64) -> (Option<u64>, Option<u64>) {
        self.0.next_gap(index)
    }

    /// Get an iterator over all ranges in the [Ordinal].
    pub fn ranges(&self) -> impl Iterator<Item = (u64, u64)> + '_ {
        self.0.ranges()
    }

    /// Get an iterator over ranges that overlap or follow `from`.
    pub fn ranges_from(&self, from: u64) -> impl Iterator<Item = (u64, u64)> + '_ {
        self.0.ranges_from(from)
    }

    /// Retrieve the first index in the [Ordinal].
    pub fn first_index(&self) -> Option<u64> {
        self.0.first_index()
    }

    /// Retrieve the last index in the [Ordinal].
    pub fn last_index(&self) -> Option<u64> {
        self.0.last_index()
    }

    /// Returns up to `max` missing items starting from `start`.
    ///
    /// This method iterates through gaps between existing ranges, collecting missing indices
    /// until either `max` items are found or there are no more gaps to fill.
    pub fn missing_items(&self, start: u64, max: usize) -> Vec<u64> {
        self.0.missing_items(start, max)
    }

    /// Prune indices older than `min` by removing entire blobs.
    ///
    /// Pruning is done at blob boundaries to avoid partial deletions. A blob is pruned only if
    /// all possible indices in that blob are less than `min`.
    pub async fn prune(mut self, min: u64) -> Result<Self, Error> {
        self.0 = self.0.prune(min).await?;
        Ok(self)
    }

    /// Write all pending entries and sync all modified [Blob]s.
    pub async fn sync(mut self) -> Result<Self, Error> {
        self.0 = self.0.sync().await?;
        Ok(self)
    }

    /// Destroy [Ordinal] and remove all data.
    pub async fn destroy(self) -> Result<(), Error> {
        self.0.destroy().await
    }
}

#[cfg(all(test, feature = "arbitrary"))]
mod conformance {
    use super::*;
    use commonware_codec::conformance::CodecConformance;

    commonware_conformance::conformance_tests! {
        CodecConformance<Record<u32>>
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::utils::codec::View;
    use commonware_runtime::deterministic::Context;

    #[test]
    fn test_record_preserves_owned_byte_fields() {
        let value = View::new(7);
        let encoded = Record::encode(&value);
        let source = IoBuf::from(encoded.clone());
        let decoded = Record::<View>::decode_valid(source).unwrap();
        assert_eq!(decoded.bytes, value.bytes);
        decoded.assert_shared();

        let mut corrupt = encoded.clone();
        corrupt[0] ^= 1;
        assert!(Record::<View>::decode_valid(corrupt.into()).is_none());
        let mut truncated = encoded;
        truncated.pop();
        assert!(Record::<View>::decode_valid(truncated.into()).is_none());
    }

    type TestOrdinal = Ordinal<Context, u64>;

    fn is_send<T: Send>(_: T) {}

    #[allow(dead_code)]
    fn assert_ordinal_futures_are_send(ordinal: TestOrdinal, key: u64) {
        is_send(ordinal.get(key));
        is_send(ordinal.put(key, 0u64));
    }

    #[allow(dead_code)]
    fn assert_ordinal_destroy_is_send(ordinal: TestOrdinal) {
        is_send(ordinal.destroy());
    }
}