use std::hash::Hash;

use im::{HashSet, Vector};

#[derive(Debug, Clone)]
pub enum Term {
  Unit,
  Var(String),
  Lam(String, Box<Term>),
  App(Box<Term>, Box<Term>),
  Annot(Box<Term>, Type),
}

#[derive(Debug, Clone)]
pub enum Type {
  Unit,
  Var(String),
  Existential(String),
  Polytype(String, Box<Type>),
  Arrow(Box<Type>, Box<Type>),
}

impl Type {
  pub fn into_mono(&self) -> Option<Monotype> {
    match self {
      Type::Unit => Some(Monotype::Unit),
      Type::Var(x) => Some(Monotype::Var(x.clone())),
      Type::Existential(x) => Some(Monotype::Existential(x.clone())),
      Type::Polytype(_, _) => None,
      Type::Arrow(a, b) => Some(Monotype::Arrow(
        Box::new(a.into_mono()?),
        Box::new(b.into_mono()?),
      )),
    }
  }

  #[inline]
  pub fn is_mono(&self) -> bool {
    self.into_mono().is_some()
  }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum FreeVar {
  Var(String),
  Existential(String),
}

impl Type {
  pub fn free_vars(&self) -> HashSet<FreeVar> {
    fn free_vars_with_context(ctx: &Context, ty: &Type) -> HashSet<FreeVar> {
      match ty {
        Type::Unit => HashSet::default(),
        Type::Var(x) if ctx.lookup_type(x).is_some() => HashSet::default(),
        Type::Var(x) => [FreeVar::Var(x.to_owned())].into_iter().collect(),
        Type::Existential(x) if ctx.lookup_existential(x).is_some() => HashSet::default(),
        Type::Existential(x) => [FreeVar::Existential(x.to_owned())].into_iter().collect(),
        Type::Polytype(x, ty_a) => {
          let new_ctx = ctx.add(vec![ContextEntry::ExistentialVar(x.to_owned())]);
          free_vars_with_context(&new_ctx, &ty_a)
        }
        Type::Arrow(ty_a, ty_b) => {
          let a_vars = free_vars_with_context(&ctx, &ty_a);
          let b_vars = free_vars_with_context(&ctx, &ty_b);
          a_vars.union(b_vars)
        }
      }
    }

    free_vars_with_context(&Context::default(), self)
  }
}

#[derive(Debug, Clone)]
pub enum Monotype {
  Unit,
  Var(String),
  Existential(String),
  Arrow(Box<Monotype>, Box<Monotype>),
}

impl Monotype {
  pub fn into_poly(&self) -> Type {
    match self {
      Monotype::Unit => Type::Unit,
      Monotype::Var(x) => Type::Var(x.clone()),
      Monotype::Existential(x) => Type::Existential(x.clone()),
      Monotype::Arrow(a, b) => Type::Arrow(Box::new(a.into_poly()), Box::new(b.into_poly())),
    }
  }
}

#[derive(Debug, Clone)]
pub enum ContextEntry {
  TypeVar(String),
  TermAnnot(String, Type),
  ExistentialVar(String),
  ExistentialSolved(String, Monotype),
  Marker(String),
}

#[derive(Debug, Clone)]
pub enum CompleteContextEntry {
  TypeVar(String),
  TermAnnot(String, Type),
  ExistentialSolved(String, Monotype),
  Marker(String),
}

#[derive(Debug, Clone, Default)]
pub struct Context(Vector<ContextEntry>);

impl Context {
  pub fn add(&self, entries: Vec<ContextEntry>) -> Context {
    let mut new_ctx = self.clone();
    for entry in entries {
      new_ctx.0.push_back(entry);
    }
    new_ctx
  }

  /// Looks up a polytype by name
  pub fn lookup_type(&self, name: impl AsRef<str>) -> Option<Type> {
    self.0.iter().find_map(|entry| match entry {
      ContextEntry::TermAnnot(n, t) if n == name.as_ref() => Some(t.clone()),
      _ => None,
    })
  }

  #[inline]
  pub fn has_type(&self, name: impl AsRef<str>) -> bool {
    self.lookup_type(name).is_some()
  }

  /// Looks up an existential variable by name
  /// - Returns Some(Some(t)) if solved
  /// - Returns Some(None) if exists but unsolved
  /// - Returns None if not found
  pub fn lookup_existential(&self, name: impl AsRef<str>) -> Option<Option<Monotype>> {
    self.0.iter().find_map(|entry| match entry {
      ContextEntry::ExistentialVar(n) if n == name.as_ref() => Some(None),
      ContextEntry::ExistentialSolved(n, t) if n == name.as_ref() => Some(Some(t.clone())),
      _ => None,
    })
  }

  #[inline]
  pub fn has_existential(&self, name: impl AsRef<str>) -> bool {
    self.lookup_existential(name).is_some()
  }
}