hoard/src/import.rs

use std::collections::BTreeSet;

use crate::prelude::*;

/*use crate::data::{
    AccountName, Datestamp, Decimal, UnitName,
    ledger::{Change, Transaction},
    spec::{AccountSpec, CsvColumnSpec, CsvImportSpec, ImportFileFormat},
};*/

#[derive(Debug)]
pub enum ImportError {
    IOError(std::io::Error),
    ConfigError(String),
    InputError(String),
    CsvError(csv::Error),
}

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

impl From<csv::Error> for ImportError {
    fn from(value: csv::Error) -> Self {
        Self::CsvError(value)
    }
}

impl From<strptime::ParseError> for ImportError {
    fn from(value: strptime::ParseError) -> Self {
        Self::InputError(value.to_string())
    }
}

impl From<rust_decimal::Error> for ImportError {
    fn from(value: rust_decimal::Error) -> Self {
        Self::InputError(value.to_string())
    }
}

fn try_parse_decimal(from: &str) -> Result<data::Decimal, ImportError> {
    // remove any '$' or units from the string
    let filtered = from
        .chars()
        .filter(|f| char::is_digit(*f, 10) || *f == '.' || *f == '-')
        .collect::<String>();

    data::Decimal::from_str_radix(filtered.as_str(), 10).map_err(Into::into)
}

fn import_from_csv(
    csv_spec: &data::spec::CsvImportSpec,
    aspec: &data::spec::AccountSpec,
    account: data::AccountName,
    target: data::AccountName,
    reader: impl std::io::Read,
) -> Result<Vec<data::Transaction>, ImportError> {
    let mut csv_reader = csv::Reader::from_reader(reader);

    // validate CSV spec
    if !csv_spec
        .cols
        .contains(&data::spec::CsvColumnSpec::Datestamp)
    {
        return Err(ImportError::ConfigError(
            "CSV config does not have a datestamp column".into(),
        ));
    }

    if !csv_spec.cols.contains(&data::spec::CsvColumnSpec::Change)
        && (!csv_spec.cols.contains(&data::spec::CsvColumnSpec::Withdraw)
            || !csv_spec.cols.contains(&data::spec::CsvColumnSpec::Deposit))
    {
        return Err(ImportError::ConfigError(
            "CSV config needs either a change column or both withdraw and deposit columns!".into(),
        ));
    }

    // strptime is silly and wants a &'static format string
    let date_format = Box::leak(csv_spec.date_format.clone().into_boxed_str());
    let date_parser = strptime::Parser::new(date_format);

    let mut txns = vec![];

    for record in csv_reader.records() {
        let record = record?;

        let mut txn_datestamp: Option<data::Datestamp> = None;
        let mut txn_title: Option<String> = None;
        let mut txn_change: Option<data::Decimal> = None;
        let mut txn_balance: Option<data::Decimal> = None;
        let mut txn_unit: Option<data::UnitName> = None;

        for (record, spec) in record.iter().zip(csv_spec.cols.iter()) {
            match spec {
                data::spec::CsvColumnSpec::Ignore => (),
                data::spec::CsvColumnSpec::Datestamp => {
                    let date = date_parser.parse(record)?.date()?;
                    txn_datestamp = Some(data::Datestamp {
                        year: date.year() as u16,
                        month: date.month(),
                        day: date.day(),
                    });
                }
                data::spec::CsvColumnSpec::Title => {
                    txn_title = Some(record.into());
                }
                data::spec::CsvColumnSpec::Deposit => {
                    if record.trim().is_empty() {
                        continue;
                    }

                    txn_change = Some(try_parse_decimal(record)?);
                }
                data::spec::CsvColumnSpec::Withdraw => {
                    if record.trim().is_empty() {
                        continue;
                    }
                    let mut dec = try_parse_decimal(record)?;
                    dec.set_sign_negative(true);
                    txn_change = Some(dec);
                }
                data::spec::CsvColumnSpec::Change => {
                    if record.trim().is_empty() {
                        continue;
                    }

                    txn_change = Some(try_parse_decimal(record)?);
                }
                data::spec::CsvColumnSpec::Balance => {
                    if record.trim().is_empty() {
                        continue;
                    }

                    txn_balance = Some(try_parse_decimal(record)?);
                }
                data::spec::CsvColumnSpec::Unit => {
                    txn_unit = Some(data::UnitName::new(record));
                }
            }
        }

        txns.push(data::Transaction {
            datestamp: txn_datestamp.unwrap(),
            title: txn_title,
            annotations: vec![],
            changes: vec![
                data::Change {
                    account: account.into(),
                    amount: txn_change.unwrap().into(),
                    balance: txn_balance.map(Into::into),
                    unit: txn_unit.or(aspec.unit).map(Into::into).unwrap(),
                    source: None,
                },
                data::Change {
                    account: target.into(),
                    amount: data::Decimal::ZERO
                        .checked_sub(txn_change.unwrap())
                        .unwrap()
                        .into(),
                    balance: None,
                    unit: txn_unit.or(aspec.unit).map(Into::into).unwrap(),
                    source: None,
                },
            ],
            source: None,
        });
    }

    Ok(txns)
}

