//! A ledger implementation to track all transactions.

use crate::{
    Chargeback, ClientId, Deposit, Dispute, LedgerError, Resolve, Transaction, TxAmount, TxId,
    Withdrawal,
};

/// A ledger of accounts, which processes transactions one at a time.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Ledger {
    accounts: std::collections::HashMap<ClientId, AccountInfo>,
    transaction_amounts: std::collections::HashMap<(ClientId, TxId), TxAmount>,
    transaction_state: std::collections::HashMap<(ClientId, TxId), TxState>,
}

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct AccountInfo {
    available_funds: TxAmount,
    held_funds: TxAmount,
    locked: bool,
}

/// Represent the state of a transaction. Here are the possible transitions:
///
/// ```graphviz
/// Processed -> Disputed
/// Disputed -> Resolved
/// Disputed -> ChargedBack
/// ```
///
/// The starting state is `Processed`.
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum TxState {
    /// A transaction was just accepted.
    Processed,
    /// A transaction dispute has been processed.
    Disputed,
    /// A transaction disputed was resolved.
    Resolved,
    /// A transaction disputed was chargedback.
    ChargedBack,
}

impl TxState {
    pub fn apply_dispute(
        &mut self,
        account: &mut AccountInfo,
        amount: TxAmount,
    ) -> LedgerResult<()> {
        if *self != Self::Processed {
            return Err(LedgerError::AlreadyDisputed);
        }

        account.apply_dispute(amount)?;
        *self = Self::Disputed;
        Ok(())
    }

    pub fn apply_resolution(
        &mut self,
        account: &mut AccountInfo,
        amount: TxAmount,
    ) -> LedgerResult<()> {
        if *self != Self::Disputed {
            return Err(LedgerError::NotDisputed);
        }

        account.apply_resolution(amount)?;
        *self = Self::Resolved;
        Ok(())
    }

    pub fn apply_chargeback(
        &mut self,
        account: &mut AccountInfo,
        amount: TxAmount,
    ) -> LedgerResult<()> {
        if *self != Self::Disputed {
            return Err(LedgerError::NotDisputed);
        }

        account.apply_chargeback(amount)?;
        *self = Self::ChargedBack;
        Ok(())
    }
}

type LedgerResult<T> = Result<T, LedgerError>;

impl Ledger {
    pub fn new() -> Self {
        Default::default()
    }

    /// Serialize a [Ledger] to CSV.
    pub fn dump_csv<W: std::io::Write>(&self, writer: &mut csv::Writer<W>) -> csv::Result<()> {
        // Keep list of accounts ordered for easier diffs
        let ordered_accounts: std::collections::BTreeMap<_, _> = self.accounts.iter().collect();
        writer.write_record(&["client", "available", "held", "total", "locked"])?;
        for (id, info) in ordered_accounts.into_iter() {
            writer.write_record(&[
                id.0.to_string(),
                info.available_funds().0.to_string(),
                info.held_funds().0.to_string(),
                info.total_funds().0.to_string(),
                info.is_locked().to_string(),
            ])?
        }
        Ok(())
    }

    pub fn process(&mut self, tx: Transaction) -> LedgerResult<()> {
        match tx {
            Transaction::Deposit(Deposit { client, tx, amount }) => self.delta(client, tx, amount),
            Transaction::Withdrawal(Withdrawal { client, tx, amount }) => {
                self.delta(client, tx, -amount)
            }
            Transaction::Dispute(tx) => self.dispute(tx),
            Transaction::Resolve(tx) => self.resolve(tx),
            Transaction::Chargeback(tx) => self.chargeback(tx),
        }
    }

    fn delta(&mut self, client: ClientId, tx: TxId, delta: TxAmount) -> LedgerResult<()> {
        let account = self.accounts.entry(client).or_default();
        account.apply_delta(delta)?;
        self.transaction_amounts.insert((client, tx), delta);
        self.transaction_state
            .insert((client, tx), TxState::Processed);
        Ok(())
    }

