hoard/src/data.rs

use std::collections::HashMap;

use itertools::Itertools;

use ariadne::Cache;
use chumsky::span::Span as _;

pub use rust_decimal::Decimal;

use crate::prelude::*;

mod format;
mod parse;
pub mod spec;

pub use format::format_ledger;
pub use parse::parse_ledger;

pub struct UnitTag;
impl stringstore::NamespaceTag for UnitTag {
    const PREFIX: &'static str = "unit";
}
pub type UnitName = stringstore::StoredString<UnitTag>;

pub struct AccountTag;
impl stringstore::NamespaceTag for AccountTag {
    const PREFIX: &'static str = "acc";
}
pub type AccountName = stringstore::StoredString<AccountTag>;

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Datestamp {
    pub year: u16,
    pub month: u8,
    pub day: u8,
}

impl std::fmt::Display for Datestamp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:04}-{:02}-{:02}", self.year, self.month, self.day)
    }
}

impl std::fmt::Debug for Datestamp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Datestamp ({self})")
    }
}

#[derive(Clone, Debug, PartialEq, PartialOrd, Ord, Eq)]
pub struct Change {
    pub source: Option<io::Span>,

    pub account: Spanned<AccountName>,
    pub amount: Spanned<Decimal>,
    pub balance: Option<Spanned<Decimal>>,
    pub unit: Spanned<UnitName>,
}

#[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
pub struct Transaction {
    pub source: Option<io::Span>,

    pub datestamp: Datestamp,
    pub title: Option<String>,
    pub annotations: Vec<String>,
    pub changes: Vec<Change>,
}

impl Transaction {
    pub fn modifies(&self, account: AccountName) -> bool {
        self.changes.iter().any(|b| b.account.as_ref() == &account)
    }

    pub fn change_for(&self, account: AccountName) -> Option<&Change> {
        self.changes.iter().find(|b| b.account.as_ref() == &account)
    }

    pub fn split_changes(
        &self,
        account: AccountName,
    ) -> Option<(&Change, impl Iterator<Item = &Change>)> {
        let index = self
            .changes
            .iter()
            .position(|b| b.account.as_ref() == &account)?;
        Some((
            &self.changes[index],
            self.changes[0..index]
                .iter()
                .chain(self.changes[index + 1..].iter()),
        ))
    }

    pub fn is_mono_unit(&self) -> bool {
        self.changes.iter().unique_by(|b| *b.unit).count() == 1
    }

    pub fn mono_unit(&self) -> Option<UnitName> {
        let mut it = self.changes.iter().unique_by(|b| *b.unit);
        let uniq = it.next()?;
        it.next().is_none().then_some(*uniq.unit)
    }

    pub fn get_annotation(&self, label: &str) -> Option<&str> {
        for anno in self.annotations.iter() {
            if let Some(body) = anno.strip_prefix(label) {
                return Some(body);
            }
        }
        None
    }
}

pub type TransactionRef = std::rc::Rc<std::cell::RefCell<Transaction>>;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
enum RawLedgerEntry {
    Transaction(Transaction),
    Comment(Spanned<String>),
}

impl RawLedgerEntry {
    pub fn as_transaction(&self) -> Option<&Transaction> {
        match self {
            Self::Transaction(tx) => Some(tx),
            _ => None,
        }
    }

    pub fn as_transaction_mut(&mut self) -> Option<&mut Transaction> {
        match self {
            Self::Transaction(tx) => Some(tx),
            _ => None,
        }
    }

