use brk_cohort::AmountBucket;
use brk_error::Result;
use brk_types::{Cents, Sats, TypeIndex};
use vecdb::VecIndex;

use crate::distribution::{
    addr::{AddrMetricsState, AddrSendPreState, HeightToAddrTypeToVec},
    cohorts::AddrCohorts,
};

use super::{super::cache::AddrLookup, transfer_address_cache::TransferAddressCache};

/// Process sent UTXOs for address cohort membership and empty-address transitions.
pub(crate) fn process_sent(
    sent_data: HeightToAddrTypeToVec<(TypeIndex, Sats)>,
    cohorts: &mut AddrCohorts,
    lookup: &mut AddrLookup<'_>,
    current_price: Cents,
    state: &mut AddrMetricsState,
    addresses: &mut TransferAddressCache,
    height_to_price: &[Cents],
) -> Result<()> {
    for (receive_height, by_type) in sent_data.into_iter() {
        let prev_price = height_to_price[receive_height.to_usize()];

        for (output_type, vec) in by_type.unwrap().into_iter() {
            let (type_received, type_seen) = addresses.sets_for(output_type);

            for (type_index, value) in vec {
                let addr_data = lookup.get_for_send(output_type, type_index);
                let pre = AddrSendPreState::capture(addr_data, output_type);

                let prev_balance = addr_data.balance();
                let is_first_encounter = type_seen.insert(type_index);
                let also_received = type_received.is_some_and(|s| s.contains(&type_index));
                let will_be_empty = addr_data.has_1_utxos();

                let prev_bucket = AmountBucket::from(prev_balance);
                let cohort_state = cohorts
                    .amount_range
                    .get_mut_by_bucket(prev_bucket)
                    .state
                    .as_mut()
                    .unwrap();

                // Mutates addr_data.spent_txo_count (+= 1). on_send_applied reads the post-spend view.
                cohort_state.send(addr_data, value, current_price, prev_price)?;
                let new_bucket = AmountBucket::from(addr_data.balance());
                let crossing_boundary = prev_bucket != new_bucket;
                state.on_send_applied(
                    output_type,
                    addr_data,
                    &pre,
                    is_first_encounter,
                    also_received,
                    will_be_empty,
                );

                if will_be_empty || crossing_boundary {
                    cohort_state.subtract(addr_data);
                }
                if will_be_empty {
                    lookup.move_to_empty(output_type, type_index);
                } else if crossing_boundary {
                    cohorts
                        .amount_range
                        .get_mut_by_bucket(new_bucket)
                        .state
                        .as_mut()
                        .unwrap()
                        .add(addr_data);
                }
            }
        }
    }

    Ok(())
}