fn recursive_order_search(
    txns: &[data::Transaction],
    account: data::AccountName,
    order: &mut Vec<usize>,
    remaining: &mut BTreeSet<usize>,
) -> bool {
    if remaining.is_empty() {
        return true;
    }

    let possibles = remaining.iter().cloned().collect::<Vec<_>>();

    if let Some(last) = order.last() {
        let Some(last_balance) = txns[*last].change_for(account).unwrap().balance.as_deref() else {
            return false;
        };
        for possible in possibles.iter() {
            // check if balances line up
            let change = txns[*possible].change_for(account).unwrap();
            let Some(nbal) = change.balance else { continue };

            if last_balance.checked_add(*change.amount) != Some(*nbal) {
                continue;
            }

            remaining.remove(possible);
            order.push(*possible);
            if recursive_order_search(txns, account, order, remaining) {
                return true;
            }
            order.pop();
            remaining.insert(*possible);
        }
    } else {
        for possible in possibles.into_iter() {
            remaining.remove(&possible);
            order.push(possible);
            if recursive_order_search(txns, account, order, remaining) {
                return true;
            }
            order.pop();
            remaining.insert(possible);
        }
    }

    false
}

fn postprocess(account: data::AccountName, transactions: &mut Vec<data::Transaction>) {
    // check if we're sorted by datestamp already
    if transactions.is_sorted_by_key(|tx| tx.datestamp) {
        // already vaguely in the right order
    } else if transactions
        .iter()
        .rev()
        .is_sorted_by_key(|tx| tx.datestamp)
    {
        // reverse everything
        transactions.reverse();
    } else {
        // otherwise try to get things vaguely sorted
        transactions.sort_by_key(|tx| tx.datestamp);
    }

    let mut to_assign = BTreeSet::from_iter(0..transactions.len());
    let mut order = vec![];

    if !recursive_order_search(transactions, account, &mut order, &mut to_assign) {
        log::warn!("Unable to determine transaction ordering!");
        return;
    }

    let mut ntransact = order
        .iter()
        .map(|v| transactions[*v].clone())
        .collect::<Vec<_>>();

    std::mem::swap(&mut ntransact, transactions);
}

pub fn import_from(
    aspec: &data::spec::AccountSpec,
    account: data::AccountName,
    target: data::AccountName,
    path: &std::path::Path,
) -> Result<Vec<data::Transaction>, ImportError> {
    let reader = std::fs::File::open(path)?;

    let mut output = match &aspec.import {
        Some(data::spec::ImportFileFormat::Csv(csv)) => {
            import_from_csv(csv, aspec, account, target, reader)
        }
        None => Err(ImportError::ConfigError(format!(
            "no import configuration for {account}"
        ))),
    }?;

    postprocess(account, &mut output);

    Ok(output)
}