brk_computer 0.11.2

use std::{
    collections::{BTreeMap, btree_map::Entry},
    fs,
    path::{Path, PathBuf},
};

use brk_error::{Error, Result};
use brk_types::{Cents, CentsCompact, CentsSats, CentsSquaredSats, Height, Sats, UrpdRaw};
use rustc_hash::FxHashMap;
use vecdb::{Bytes, unlikely};

use super::{Accumulate, CachedUnrealizedState, UnrealizedState};
use crate::distribution::state::pending::{
    PendingCapDelta, PendingCapitalizedCapRawDelta, PendingDelta,
};

/// Type alias for the price-to-sats map used in cost basis data.
pub(super) type CostBasisMap = BTreeMap<CentsCompact, Sats>;

const STATE_TO_KEEP: usize = 10;

/// Common interface for cost basis tracking.
///
/// Implemented by `CostBasisRaw` (scalars only) and `CostBasisData` (full map + scalars).
pub trait CostBasisOps: Send + Sync + 'static {
    fn create(path: &Path, name: &str) -> Self;
    fn import_at_or_before(&mut self, height: Height) -> Result<Height>;
    fn cap_raw(&self) -> CentsSats;
    fn capitalized_cap_raw(&self) -> CentsSquaredSats;
    fn increment(
        &mut self,
        price: Cents,
        sats: Sats,
        price_sats: CentsSats,
        capitalized_cap: CentsSquaredSats,
    );
    fn decrement(
        &mut self,
        price: Cents,
        sats: Sats,
        price_sats: CentsSats,
        capitalized_cap: CentsSquaredSats,
    );
    fn apply_pending(&mut self);
    fn init(&mut self);
    fn clean(&mut self) -> Result<()>;
    fn write(&mut self, height: Height, cleanup: bool) -> Result<()>;
}

// ─── CostBasisRaw ───────────────────────────────────────────────────────────

#[derive(Clone, Default, Debug)]
struct RawState {
    cap_raw: CentsSats,
}

impl RawState {
    fn serialize(&self) -> Vec<u8> {
        self.cap_raw.to_bytes().to_vec()
    }

    fn deserialize(data: &[u8]) -> Result<Self> {
        Ok(Self {
            cap_raw: CentsSats::from_bytes(&data[0..16])?,
        })
    }
}

/// Lightweight cost basis tracking: only cap_raw scalar.
/// No BTreeMap, no unrealized computation, no pending map.
/// Used by cohorts that only need realized cap on restart (amount_range, address).
#[derive(Clone, Debug)]
pub struct CostBasisRaw {
    pathbuf: PathBuf,
    state: Option<RawState>,
    pending_cap: PendingCapDelta,
}

impl CostBasisRaw {
    #[inline]
    pub(crate) fn increment_cap(&mut self, value: CentsSats) {
        self.pending_cap.inc += value;
    }

    #[inline]
    pub(crate) fn decrement_cap(&mut self, value: CentsSats) {
        self.pending_cap.dec += value;
    }

    pub(super) fn path_state(&self, height: Height) -> PathBuf {
        self.pathbuf.join(height.to_string())
    }

    pub(super) fn read_dir(
        &self,
        keep_only_before: Option<Height>,
    ) -> Result<BTreeMap<Height, PathBuf>> {
        if !self.pathbuf.exists() {
            return Ok(BTreeMap::new());
        }
        Ok(fs::read_dir(&self.pathbuf)?
            .filter_map(|entry| {
                let path = entry.ok()?.path();
                let name = path.file_name()?.to_str()?;
                if let Ok(h) = name.parse::<u32>().map(Height::from) {
                    if keep_only_before.is_none_or(|height| h < height) {
                        Some((h, path))
                    } else {
                        let _ = fs::remove_file(path);
                        None
                    }
                } else {
                    None
                }
            })
            .collect())
    }

    fn apply_pending_cap(&mut self) {
        if self.pending_cap.is_zero() {
            return;
        }
        let state = self.state.as_mut().unwrap();

        state.cap_raw += self.pending_cap.inc;
        if unlikely(state.cap_raw.inner() < self.pending_cap.dec.inner()) {
            panic!(
                "CostBasis cap_raw underflow!\n\
                Path: {:?}\n\
                Current cap_raw (after increments): {}\n\
                Trying to decrement by: {}",
                self.pathbuf, state.cap_raw, self.pending_cap.dec
            );
        }
        state.cap_raw -= self.pending_cap.dec;

        self.pending_cap = PendingCapDelta::default();
    }

