hoard/src/data.rs

use std::collections::HashMap;

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

pub use rust_decimal::Decimal;

pub mod ledger;
pub mod spec;

pub struct LocationTag;
impl stringstore::NamespaceTag for LocationTag {
    const PREFIX: &'static str = "loc";
}
pub type SourceFile = stringstore::StoredOsString<LocationTag>;

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, Debug, PartialEq, Hash)]
pub struct DataSource {
    file: SourceFile,
    range: std::ops::Range<usize>,
}

/// Helper for accessing data on the filesystem
#[derive(Default)]
pub struct FilesystemData {
    file_data: HashMap<SourceFile, ariadne::Source<std::rc::Rc<str>>>,
}

impl std::fmt::Debug for FilesystemData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FilesystemData").finish()
    }
}

impl ariadne::Cache<SourceFile> for FilesystemData {
    type Storage = std::rc::Rc<str>;
    fn fetch(
        &mut self,
        id: &SourceFile,
    ) -> Result<&ariadne::Source<Self::Storage>, impl std::fmt::Debug> {
        if !self.file_data.contains_key(id) {
            match std::fs::read_to_string(id.as_str()) {
                Ok(data) => {
                    let data: std::rc::Rc<_> = std::rc::Rc::from(data.into_boxed_str());
                    self.file_data.insert(*id, ariadne::Source::from(data));
                }
                Err(e) => return Err(e),
            }
        }

        Ok(self.file_data.get(id).unwrap())
    }

    fn display<'a>(&self, id: &'a SourceFile) -> Option<impl std::fmt::Display + 'a> {
        Some(id.as_str().to_string_lossy())
    }
}

#[derive(Debug)]
pub enum DataError {
    IOError(std::io::Error),
    Report(Box<ariadne::Report<'static, Span>>),
    Validation(String),
}

impl std::fmt::Display for DataError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        <Self as std::fmt::Debug>::fmt(self, f)
    }
}

impl From<std::io::Error> for DataError {
    fn from(value: std::io::Error) -> Self {
        Self::IOError(value)
    }
}

impl From<ariadne::Report<'static, Span>> for DataError {
    fn from(value: ariadne::Report<'static, Span>) -> Self {
        Self::Report(value.into())
    }
}

impl std::error::Error for DataError {}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Span {
    range: (usize, usize),
    context: Option<SourceFile>,
}

impl Default for Span {
    fn default() -> Self {
        Self { range: (0,0), context: None }
    }
}

impl chumsky::span::Span for Span {
    type Offset = usize;
    type Context = SourceFile;

    fn new(context: Self::Context, range: std::ops::Range<Self::Offset>) -> Self {
        Self {
            context: Some(context),
            range: (range.start, range.end),
        }
    }

    fn start(&self) -> Self::Offset {
        self.range.0
    }

    fn end(&self) -> Self::Offset {
        self.range.1
    }

    fn context(&self) -> Self::Context {
        self.context.unwrap()
    }

    fn to_end(&self) -> Self {
        Self {
            context: self.context,
            range: (self.range.1, self.range.1),
        }
    }
}

impl ariadne::Span for Span {
    type SourceId = SourceFile;
    fn source(&self) -> &Self::SourceId {
        self.context.as_ref().unwrap()
    }
    fn start(&self) -> usize {
        self.range.0
    }
    fn end(&self) -> usize {
        self.range.1
    }
}

#[derive(Debug, Clone, Copy)]
pub struct Spanned<T>(pub T, pub Span);

impl<T> Spanned<T> {
    pub fn new(t: T, span: Span) -> Self {
        Self(t, span)
    }

    pub fn span(&self) -> Span {
        self.1
    }
}

impl<T> std::ops::Deref for Spanned<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> std::ops::DerefMut for Spanned<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<T> AsRef<T> for Spanned<T> {
    fn as_ref(&self) -> &T {
        &self.0
    }
}

impl<T: PartialEq> PartialEq for Spanned<T> {
    fn eq(&self, other: &Self) -> bool {
        self.0.eq(&other.0)
    }
}

impl<T: Eq> Eq for Spanned<T> {}

impl<T> From<T> for Spanned<T> {
    fn from(value: T) -> Self {
        Self(value, Span::default())
    }
}

impl<T: PartialOrd> PartialOrd for Spanned<T> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

impl<T: Ord> Ord for Spanned<T> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.0.cmp(&other.0)
    }
}

impl<T: std::fmt::Display> std::fmt::Display for Spanned<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

#[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(Debug)]
pub struct Root {
    path: std::path::PathBuf,
    spec_root: spec::SpecRoot,

    ledger_data: Vec<ledger::LedgerEntry>,

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

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

        match toml::from_str::<spec::SpecRoot>(root_data.text()) {
            Ok(mut spec_root) => {
                let initial_name = AccountName::from("initial");
                if let std::collections::hash_map::Entry::Vacant(ve) =
                    spec_root.accounts.entry(initial_name)
                {
                    ve.insert(spec::AccountSpec {
                        title: Some(String::from("initial balances")),
                        description: None,
                        annotations: None,
                        unit: None,
                        import: None,
                    });
                } else {
                    return Err(DataError::Validation(String::from(
                        "cannot define 'initial' account, as it is a built-in",
                    )));
                }

                let mut r = Self {
                    path: path.into(),
                    spec_root,
                    ledger_data: vec![],
                    account_ledger_data: Default::default(),
                };

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

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

                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,
                    Span::new(sf, range.clone()),
                )
                .with_label(ariadne::Label::new(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 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 s = SourceFile::new(path.as_os_str());
            let data = fsdata.fetch(&s).unwrap();
            self.ledger_data
                .extend(ledger::parse_ledger(s, &self.spec_root, data.text())?);
        }

        Ok(())
    }

    fn load_ledgers(&mut self, fsdata: &mut FilesystemData) -> Result<(), DataError> {
        let mut ledger_path = self.path.to_owned();
        ledger_path.pop();
        ledger_path.push(&self.spec_root.ledger_path);

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

        self.ledger_data.sort();

        Ok(())
    }

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

        for txns in self.account_ledger_data.values_mut() {
            txns.sort_by_key(|txn| txn.datestamp);
        }
    }

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

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

    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)
    }
}