use std::cell::RefCell;
use std::fmt;
use std::rc::Rc;

use bound::Bound;
use hashbrown::HashMap;
use hashbrown::hash_map::Entry;

use crate::api;
use crate::expr::Expr;

#[derive(Default)]
pub struct ExprStoreData {
	exprs: RefCell<HashMap<api::ExprTicket, (u32, Expr)>>,
	parent: Option<ExprStore>,
}
#[derive(Clone, Default)]
pub struct ExprStore(Rc<ExprStoreData>);
impl ExprStore {
	/// If tracking_parent is false, get_expr can fall back to the parent if none
	/// is found here.
	///
	/// If tracking_parent is true, get_expr can still fall back to the parent,
	/// but operations on the parent can access the child exprs too until this
	/// store is dropped.
	#[must_use]
	pub fn derive(&self) -> Self {
		Self(Rc::new(ExprStoreData { exprs: RefCell::default(), parent: Some(self.clone()) }))
	}
	pub fn give_expr(&self, expr: Expr) {
		if let Some(parent) = self.0.parent.as_ref() {
			parent.give_expr(expr.clone())
		}
		match self.0.exprs.borrow_mut().entry(expr.id()) {
			Entry::Occupied(mut oe) => oe.get_mut().0 += 1,
			Entry::Vacant(v) => {
				v.insert((1, expr));
			},
		}
	}
	pub fn take_expr(&self, ticket: api::ExprTicket) -> Option<Expr> {
		if let Some(parent) = self.0.parent.as_ref() {
			parent.take_expr(ticket);
		}
		match self.0.exprs.borrow_mut().entry(ticket) {
			Entry::Vacant(_) => panic!("Attempted to double-take expression"),
			Entry::Occupied(oe) if oe.get().0 == 1 => Some(oe.remove().1),
			Entry::Occupied(mut oe) => {
				oe.get_mut().0 -= 1;
				Some(oe.get().1.clone())
			},
		}
	}
	#[must_use]
	pub fn get_expr(&self, ticket: api::ExprTicket) -> Option<Expr> {
		(self.0.exprs.borrow().get(&ticket).map(|(_, expr)| expr.clone()))
			.or_else(|| self.0.parent.as_ref()?.get_expr(ticket))
	}
	pub fn iter(&self) -> impl Iterator<Item = (u32, Expr)> {
		let r = Bound::new(self.clone(), |this| this.0.exprs.borrow());
		let mut iter = Bound::new(r, |r| r.values());
		std::iter::from_fn(move || iter.wrapped_mut().next().cloned())
	}
}
impl fmt::Display for ExprStore {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let r = self.0.exprs.borrow();
		let rc: u32 = r.values().map(|v| v.0).sum();
		write!(f, "Store holding {rc} refs to {} exprs", r.len())
	}
}
impl Drop for ExprStore {
	fn drop(&mut self) {
		if 1 < Rc::strong_count(&self.0) {
			return;
		}
		if let Some(parent) = self.0.parent.as_ref() {
			for (id, (count, _)) in self.0.exprs.borrow().iter() {
				for _ in 0..*count {
					parent.take_expr(*id);
				}
			}
		}
	}
}