namespace Mlang

abbrev Name := String

inductive Effect where
  | error : Effect
  deriving Repr, Inhabited, BEq

abbrev Effects := List Effect

def Effects.contains (effects : Effects) (effect : Effect) : Bool :=
  effects.any (· == effect)

def Effects.insert (effects : Effects) (effect : Effect) : Effects :=
  if effects.contains effect then effects else effect :: effects

def Effects.union (lhs rhs : Effects) : Effects :=
  rhs.foldl Effects.insert lhs

def Effects.erase (effects : Effects) (target : Effect) : Effects :=
  effects.filter (· != target)

def Effects.render : Effects → String
  | [] => "{}"
  | effects =>
      let names := effects.reverse.map (fun
        | .error => "Error")
      "{" ++ String.intercalate ", " names ++ "}"

inductive Ty where
  | int : Ty
  | bool : Ty
  | funTy : Ty → Effects → Ty → Ty
  deriving Repr, Inhabited, BEq

inductive Expr where
  | int : Int → Expr
  | bool : Bool → Expr
  | var : Name → Expr
  | add : Expr → Expr → Expr
  | sub : Expr → Expr → Expr
  | mul : Expr → Expr → Expr
  | div : Expr → Expr → Expr
  | eq : Expr → Expr → Expr
  | letE : Name → Expr → Expr → Expr
  | ifE : Expr → Expr → Expr → Expr
  | tryE : Expr → Expr → Expr
  | lam : Name → Ty → Expr → Expr
  | app : Expr → Expr → Expr
  deriving Repr, Inhabited

inductive Value where
  | int : Int → Value
  | bool : Bool → Value
  | closure : Name → Expr → List (Name × Value) → Value
  deriving Repr, Inhabited

abbrev Env := List (Name × Value)

inductive RuntimeError where
  | unboundVariable : Name → RuntimeError
  | typeError : String → RuntimeError
  | divisionByZero : RuntimeError
  deriving Repr, Inhabited

abbrev EvalM := Except RuntimeError

def Env.lookup (env : Env) (name : Name) : Option Value :=
  match env with
  | [] => none
  | (n, v) :: rest => if n = name then some v else rest.lookup name

def expectInt : Value → EvalM Int
  | .int n => .ok n
  | _ => .error (.typeError "expected an integer")

def expectBool : Value → EvalM Bool
  | .bool b => .ok b
  | _ => .error (.typeError "expected a boolean")

