hoard/src/import.rs

use std::collections::HashMap;

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<Decimal, ImportError> {
    // remove any '$' or units from the string
    let filtered = from
        .chars()
        .filter(|f| char::is_digit(*f, 10) || *f == '.' || *f == '-')
        .collect::<String>();

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

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

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

    if !csv_spec.cols.contains(&CsvColumnSpec::Change)
        && (!csv_spec.cols.contains(&CsvColumnSpec::Withdraw)
            || !csv_spec.cols.contains(&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 unbalanced = AccountName::new("unbalanced");

    let mut txns = vec![];

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

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

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

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

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

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

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

    Ok(txns)
}

fn postprocess(account: AccountName, transactions: &mut Vec<Transaction>) {
    // check if we need to re-order transactions due to balances not lining up because of ordering
    let mut running_balances = HashMap::<UnitName, Decimal>::new();
    let mut idx = 0;

    // first get things vaguely sorted
    transactions.sort_by_key(|tx| tx.datestamp);

    let check_for_match = |running_balances: &mut HashMap<UnitName, Decimal>, change: &Change| {
        let bal = *change.balance.unwrap();
        match running_balances.entry(*change.unit) {
            std::collections::hash_map::Entry::Vacant(entry) => {
                entry.insert(bal);
                return true
            },
            std::collections::hash_map::Entry::Occupied(mut entry) => {
                let rbal = entry.get_mut();
                let new_rbal = rbal.checked_add(*change.amount).unwrap();
                if new_rbal != bal {
                    return false
                } else {
                    *rbal = new_rbal;
                    return true
                }
            },
        }
    };


    let mut removed : Vec<Transaction> = vec![];

    'outer: loop {
        for ridx in 0..removed.len() {
            if check_for_match(&mut running_balances, removed[ridx].change_for(account).unwrap()) {
                transactions.insert(idx, removed.remove(ridx));
                log::trace!("pulling transaction out of removed");
                idx += 1;
                continue 'outer
            }
        }

        if idx >= transactions.len() {
            break
        }

        let tx = &transactions[idx];
        let change = tx.change_for(account).unwrap();
        if change.balance.is_none() {
            idx += 1;
            continue
        };

        if check_for_match(&mut running_balances, change) {
            log::trace!("transaction is good! balance is now: {}", running_balances[&*change.unit]);
            idx += 1;
        } else {
            log::trace!("shifting transaction to removed");
            removed.push(transactions.remove(idx));
        }
    }

    if removed.len() > 0 {
        log::error!("Not all transactions are consistent!");
    }
}

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

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

    postprocess(target, &mut output);

    Ok(output)
}