use anyhow::{bail, Result};

use crate::DEPTH;

use crate::data::FreeVar;
use crate::{
  data::{Context, ContextEntry, Term, Type},
  gensym::gensym_prefix,
};

// Figure 8. Applying a context, as a substitution, to a type

#[trace]
pub fn app_ctx(ctx: &Context, ty: &Type) -> Result<Type> {
  match ty {
    Type::Existential(a) => match ctx.lookup_existential(a) {
      Some(Some(m)) => Ok(m.into_poly()),
      Some(None) => Ok(Type::Existential(a.clone())),
      None => bail!("existential variable {a} doesn't exist in context"),
    },

    Type::Polytype(_, _) => todo!(),
    Type::Arrow(a, b) => Ok(Type::Arrow(
      Box::new(app_ctx(ctx, a)?),
      Box::new(app_ctx(ctx, b)?),
    )),
    _ => todo!("shiiet"),
  }
}

// Figure 9. Algorithmic subtyping

/// Under input context Γ , type A is a subtype of B, with output context ∆
#[trace]
pub fn subtype(ctx: &Context, a: &Type, b: &Type) -> Result<Context> {
  match (a, b) {
    // <:Unit rule
    (Type::Unit, Type::Unit) => Ok(ctx.clone()),

    // <:Var rule
    (Type::Var(x), Type::Var(y)) if x == y => {
      // Ensure that the name exists in the context
      if ctx.lookup_type(x).is_none() {
        bail!("name {x} not in context");
      }

      Ok(ctx.clone())
    }

    // <:Exvar rule
    (Type::Existential(x), Type::Existential(y)) if x == y => {
      // Ensure that the name exists in the context
      if ctx.lookup_existential(x).is_none() {
        bail!("name {x} not in context");
      }

      Ok(ctx.clone())
    }

    (Type::Existential(a), ty_a)
      if !b.free_vars().contains(&FreeVar::Existential(a.to_string()))
        && ctx.lookup_existential(a).is_some() =>
    {
      let ctx_delta = instantiate_left(ctx, a, ty_a)?;
      Ok(ctx_delta)
    }
    (Type::Existential(_), Type::Unit) => todo!(),
    (Type::Existential(_), Type::Var(_)) => todo!(),
    (Type::Existential(_), Type::Polytype(_, _)) => todo!(),
    (Type::Existential(_), Type::Arrow(_, _)) => todo!(),
    (Type::Polytype(_, _), Type::Unit) => todo!(),
    (Type::Polytype(_, _), Type::Var(_)) => todo!(),
    (Type::Polytype(_, _), Type::Existential(_)) => todo!(),
    (Type::Polytype(_, _), Type::Polytype(_, _)) => todo!(),
    (Type::Polytype(_, _), Type::Arrow(_, _)) => todo!(),
    (Type::Arrow(_, _), Type::Unit) => todo!(),
    (Type::Arrow(_, _), Type::Var(_)) => todo!(),
    (Type::Arrow(_, _), Type::Existential(_)) => todo!(),
    (Type::Arrow(_, _), Type::Polytype(_, _)) => todo!(),
    (Type::Arrow(_, _), Type::Arrow(_, _)) => todo!(),

    _ => bail!("subtyping relation failed between {a:?} and {b:?} (ctx = {ctx:?})"),
  }
}

// Figure 10. Instantiation

#[trace]
pub fn instantiate_left(ctx: &Context, a: &str, ty_a: &Type) -> Result<Context> {
  match ty_a {
    // InstLReach rule
    Type::Existential(b) if ctx.has_existential(a) && ctx.has_existential(b) => {
      // TODO: WTF?
      todo!()
    }
    Type::Existential(b) => todo!(),

    Type::Unit => todo!(),
    Type::Var(_) => todo!(),
    Type::Polytype(_, _) => todo!(),
    Type::Arrow(_, _) => todo!(),
  }
}

// Figure 11. Algorithmic typing

#[trace]
pub fn typecheck(ctx: &Context, term: &Term, ty: &Type) -> Result<Context> {
  match (term, ty) {
    // 1I rule
    (Term::Unit, Type::Unit) => Ok(ctx.clone()),

    // ∀I rule
    (e, Type::Polytype(x, tyA)) => todo!(),

    // →I rule
    (Term::Lam(x, e), Type::Arrow(ty_a, ty_b)) => {
      let mut aug_ctx = ctx.clone();
      todo!()
    }

    // Sub rule
    (term, ty) => {
      let (ty_a, ctx_theta) = synthesize(ctx, term)?;
      let a = app_ctx(&ctx_theta, &ty_a)?;
      let b = app_ctx(&ctx_theta, ty)?;
      let ctx_delta = subtype(&ctx_theta, &a, &b)?;
      Ok(ctx_delta)
    }
  }
}

#[trace]
pub fn synthesize(ctx: &Context, term: &Term) -> Result<(Type, Context)> {
  match term {
    // Var rule
    Term::Var(name) => {
      let ty = match ctx.lookup_type(name) {
        Some(v) => v,
        None => bail!("could not find name {name}"),
      };
      Ok((ty, ctx.clone()))
    }

    // 1I⇒ rule
    Term::Unit => Ok((Type::Unit, ctx.clone())),

    // Anno rule
    Term::Annot(e, ty_a) => {
      // Make sure the type is well formed
      // TODO: Have an actual "well-formed" check

      let ctx_delta = typecheck(ctx, &e, ty_a)?;
      Ok((ty_a.clone(), ctx_delta))
    }

    // →I⇒ rule
    // > Rule →I⇒ corresponds to Decl→I⇒, one of the guessing rules, so we
    // > create new existential variables ^α (for the function domain) and ^β (for
    // > the codomain) and check the function body against ^β.  As in ∀App, we do
    // > not place a marker before ^α, because ^α and ^β appear in the output type
    // > (λx. e ⇒ ^α → ^β).
    Term::Lam(x, e) => {
      let ex_a_s = gensym_prefix("ex");
      let ex_b_s = gensym_prefix("ex");
      let ex_a = Type::Existential(ex_a_s.clone());
      let ex_b = Type::Existential(ex_b_s.clone());
      let aug_ctx = ctx.add(vec![
        ContextEntry::ExistentialVar(ex_a_s.clone()),
        ContextEntry::ExistentialVar(ex_b_s.clone()),
        ContextEntry::TermAnnot(x.clone(), ex_a.clone()),
      ]);

      let wtf_ctx = typecheck(&aug_ctx, &e, &ex_b)?;
      println!("the wtf ctx is {wtf_ctx:?}");
      Ok((Type::Arrow(Box::new(ex_a), Box::new(ex_b)), ctx.clone()))
    }

    // →E rule
    Term::App(e1, e2) => {
      let (ty_a, ctx_theta) = synthesize(ctx, e1)?;
      let app_a = app_ctx(&ctx_theta, &ty_a)?;
      let (ty_c, ctx_delta) = app_synthesize(&ctx_theta, &app_a, &e2)?;
      Ok((ty_c, ctx_delta))
    }
  }
}

#[trace]
pub fn app_synthesize(ctx: &Context, fun_ty: &Type, term: &Term) -> Result<(Type, Context)> {
  match (fun_ty, term) {
    // →App rule
    (Type::Arrow(ty_a, ty_c), e) => {
      typecheck(ctx, e, ty_a)?;
      Ok((*ty_c.clone(), ctx.clone()))
    }

    // ∀App rule
    (Type::Polytype(_, _), _) => todo!(),

    // âApp rule
    (Type::Existential(_), _) => todo!(),

    _ => bail!("trying to appSynthesize with a non-function type"),
  }
}