    fn write_and_cleanup(&mut self, height: Height, cleanup: bool) -> Result<()> {
        if cleanup {
            let files = self.read_dir(Some(height))?;
            for (_, path) in files
                .iter()
                .take(files.len().saturating_sub(STATE_TO_KEEP - 1))
            {
                fs::remove_file(path)?;
            }
        }
        Ok(())
    }
}

impl CostBasisOps for CostBasisRaw {
    fn create(path: &Path, name: &str) -> Self {
        Self {
            pathbuf: path.join(name).join("cost_basis"),
            state: None,
            pending_cap: PendingCapDelta::default(),
        }
    }

    fn import_at_or_before(&mut self, height: Height) -> Result<Height> {
        let files = self.read_dir(None)?;
        let (&height, path) = files.range(..=height).next_back().ok_or(Error::NotFound(
            "No cost basis state found at or before height".into(),
        ))?;
        let data = fs::read(path)?;
        // Handle both formats: full (map + raw at end) and raw-only (16 bytes).
        self.state = Some(if data.len() == 16 {
            RawState::deserialize(&data)?
        } else {
            let (_, rest) = UrpdRaw::deserialize_with_rest(&data)?;
            RawState::deserialize(rest)?
        });
        self.pending_cap = PendingCapDelta::default();
        Ok(height)
    }

    fn cap_raw(&self) -> CentsSats {
        debug_assert!(self.pending_cap.is_zero());
        self.state.as_ref().unwrap().cap_raw
    }

    fn capitalized_cap_raw(&self) -> CentsSquaredSats {
        CentsSquaredSats::ZERO
    }

    #[inline]
    fn increment(
        &mut self,
        _price: Cents,
        _sats: Sats,
        price_sats: CentsSats,
        _capitalized_cap: CentsSquaredSats,
    ) {
        self.increment_cap(price_sats);
    }

    #[inline]
    fn decrement(
        &mut self,
        _price: Cents,
        _sats: Sats,
        price_sats: CentsSats,
        _capitalized_cap: CentsSquaredSats,
    ) {
        self.decrement_cap(price_sats);
    }

    fn apply_pending(&mut self) {
        self.apply_pending_cap();
    }

    fn init(&mut self) {
        self.state.replace(RawState::default());
        self.pending_cap = PendingCapDelta::default();
    }

    fn clean(&mut self) -> Result<()> {
        let _ = fs::remove_dir_all(&self.pathbuf);
        fs::create_dir_all(&self.pathbuf)?;
        Ok(())
    }

    fn write(&mut self, height: Height, cleanup: bool) -> Result<()> {
        self.apply_pending_cap();
        self.write_and_cleanup(height, cleanup)?;
        fs::write(
            self.path_state(height),
            self.state.as_ref().unwrap().serialize(),
        )?;
        Ok(())
    }
}

// ─── CostBasisData ──────────────────────────────────────────────────────────

/// Full cost basis tracking: BTreeMap distribution + raw scalars.
/// Composes `CostBasisRaw` for scalar tracking, adds map, pending, and cache.
///
/// Generic over the accumulator `S`:
/// - `WithCapital`: tracks all fields including invested capital + capitalized cap (128 bytes)
/// - `WithoutCapital`: tracks only supply + unrealized profit/loss (64 bytes, 1 cache line)
///   and writes a compact checkpoint without a capitalized-cap slot
#[derive(Clone, Debug)]
pub struct CostBasisData<S: Accumulate> {
    raw: CostBasisRaw,
    map: Option<UrpdRaw>,
    pending: FxHashMap<CentsCompact, PendingDelta>,
    cache: Option<CachedUnrealizedState<S>>,
    capitalized_cap_raw: CentsSquaredSats,
    pending_capitalized_cap: PendingCapitalizedCapRawDelta,
}

impl<S: Accumulate> CostBasisData<S> {
    pub(crate) fn map(&self) -> &CostBasisMap {
        debug_assert!(self.pending.is_empty() && self.raw.pending_cap.is_zero());
        &self.map.as_ref().unwrap().map
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.pending.is_empty() && self.map.as_ref().unwrap().map.is_empty()
    }

