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

use bone_kernel::{BrepError, BrepSolid, ExtrudeFeature};
use bone_solver::SolverError;
use bone_types::{FeatureId, GeometryGeneration};

use crate::Sketch;
use crate::profile::build_profile;

#[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, thiserror::Error)]
pub enum ExtrudeError {
    #[error(transparent)]
    UnsolvedSketch(SolverError),
    #[error(transparent)]
    Kernel(BrepError),
}

#[derive(Clone)]
pub struct EvaluatedExtrude {
    result: Result<BrepSolid, ExtrudeError>,
    generation: Option<GeometryGeneration>,
}

impl EvaluatedExtrude {
    pub fn result(&self) -> &Result<BrepSolid, ExtrudeError> {
        &self.result
    }

    #[must_use]
    pub fn solid(&self) -> Option<&BrepSolid> {
        self.result.as_ref().ok()
    }

    #[must_use]
    pub fn generation(&self) -> Option<GeometryGeneration> {
        self.generation
    }
}

impl core::fmt::Debug for EvaluatedExtrude {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let mut state = f.debug_struct("EvaluatedExtrude");
        state.field("generation", &self.generation);
        match &self.result {
            Ok(solid) => state.field("faces", &solid.iter_faces().count()),
            Err(error) => state.field("error", error),
        };
        state.finish()
    }
}

#[must_use]
pub fn evaluate_extrude(
    extrude: FeatureId,
    sketch: &EvaluatedSketch,
    feature: &ExtrudeFeature,
) -> EvaluatedExtrude {
    let result = match sketch {
        EvaluatedSketch::Solved(solved) => build_profile(solved)
            .and_then(|profile| bone_kernel::evaluate_extrude(extrude, &profile, feature))
            .map_err(ExtrudeError::Kernel),
        EvaluatedSketch::Failed(error) => Err(ExtrudeError::UnsolvedSketch(error.clone())),
    };
    let generation = result
        .as_ref()
        .ok()
        .map(|solid| GeometryGeneration::from_solid_key(solid.content_key()));
    EvaluatedExtrude { result, generation }
}

#[derive(Clone, Debug, Default)]
pub struct FeatureCache {
    sketches: BTreeMap<FeatureId, Cached<Sketch, EvaluatedSketch>>,
    extrudes: BTreeMap<FeatureId, Cached<(EvaluatedSketch, ExtrudeFeature), EvaluatedExtrude>>,
}

#[derive(Clone, Debug)]
struct Cached<I, O> {
    input: I,
    output: O,
}

fn memoize<I, O>(
    cache: &mut BTreeMap<FeatureId, Cached<I, O>>,
    feature: FeatureId,
    matches: impl FnOnce(&I) -> bool,
    build_input: impl FnOnce() -> I,
    eval: impl FnOnce(&I) -> O,
) -> &O {
    let cached = match cache.entry(feature) {
        Entry::Occupied(mut slot) => {
            if !matches(&slot.get().input) {
                let input = build_input();
                let output = eval(&input);
                *slot.get_mut() = Cached { input, output };
            }
            slot.into_mut()
        }
        Entry::Vacant(slot) => {
            let input = build_input();
            let output = eval(&input);
            slot.insert(Cached { input, output })
        }
    };
    &cached.output
}

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

    pub fn evaluate(&mut self, feature: FeatureId, input: &Sketch) -> &EvaluatedSketch {
        memoize(
            &mut self.sketches,
            feature,
            |cached| cached == input,
            || input.clone(),
            evaluate_sketch,
        )
    }

    pub fn evaluate_extrude(
        &mut self,
        feature: FeatureId,
        sketch: &EvaluatedSketch,
        extrude: &ExtrudeFeature,
    ) -> &EvaluatedExtrude {
        memoize(
            &mut self.extrudes,
            feature,
            |(cached_sketch, cached_extrude)| cached_sketch == sketch && cached_extrude == extrude,
            || (sketch.clone(), *extrude),
            |(sketch, extrude)| evaluate_extrude(feature, sketch, extrude),
        )
    }

    #[must_use]
    pub fn contains(&self, feature: FeatureId) -> bool {
        self.sketches.contains_key(&feature) || self.extrudes.contains_key(&feature)
    }

    pub fn invalidate(&mut self, feature: FeatureId) -> bool {
        let sketch = self.sketches.remove(&feature).is_some();
        let extrude = self.extrudes.remove(&feature).is_some();
        sketch || extrude
    }

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

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

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