    pub fn span(&self) -> io::Span {
        match self {
            Self::Transaction(ts) => ts.source.unwrap_or_default(),
            Self::Comment(c) => c.span(),
        }
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum LedgerEntry {
    Transaction(TransactionRef),
    Comment(Spanned<String>),
}

impl LedgerEntry {
    pub fn as_transaction(&self) -> Option<&TransactionRef> {
        match self {
            Self::Transaction(tx) => Some(tx),
            _ => None,
        }
    }

    pub fn as_transaction_mut(&mut self) -> Option<&mut Transaction> {
        match self {
            Self::Transaction(tx) => Some(tx),
            _ => None,
        }
    }

    pub fn span(&self) -> io::Span {
        match self {
            Self::Transaction(ts) => ts.source.unwrap_or_default(),
            Self::Comment(c) => c.span(),
        }
    }
}

#[derive(Debug)]
pub struct Hoard {
    path: std::path::PathBuf,
    spec_root: spec::SpecRoot,

    comments: Vec<Spanned<String>>,

    raw_ledger_data: Vec<RawLedgerEntry>,

    account_ledger_data: HashMap<AccountName, Vec<Transaction>>,
}

impl Hoard {
    pub fn load(
        fsdata: &mut io::FilesystemData,
        path: &std::path::Path,
        check_level: check::CheckLevel,
    ) -> Result<Self, DataError> {
        let sf = io::SourceFile::new(path.as_os_str());
        let root_data = fsdata.fetch(&sf).unwrap();

        match toml::from_str::<spec::SpecRoot>(root_data.text()) {
            Ok(spec_root) => {
                let mut r = Self {
                    path: path.into(),
                    spec_root,
                    raw_ledger_data: vec![],
                    comments: vec![],
                    account_ledger_data: Default::default(),
                };

                r.load_ledgers(fsdata)?;
                r.preprocess_ledger_data();

                crate::check::run_checks(&mut r, check_level)?;

                Ok(r)
            }
            Err(te) => {
                let Some(range) = te.span() else {
                    panic!("TOML parse error with no range: {te}");
                };

                let report = ariadne::Report::build(
                    ariadne::ReportKind::Error,
                    io::Span::new(sf, range.clone()),
                )
                .with_label(
                    ariadne::Label::new(io::Span::new(sf, range)).with_message(te.message()),
                )
                .with_message("Failed to parse root TOML")
                .finish();

                Err(report.into())
            }
        }
    }

    fn load_ledger(
        &mut self,
        fsdata: &mut io::FilesystemData,
        path: &mut std::path::PathBuf,
    ) -> Result<(), DataError> {
        log::debug!("Loading ledger data from {}", path.display());

        let md = std::fs::metadata(path.as_path()).map_err(DataError::IOError)?;

        if md.is_dir() {
            // recurse
            for de in std::fs::read_dir(path.as_path()).map_err(DataError::IOError)? {
                let de = de.map_err(DataError::IOError)?;
                path.push(de.file_name());

                self.load_ledger(fsdata, path)?;

                path.pop();
            }
        } else {
            let path = std::fs::canonicalize(path)?;
            let Some(filename) = path.file_name() else {
                return Ok(());
            };
            // skip filenames beginning with a dot
            if filename.as_encoded_bytes()[0] == b'.' {
                log::info!("Skipping file {}", path.display());
                return Ok(());
            }

            let sf = io::SourceFile::new_from_string(path.into_os_string());
            if let Ok(data) = fsdata.fetch(&sf) {
                self.raw_ledger_data
                    .extend(parse_ledger(sf, &self.spec_root, data.text())?);
            } else {
                log::error!(
                    "Failed to load data from {}",
                    std::path::Path::new(sf.as_str()).display()
                );
            }
        }

        Ok(())
    }

    fn load_ledgers(&mut self, fsdata: &mut io::FilesystemData) -> Result<(), DataError> {
        let mut ledger_path = std::fs::canonicalize(self.path.as_path())?;
        ledger_path.pop();
        ledger_path.push(&self.spec_root.ledger_path);
        let mut ledger_path = std::fs::canonicalize(ledger_path)?;

        self.load_ledger(fsdata, &mut ledger_path)?;

        Ok(())
    }

    fn preprocess_ledger_data(&mut self) {
        for entry in &self.raw_ledger_data {
            let RawLedgerEntry::Transaction(tx) = &entry else {
                continue;
            };
            for bal in &tx.changes {
                self.account_ledger_data
                    .entry(*bal.account)
                    .or_default()
                    .push(tx.clone());
            }
        }
    }

    pub fn all_ledger_data(&self) -> &[LedgerEntry] {
        self.ledger_data.as_slice()
    }

    pub fn all_ledger_data_mut(&mut self) -> &mut [LedgerEntry] {
        self.ledger_data.as_mut_slice()
    }

    pub fn ledger_data_for(&self, aname: AccountName) -> Option<&[Transaction]> {
        self.account_ledger_data.get(&aname).map(Vec::as_slice)
    }

    pub fn ledger_data_for_mut(
        &mut self,
        aname: AccountName,
    ) -> Option<&mut [Transaction]> {
        self.account_ledger_data
            .get_mut(&aname)
            .map(Vec::as_mut_slice)
    }

    pub fn ledger_data_from(&self, source: io::SourceFile) -> impl Iterator<Item = &LedgerEntry> {
        self.all_ledger_data()
            .iter()
            .filter(move |le| le.span().context == Some(source))
    }

    pub fn account_names(&self) -> impl Iterator<Item = AccountName> {
        self.spec_root.accounts.keys().cloned()
    }

    pub fn account_spec(&self, aname: AccountName) -> Option<&spec::AccountSpec> {
        self.spec_root.accounts.get(&aname)
    }

    pub fn unit_spec(&self, unit: UnitName) -> Option<&spec::UnitSpec> {
        self.spec_root.units.get(&unit)
    }

    pub fn spec_root(&self) -> &spec::SpecRoot {
        &self.spec_root
    }

    pub fn balance(&self, aname: AccountName) -> Option<HashMap<UnitName, Decimal>> {
        let mut running = HashMap::<UnitName, Decimal>::new();

        for le in self.ledger_data_for(aname)? {
            for b in &le.changes {
                if *b.account != aname {
                    continue;
                }
                let v = running.entry(*b.unit).or_default();
                *v = v.checked_add(*b.amount)?;
            }
        }

        Some(running)
    }
}