    fn dispute(&mut self, Dispute { client, tx }: Dispute) -> LedgerResult<()> {
        let (current_state, account, amount) = self.get_past_transaction_info(client, tx)?;
        current_state.apply_dispute(account, amount)
    }

    fn resolve(&mut self, Resolve { client, tx }: Resolve) -> LedgerResult<()> {
        let (current_state, account, amount) = self.get_past_transaction_info(client, tx)?;
        current_state.apply_resolution(account, amount)
    }

    fn chargeback(&mut self, Chargeback { client, tx }: Chargeback) -> LedgerResult<()> {
        let (current_state, account, amount) = self.get_past_transaction_info(client, tx)?;
        current_state.apply_chargeback(account, amount)
    }

    fn get_past_transaction_info(
        &mut self,
        client: ClientId,
        tx: TxId,
    ) -> LedgerResult<(&mut TxState, &mut AccountInfo, TxAmount)> {
        let current_state = self
            .transaction_state
            .get_mut(&(client, tx))
            .ok_or(LedgerError::UnknownTx(client, tx))?;
        let account = self
            .accounts
            .get_mut(&client)
            .expect("a processed transaction should have its account recorded");
        let amount = self
            .transaction_amounts
            .get(&(client, tx))
            .cloned()
            .expect("a processed transaction should have its amount recorded");
        Ok((current_state, account, amount))
    }
}

impl AccountInfo {
    /// Whether or not an account has been locked.
    pub fn is_locked(&self) -> bool {
        self.locked
    }

    /// The funds that are usable on this account.
    pub fn available_funds(&self) -> TxAmount {
        self.available_funds
    }

    /// The funds that have been locked pending resolution of dispute.
    pub fn held_funds(&self) -> TxAmount {
        self.held_funds
    }

    /// The totals funds on an account, i.e: available funds and held funds.
    pub fn total_funds(&self) -> TxAmount {
        self.available_funds + self.held_funds
    }

    pub fn apply_delta(&mut self, delta: TxAmount) -> LedgerResult<()> {
        self.check_frozen()?;
        let new_balance = self.available_funds() + delta;
        if new_balance < TxAmount::ZERO {
            return Err(LedgerError::NotEnoughFunds);
        }
        self.available_funds = new_balance;
        Ok(())
    }

    pub fn apply_dispute(&mut self, delta: TxAmount) -> LedgerResult<()> {
        self.check_frozen()?;
        // FIXME: should we check for negative funds?
        self.available_funds -= delta;
        self.held_funds += delta;
        Ok(())
    }

    pub fn apply_resolution(&mut self, delta: TxAmount) -> LedgerResult<()> {
        self.check_frozen()?;
        // FIXME: should we check for negative funds?
        self.available_funds += delta;
        self.held_funds -= delta;
        Ok(())
    }

    pub fn apply_chargeback(&mut self, delta: TxAmount) -> LedgerResult<()> {
        self.check_frozen()?;
        // FIXME: should we check for negative funds?
        self.held_funds -= delta;
        self.locked = true;
        Ok(())
    }

    fn check_frozen(&self) -> LedgerResult<()> {
        if self.is_locked() {
            Err(LedgerError::FrozenAccount)
        } else {
            Ok(())
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use expect_test::{expect, Expect};

    macro_rules! inline_csv {
        ($line:literal) => {
            $line
        };
        ($line:literal, $($lines:literal),+ $(,)?) => {
            concat!($line, "\n", inline_csv!($($lines),+))
        };
    }

    fn process_transactions(input: &str) -> Result<Ledger, LedgerError> {
        let mut ledger = Ledger::new();
        for tx in Transaction::configured_csv_reader_builder()
            .from_reader(input.as_bytes())
            .into_deserialize()
        {
            ledger.process(tx.unwrap())?
        }
        Ok(ledger)
    }

    fn check_ledger(ledger: &Ledger, expect: Expect) {
        let mut writer = csv::Writer::from_writer(vec![]);
        ledger.dump_csv(&mut writer).unwrap();
        let actual = String::from_utf8(writer.into_inner().unwrap()).unwrap();
        expect.assert_eq(&actual);
    }

    #[test]
    fn deposit_single_account() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "deposit,         1,  2,   2.0",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,3.0,0,3.0,false
            "#]],
        );
    }

