hoard/src/data.rs

use std::collections::HashMap;

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

pub use rust_decimal::Decimal;

use crate::{
    check::CheckLevel,
    io::{FilesystemData, SourceFile},
};

pub mod ledger;
pub mod spec;

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(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 Span {
    pub fn null_for_file(source: SourceFile) -> Self {
        Self {
            range: (usize::MAX, usize::MAX),
            context: Some(source),
        }
    }
}

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,
        check_level: CheckLevel,
    ) -> 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(spec_root) => {
                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, 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,
                    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 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 = SourceFile::new_from_string(path.into_os_string());
            if let Ok(data) = fsdata.fetch(&sf) {
                self.ledger_data
                    .extend(ledger::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 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.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 ledger_data_from(
        &self,
        source: SourceFile,
    ) -> impl Iterator<Item = &ledger::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)
    }
}