    pub(crate) fn for_each_pending(&self, mut f: impl FnMut(&CentsCompact, &PendingDelta)) {
        self.pending.iter().for_each(|(k, v)| f(k, v));
    }

    pub(crate) fn compute_unrealized_state(&mut self, height_price: Cents) -> UnrealizedState {
        if self.is_empty() {
            return UnrealizedState::ZERO;
        }

        let map = &self.map.as_ref().unwrap().map;

        if let Some(cache) = self.cache.as_mut() {
            cache.get_at_price(height_price, map)
        } else {
            let cache = CachedUnrealizedState::compute_fresh(height_price, map);
            let state = cache.current_state();
            self.cache = Some(cache);
            state
        }
    }

    fn apply_map_pending(&mut self) {
        if self.pending.is_empty() {
            return;
        }
        let map = &mut self.map.as_mut().unwrap().map;
        for (cents, PendingDelta { inc, dec }) in self.pending.drain() {
            match map.entry(cents) {
                Entry::Occupied(mut e) => {
                    *e.get_mut() += inc;
                    if unlikely(*e.get() < dec) {
                        panic!(
                            "CostBasisData::apply_pending underflow!\n\
                            Path: {:?}\n\
                            Price: {}\n\
                            Current + increments: {}\n\
                            Trying to decrement by: {}",
                            self.raw.pathbuf,
                            cents.to_dollars(),
                            e.get(),
                            dec
                        );
                    }
                    *e.get_mut() -= dec;
                    if *e.get() == Sats::ZERO {
                        e.remove();
                    }
                }
                Entry::Vacant(e) => {
                    if unlikely(inc < dec) {
                        panic!(
                            "CostBasisData::apply_pending underflow (new entry)!\n\
                            Path: {:?}\n\
                            Price: {}\n\
                            Increment: {}\n\
                            Trying to decrement by: {}",
                            self.raw.pathbuf,
                            cents.to_dollars(),
                            inc,
                            dec
                        );
                    }
                    let val = inc - dec;
                    if val != Sats::ZERO {
                        e.insert(val);
                    }
                }
            }
        }
    }
}

impl<S: Accumulate> CostBasisOps for CostBasisData<S> {
    fn create(path: &Path, name: &str) -> Self {
        Self {
            raw: CostBasisRaw::create(path, name),
            map: None,
            pending: FxHashMap::default(),
            cache: None,
            capitalized_cap_raw: CentsSquaredSats::ZERO,
            pending_capitalized_cap: PendingCapitalizedCapRawDelta::default(),
        }
    }

    fn import_at_or_before(&mut self, height: Height) -> Result<Height> {
        let files = self.raw.read_dir(None)?;
        let (&height, path) = files.range(..=height).next_back().ok_or(Error::NotFound(
            "No cost basis state found at or before height".into(),
        ))?;
        let data = fs::read(path)?;
        let (base, rest) = UrpdRaw::deserialize_with_rest(&data)?;
        debug_assert!(
            rest.len() >= if S::TRACK_CAPITAL { 32 } else { 16 },
            "CostBasisData state too short: {} bytes",
            rest.len()
        );
        self.map = Some(base);
        self.raw.state = Some(RawState::deserialize(rest)?);
        self.capitalized_cap_raw = if S::TRACK_CAPITAL {
            CentsSquaredSats::from_bytes(&rest[16..32])?
        } else {
            CentsSquaredSats::ZERO
        };
        self.pending.clear();
        self.raw.pending_cap = PendingCapDelta::default();
        self.pending_capitalized_cap = PendingCapitalizedCapRawDelta::default();
        self.cache = None;
        Ok(height)
    }

    fn cap_raw(&self) -> CentsSats {
        self.raw.cap_raw()
    }

    fn capitalized_cap_raw(&self) -> CentsSquaredSats {
        self.capitalized_cap_raw
    }

    #[inline]
    fn increment(
        &mut self,
        price: Cents,
        sats: Sats,
        price_sats: CentsSats,
        capitalized_cap: CentsSquaredSats,
    ) {
        self.pending.entry(price.into()).or_default().inc += sats;
        self.raw.increment_cap(price_sats);
        if S::TRACK_CAPITAL && capitalized_cap != CentsSquaredSats::ZERO {
            self.pending_capitalized_cap.inc += capitalized_cap;
        }
        if let Some(cache) = self.cache.as_mut() {
            cache.on_receive(price, sats);
        }
    }