    #[test]
    fn deposit_multiple_accounts() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "deposit,         2,  2,   1.0",
            "deposit,         1,  3,   2.0",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,3.0,0,3.0,false
                2,1.0,0,1.0,false
            "#]],
        );
    }

    #[test]
    fn deposit_and_withdrawal() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "deposit,         2,  2,   1.0",
            "deposit,         1,  3,   2.0",
            "withdrawal,      1,  4,   1.5",
            "withdrawal,      2,  5,   1.0",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,1.5,0,1.5,false
                2,0.0,0,0.0,false
            "#]],
        );
    }

    #[test]
    fn deposit_and_withdrawal_not_enough_funds() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         2,  2,   1.0",
            "withdrawal,      2,  5,   3.0",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::NotEnoughFunds);
    }

    #[test]
    fn dispute_deposit() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,0.0,1.0,1.0,false
            "#]],
        );
    }

    #[test]
    fn dispute_withdrawal() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "withdrawal,      1,  2,   1.0",
            "dispute,         1,  2",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,1.0,-1.0,0.0,false
            "#]],
        );
    }

    #[test]
    fn resolve_dispute() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "resolve,         1,  1",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,1.0,0.0,1.0,false
            "#]],
        );
    }

    #[test]
    fn chargeback_dispute() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "chargeback,      1,  1",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,0.0,0.0,0.0,true
            "#]],
        );
    }

    #[test]
    fn dispute_then_withdrawal() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "withdrawal,      1,  2,   1.0",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::NotEnoughFunds);
    }

    #[test]
    fn dispute_then_resolution() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "resolve,         1,  1",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,1.0,0.0,1.0,false
            "#]],
        );
    }

    #[test]
    fn dispute_then_chargeback() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "chargeback,      1,  1",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,0.0,0.0,0.0,true
            "#]],
        );
    }

    #[test]
    fn dispute_then_resolution_then_withdrawal() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "resolve,         1,  1",
            "withdrawal,      1,  2,   1.0",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
                client,available,held,total,locked
                1,0.0,0.0,0.0,false
            "#]],
        );
    }

    #[test]
    fn dispute_then_chargeback_then_withdrawal() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "chargeback,      1,  1",
            "withdrawal,      1,  2,   1.0",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::FrozenAccount);
    }

    #[test]
    fn dispute_then_deposit_then_dispute_then_resolve_then_chargeback() {
        let ledger = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "deposit,         1,  2,   1.0",
            "dispute,         1,  2",
            "resolve,         1,  2",
            "chargeback,      1,  1",
        ))
        .unwrap();
        check_ledger(
            &ledger,
            expect![[r#"
            client,available,held,total,locked
            1,1.0,0.0,1.0,true
        "#]],
        );
    }

    #[test]
    fn unknown_client_and_transaction() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         2,  2",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::UnknownTx(ClientId(2), TxId(2)));
    }

    #[test]
    fn unknown_client() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         2,  1",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::UnknownTx(ClientId(2), TxId(1)));
    }

    #[test]
    fn unknown_transaction() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  2",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::UnknownTx(ClientId(1), TxId(2)));
    }

    #[test]
    fn double_dispute() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "dispute,         1,  1",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::AlreadyDisputed);
    }

    #[test]
    fn dispute_after_resolution() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "resolve,         1,  1",
            "dispute,         1,  1",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::AlreadyDisputed);
    }

    #[test]
    fn dispute_after_chargeback() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "dispute,         1,  1",
            "chargeback,      1,  1",
            "dispute,         1,  1",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::AlreadyDisputed);
    }

    #[test]
    fn resolution_no_dispute() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "resolve,         1,  1",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::NotDisputed);
    }

    #[test]
    fn chargeback_no_dispute() {
        let error = process_transactions(inline_csv!(
            "type,       client, tx, amount",
            "deposit,         1,  1,   1.0",
            "chargeback,      1,  1",
        ))
        .unwrap_err();
        assert_eq!(error, LedgerError::NotDisputed);
    }
}