partial def eval (env : Env) : Expr → EvalM Value
  | .int n => .ok (.int n)
  | .bool b => .ok (.bool b)
  | .var name =>
      match env.lookup name with
      | some v => .ok v
      | none => .error (.unboundVariable name)
  | .add lhs rhs => do
      let l ← expectInt (← eval env lhs)
      let r ← expectInt (← eval env rhs)
      pure (.int (l + r))
  | .sub lhs rhs => do
      let l ← expectInt (← eval env lhs)
      let r ← expectInt (← eval env rhs)
      pure (.int (l - r))
  | .mul lhs rhs => do
      let l ← expectInt (← eval env lhs)
      let r ← expectInt (← eval env rhs)
      pure (.int (l * r))
  | .div lhs rhs => do
      let l ← expectInt (← eval env lhs)
      let r ← expectInt (← eval env rhs)
      if r = 0 then
        .error .divisionByZero
      else
        pure (.int (l / r))
  | .eq lhs rhs => do
      let l ← expectInt (← eval env lhs)
      let r ← expectInt (← eval env rhs)
      pure (.bool (l = r))
  | .letE name value body => do
      let value' ← eval env value
      eval ((name, value') :: env) body
  | .ifE cond thenBranch elseBranch => do
      let c ← expectBool (← eval env cond)
      if c then
        eval env thenBranch
      else
        eval env elseBranch
  | .tryE body handler =>
      match eval env body with
      | .ok value => .ok value
      | .error _ => eval env handler
  | .lam param _ body =>
      pure (.closure param body env)
  | .app fn arg => do
      let fnVal ← eval env fn
      let argVal ← eval env arg
      match fnVal with
      | .closure param body closureEnv =>
          eval ((param, argVal) :: closureEnv) body
      | _ =>
          .error (.typeError "expected a function")

def renderValue : Value → String
  | .int n => toString n
  | .bool b => toString b
  | .closure _ _ _ => "<closure>"

abbrev TyEnv := List (Name × Ty)

structure Judgment where
  ty : Ty
  effects : Effects
  deriving Repr, Inhabited

inductive TypeError where
  | unboundVariable : Name → TypeError
  | mismatch : Ty → Ty → TypeError
  | expectedFunction : Ty → TypeError
  deriving Repr, Inhabited

abbrev CheckM := Except TypeError

def TyEnv.lookup (env : TyEnv) (name : Name) : Option Ty :=
  match env with
  | [] => none
  | (n, ty) :: rest => if n = name then some ty else rest.lookup name

def ensureType (expected actual : Ty) : CheckM Unit :=
  if expected == actual then
    pure ()
  else
    throw (.mismatch expected actual)

def pureJudgment (ty : Ty) : Judgment :=
  { ty := ty, effects := [] }

partial def inferType (env : TyEnv) : Expr → CheckM Judgment
  | .int _ => pure (pureJudgment .int)
  | .bool _ => pure (pureJudgment .bool)
  | .var name =>
      match env.lookup name with
      | some ty => pure (pureJudgment ty)
      | none => throw (.unboundVariable name)
  | .add lhs rhs
  | .sub lhs rhs
  | .mul lhs rhs => do
      let lhsJ ← inferType env lhs
      let rhsJ ← inferType env rhs
      ensureType .int lhsJ.ty
      ensureType .int rhsJ.ty
      pure { ty := .int, effects := lhsJ.effects.union rhsJ.effects }
  | .div lhs rhs => do
      let lhsJ ← inferType env lhs
      let rhsJ ← inferType env rhs
      ensureType .int lhsJ.ty
      ensureType .int rhsJ.ty
      pure {
        ty := .int
        effects := (lhsJ.effects.union rhsJ.effects).insert .error
      }
  | .eq lhs rhs => do
      let lhsJ ← inferType env lhs
      let rhsJ ← inferType env rhs
      ensureType .int lhsJ.ty
      ensureType .int rhsJ.ty
      pure { ty := .bool, effects := lhsJ.effects.union rhsJ.effects }
  | .letE name value body => do
      let valueJ ← inferType env value
      let bodyJ ← inferType ((name, valueJ.ty) :: env) body
      pure {
        ty := bodyJ.ty
        effects := valueJ.effects.union bodyJ.effects
      }
  | .ifE cond thenBranch elseBranch => do
      let condJ ← inferType env cond
      ensureType .bool condJ.ty
      let thenJ ← inferType env thenBranch
      let elseJ ← inferType env elseBranch
      ensureType thenJ.ty elseJ.ty
      pure {
        ty := thenJ.ty
        effects := condJ.effects.union (thenJ.effects.union elseJ.effects)
      }
  | .tryE body handler => do
      let bodyJ ← inferType env body
      let handlerJ ← inferType env handler
      ensureType bodyJ.ty handlerJ.ty
      pure {
        ty := bodyJ.ty
        effects := (bodyJ.effects.erase .error).union handlerJ.effects
      }
  | .lam param paramTy body => do
      let bodyJ ← inferType ((param, paramTy) :: env) body
      pure (pureJudgment (.funTy paramTy bodyJ.effects bodyJ.ty))
  | .app fn arg => do
      let fnJ ← inferType env fn
      let argJ ← inferType env arg
      match fnJ.ty with
      | .funTy paramTy latentEffects resultTy =>
          ensureType paramTy argJ.ty
          pure {
            ty := resultTy
            effects := fnJ.effects.union (argJ.effects.union latentEffects)
          }
      | _ =>
          throw (.expectedFunction fnJ.ty)

def renderType : Ty → String
  | .int => "Int"
  | .bool => "Bool"
  | .funTy lhs effects rhs => s!"({renderType lhs} -> {renderType rhs} ! {effects.render})"

def renderJudgment (judgment : Judgment) : String :=
  s!"{renderType judgment.ty} ! {judgment.effects.render}"

def sampleProgram : Expr :=
  .letE "inc" (.lam "x" .int (.add (.var "x") (.int 1)))
    (.app (.var "inc") (.int 41))

def lexicalScopeProgram : Expr :=
  .letE "x" (.int 10)
    (.letE "f" (.lam "y" .int (.add (.var "x") (.var "y")))
      (.letE "x" (.int 100)
        (.app (.var "f") (.int 5))))

def conditionalProgram : Expr :=
  .ifE (.eq (.mul (.int 6) (.int 7)) (.int 42))
    (.int 1)
    (.int 0)

end Mlang