    #[inline]
    fn decrement(
        &mut self,
        price: Cents,
        sats: Sats,
        price_sats: CentsSats,
        capitalized_cap: CentsSquaredSats,
    ) {
        self.pending.entry(price.into()).or_default().dec += sats;
        self.raw.decrement_cap(price_sats);
        if S::TRACK_CAPITAL && capitalized_cap != CentsSquaredSats::ZERO {
            self.pending_capitalized_cap.dec += capitalized_cap;
        }
        if let Some(cache) = self.cache.as_mut() {
            cache.on_send(price, sats);
        }
    }

    fn apply_pending(&mut self) {
        self.apply_map_pending();
        self.raw.apply_pending_cap();
        if S::TRACK_CAPITAL {
            self.capitalized_cap_raw += self.pending_capitalized_cap.inc;
            debug_assert!(
                self.capitalized_cap_raw >= self.pending_capitalized_cap.dec,
                "CostBasis capitalized_cap_raw underflow!\n\
                Path: {:?}\n\
                Current (after increments): {:?}\n\
                Trying to decrement by: {:?}",
                self.raw.pathbuf,
                self.capitalized_cap_raw,
                self.pending_capitalized_cap.dec
            );
            self.capitalized_cap_raw -= self.pending_capitalized_cap.dec;
            self.pending_capitalized_cap = PendingCapitalizedCapRawDelta::default();
        }
    }

    fn init(&mut self) {
        self.raw.init();
        self.map.replace(UrpdRaw::default());
        self.pending.clear();
        self.cache = None;
        self.capitalized_cap_raw = CentsSquaredSats::ZERO;
        self.pending_capitalized_cap = PendingCapitalizedCapRawDelta::default();
    }

    fn clean(&mut self) -> Result<()> {
        self.raw.clean()?;
        self.cache = None;
        Ok(())
    }

    fn write(&mut self, height: Height, cleanup: bool) -> Result<()> {
        self.apply_pending();
        self.raw.write_and_cleanup(height, cleanup)?;

        let raw_state = self.raw.state.as_ref().unwrap();
        let mut buffer = self.map.as_ref().unwrap().serialize()?;
        buffer.extend(raw_state.cap_raw.to_bytes());
        if S::TRACK_CAPITAL {
            buffer.extend(self.capitalized_cap_raw.to_bytes());
        }
        fs::write(self.raw.path_state(height), buffer)?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::{
        fs,
        sync::atomic::{AtomicU64, Ordering},
    };

    use super::*;
    use crate::distribution::state::{WithCapital, WithoutCapital};

    static NEXT_PATH_ID: AtomicU64 = AtomicU64::new(0);

    fn test_path() -> PathBuf {
        std::env::temp_dir().join(format!(
            "brk-cost-basis-data-{}-{}",
            std::process::id(),
            NEXT_PATH_ID.fetch_add(1, Ordering::Relaxed)
        ))
    }

    #[test]
    fn checkpoint_layout_tracks_capital_capability() {
        let root = test_path();

        let mut compact = CostBasisData::<WithoutCapital>::create(&root, "compact");
        compact.clean().unwrap();
        compact.init();
        compact.write(Height::ZERO, false).unwrap();

        let mut full = CostBasisData::<WithCapital>::create(&root, "full");
        full.clean().unwrap();
        full.init();
        full.write(Height::ZERO, false).unwrap();

        let compact_len = fs::metadata(compact.raw.path_state(Height::ZERO))
            .unwrap()
            .len();
        let full_len = fs::metadata(full.raw.path_state(Height::ZERO))
            .unwrap()
            .len();
        assert_eq!(full_len, compact_len + 16);

        let mut compact_reader = CostBasisData::<WithoutCapital>::create(&root, "compact");
        assert_eq!(
            compact_reader.import_at_or_before(Height::ZERO).unwrap(),
            Height::ZERO
        );

        let mut legacy_reader = CostBasisData::<WithoutCapital>::create(&root, "full");
        assert_eq!(
            legacy_reader.import_at_or_before(Height::ZERO).unwrap(),
            Height::ZERO
        );
        assert_eq!(legacy_reader.capitalized_cap_raw(), CentsSquaredSats::ZERO);

        fs::remove_dir_all(root).unwrap();
    }
}