use std::collections::{BTreeMap, BTreeSet, btree_map::Entry};

use bone_solver::SolverError;
use bone_types::FeatureId;

use crate::Sketch;

#[derive(Clone, Debug, PartialEq)]
pub enum EvaluatedSketch {
    Solved(Sketch),
    Failed(SolverError),
}

#[must_use]
pub fn evaluate_sketch(input: &Sketch) -> EvaluatedSketch {
    match input.solve() {
        Ok(solved) => EvaluatedSketch::Solved(solved),
        Err(error) => EvaluatedSketch::Failed(error),
    }
}

#[derive(Clone, Debug, Default)]
pub struct FeatureCache {
    entries: BTreeMap<FeatureId, CachedSketch>,
}

#[derive(Clone, Debug)]
struct CachedSketch {
    input: Sketch,
    output: EvaluatedSketch,
}

impl FeatureCache {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    pub fn evaluate(&mut self, feature: FeatureId, input: &Sketch) -> &EvaluatedSketch {
        let cached = match self.entries.entry(feature) {
            Entry::Occupied(mut slot) => {
                if &slot.get().input != input {
                    *slot.get_mut() = CachedSketch {
                        input: input.clone(),
                        output: evaluate_sketch(input),
                    };
                }
                slot.into_mut()
            }
            Entry::Vacant(slot) => slot.insert(CachedSketch {
                input: input.clone(),
                output: evaluate_sketch(input),
            }),
        };
        &cached.output
    }

    #[must_use]
    pub fn lookup(&self, feature: FeatureId) -> Option<&EvaluatedSketch> {
        self.entries.get(&feature).map(|cached| &cached.output)
    }

    pub fn invalidate(&mut self, feature: FeatureId) -> bool {
        self.entries.remove(&feature).is_some()
    }

    pub fn retain(&mut self, live: impl IntoIterator<Item = FeatureId>) {
        let keep: BTreeSet<FeatureId> = live.into_iter().collect();
        self.entries.retain(|id, _| keep.contains(id));
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }
}
