hoard/src/check.rs

use std::collections::BTreeMap;

use ariadne::Span as ASpan;
use chumsky::span::Span as CSpan;

// use crate::data::{DataError, Decimal, Hoard, Span, Spanned, UnitName};
use crate::prelude::*;

#[derive(Default, PartialEq, PartialOrd)]
pub enum CheckLevel {
    /// Check individual transactions
    WellFormed,
    /// Check relations between transactions
    #[default]
    Consistent,
}

fn check_equal_sum(root: &data::Hoard) -> Result<(), DataError> {
    log::trace!("Checking for equal sums in monounit ledger entries...");

    for le in root.all_raw_ledger_data() {
        let Some(tx) = le.as_transaction() else {
            continue;
        };
        let Some(mono) = tx.mono_unit() else { continue };

        let net = tx
            .changes
            .iter()
            .try_fold(data::Decimal::ZERO, |acc, b| acc.checked_add(*b.amount));
        if net != Some(data::Decimal::ZERO) {
            let report =
                ariadne::Report::build(ariadne::ReportKind::Error, tx.source.unwrap_or_default())
                    .with_labels(tx.changes.iter().map(|v| {
                        let span = v.amount.span().union(v.unit.span());
                        ariadne::Label::new(span).with_message("change here")
                    }));

            let report = if let Some(net) = net {
                report
                    .with_message("imbalanced transaction")
                    .with_note(format!("net value: {net} {mono}"))
            } else {
                report.with_note("numeric overflow")
            };

            return Err(report.finish().into());
        }
    }

    Ok(())
}

fn check_precision(root: &data::Hoard) -> Result<(), DataError> {
    log::trace!("Checking for precision errors in ledger entries...");

    for le in root.all_raw_ledger_data() {
        let Some(tx) = le.as_transaction() else {
            continue;
        };

        for change in &tx.changes {
            let unit_spec = root.unit_spec(*change.unit).unwrap();
            let Some(precision) = unit_spec.precision else {
                continue;
            };
            let amount_scale = change.amount.scale();
            if amount_scale > precision {
                return Err(ariadne::Report::build(
                    ariadne::ReportKind::Error,
                    change.amount.span(),
                )
                .with_label(
                    ariadne::Label::new(change.amount.span()).with_message(format!(
                        "{} has {} digits of precision, not {}",
                        change.unit,
                        precision,
                        change.amount.0.scale()
                    )),
                )
                .finish()
                .into());
            }

            let Some(balance) = change.balance.as_ref() else {
                continue;
            };
            let balance_scale = balance.scale();
            if balance_scale > precision {
                return Err(
                    ariadne::Report::build(ariadne::ReportKind::Error, balance.span())
                        .with_label(ariadne::Label::new(balance.span()).with_message(format!(
                            "{} has {} digits of precision, not {}",
                            change.unit,
                            precision,
                            change.amount.0.scale()
                        )))
                        .finish()
                        .into(),
                );
            }
        }
    }

    Ok(())
}

fn check_balances(root: &data::Hoard) -> Result<(), DataError> {
    log::trace!("Checking balance consistency...");

    for account in root.account_names() {
        let Some(ledger) = root.ledger_data_for(account) else {
            continue;
        };

        let mut running_balance = BTreeMap::<data::UnitName, Spanned<data::Decimal>>::new();
        for txn_ref in ledger {
            let txn = txn_ref.borrow();
            let change = txn.change_for(account).unwrap();
            let bal = running_balance
                .entry(*change.unit)
                .or_insert_with(|| Spanned::new(data::Decimal::default(), io::Span::default()));
            let last_span = bal.span();
            bal.0 = bal.checked_add(*change.amount).unwrap();
            bal.1 = change.source.unwrap();

            if let Some(sbal) = change.balance.as_ref() {
                if **sbal != bal.0 {
                    let report = ariadne::Report::build(
                        ariadne::ReportKind::Error,
                        txn.source.unwrap_or_default(),
                    )
                    .with_label(ariadne::Label::new(sbal.span()).with_message(format!(
                        "Calculated balance is {} {}, specified balance is {} {}",
                        bal, change.unit, sbal, change.unit
                    )));
                    let report = if !last_span.is_empty() {
                        report.with_label(
                            ariadne::Label::new(last_span)
                                .with_message("Last balance is from here"),
                        )
                    } else {
                        report
                    };

                    return Err(report.finish().into());
                }
            }
        }
    }
    Ok(())
}

pub fn run_checks(root: &mut data::Hoard, level: CheckLevel) -> Result<(), DataError> {
    if level >= CheckLevel::WellFormed {
        log::debug!("Running transaction well-formedness checks...");
        check_precision(root)?;
        check_equal_sum(root)?;
    }
    if level >= CheckLevel::Consistent {
        log::debug!("Running ledger consistency checks...");
        check_balances(root)?;
    }
    Ok(())
}