use brk_types::{OutputType, Sats, TypeIndex};
use rustc_hash::FxHashSet;

type Address = (OutputType, TypeIndex);

#[derive(Default)]
pub(super) struct Candidate {
    values: FxHashSet<Sats>,
    input_addresses: FxHashSet<Address>,
    output_addresses: FxHashSet<Address>,
    zero_values: usize,
    address_reuse: bool,
}

impl Candidate {
    pub(super) fn clear(&mut self) {
        self.values.clear();
        self.input_addresses.clear();
        self.output_addresses.clear();
        self.zero_values = 0;
        self.address_reuse = false;
    }

    pub(super) fn add_input(
        &mut self,
        value: Sats,
        output_type: OutputType,
        type_index: TypeIndex,
    ) {
        self.add_value(value);
        if has_script_address(output_type) {
            self.address_reuse |= !self.input_addresses.insert((output_type, type_index));
        }
    }

    pub(super) fn add_output(
        &mut self,
        value: Sats,
        output_type: OutputType,
        type_index: TypeIndex,
    ) {
        self.add_value(value);
        if has_script_address(output_type) {
            let address = (output_type, type_index);
            self.address_reuse |=
                self.input_addresses.contains(&address) || !self.output_addresses.insert(address);
        }
    }

    pub(super) fn is_match(&self, input_count: usize, output_count: usize) -> bool {
        !self.address_reuse
            && self.values.len() + self.zero_values <= (input_count + output_count) / 2
    }

    fn add_value(&mut self, value: Sats) {
        if value.is_zero() {
            self.zero_values += 1;
        } else {
            self.values.insert(value);
        }
    }
}

fn has_script_address(output_type: OutputType) -> bool {
    matches!(
        output_type,
        OutputType::P2PK65
            | OutputType::P2PK33
            | OutputType::P2PKH
            | OutputType::P2SH
            | OutputType::P2WPKH
            | OutputType::P2WSH
            | OutputType::P2TR
            | OutputType::P2A
    )
}

#[cfg(test)]
mod tests {
    use brk_types::{OutputType, Sats, TypeIndex};

    use super::Candidate;

    #[test]
    fn repeated_values_match_without_address_reuse() {
        let mut candidate = Candidate::default();
        for index in 0usize..5 {
            candidate.add_input(
                Sats::new(10_000),
                OutputType::P2WPKH,
                TypeIndex::from(index),
            );
            candidate.add_output(
                Sats::new(9_000),
                OutputType::P2WPKH,
                TypeIndex::from(index + 10),
            );
        }
        assert!(candidate.is_match(5, 5));
    }

    #[test]
    fn reused_output_address_does_not_match() {
        let mut candidate = Candidate::default();
        for _ in 0..5 {
            candidate.add_output(
                Sats::new(9_000),
                OutputType::P2WPKH,
                TypeIndex::from(1usize),
            );
        }
        assert!(!candidate.is_match(5, 5));
    }

    #[test]
    fn reused_input_address_does_not_match() {
        let mut candidate = Candidate::default();
        for _ in 0..5 {
            candidate.add_input(
                Sats::new(10_000),
                OutputType::P2PK33,
                TypeIndex::from(1usize),
            );
        }
        assert!(!candidate.is_match(5, 5));
    }
}
