import Mlang.Nickel
import Mlang.Http
import Mlang.NickelHost

namespace Mlang

open Nickel

abbrev Name := String

inductive Effect where
  | error : Effect
  | io : 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"
        | .io => "IO")
      "{" ++ String.intercalate ", " names ++ "}"

inductive Ty where
  | int : Ty
  | bool : Ty
  | string : Ty
  | data : Ty
  | dataInt : Ty
  | dataBool : Ty
  | dataString : Ty
  | dataNull : Ty
  | dataArray : Ty → Ty
  | dataRecord : List (String × Ty) → Ty
  | list : Ty
  | result : Ty
  | error : Ty
  | funTy : Ty → Effects → Ty → Ty
  deriving Repr, Inhabited, BEq

inductive Data where
  | null : Data
  | bool : Bool → Data
  | int : Int → Data
  | string : String → Data
  | array : List Data → Data
  | record : List (String × Data) → Data
  deriving Repr, Inhabited, BEq

inductive Expr where
  | int : Int → Expr
  | bool : Bool → Expr
  | string : String → Expr
  | null : Expr
  | var : Name → Expr
  | arrayE : List Expr → Expr
  | recordE : List (String × Expr) → 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 → Name → Expr → Expr
  | parMapE : Name → Expr → Expr → Expr
  | lam : Name → Ty → Expr → Expr
  | app : Expr → Expr → Expr
  deriving Repr, Inhabited

inductive Value where
  | int : Int → Value
  | bool : Bool → Value
  | string : String → Value
  | data : Data → Value
  | list : List Value → Value
  | resultOk : Value → Value
  | resultErr : String → Value
  | error : String → Value
  | builtinReadFile : Value
  | builtinHttpGet : Value
  | builtinReadFileNickel : Value
  | builtinParseNickel : Value
  | builtinParseJson : Value
  | builtinDataGet : Value
  | builtinDataGetField : Data → Value
  | builtinDataAt : Value
  | builtinDataAtIndex : Data → Value
  | builtinDataAsInt : Value
  | builtinDataAsString : Value
  | builtinDataAsBool : Value
  | builtinDataToJson : Value
  | builtinDataToNickel : 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
  | ioError : String → RuntimeError
  deriving Repr, Inhabited

abbrev EvalM := EIO 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 => pure n
  | _ => throw (.typeError "expected an integer")

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

def expectString : Value → EvalM String
  | .string s => pure s
  | _ => throw (.typeError "expected a string")

def expectData : Value → EvalM Data
  | .data term => pure term
  | _ => throw (.typeError "expected a data value")

def valueToData : Value → EvalM Data
  | .data term => pure term
  | .int n => pure (.int n)
  | .bool b => pure (.bool b)
  | .string s => pure (.string s)
  | _ => throw (.typeError "data literals only support Int, Bool, String, and Data values")

partial def dataType : Data → Ty
  | .null => .dataNull
  | .bool _ => .dataBool
  | .int _ => .dataInt
  | .string _ => .dataString
  | .array [] => .dataArray .data
  | .array (x :: xs) =>
      let itemTy := dataType x
      if xs.all (fun item => dataType item == itemTy) then
        .dataArray itemTy
      else
        .dataArray .data
  | .record fields =>
      .dataRecord (fields.map (fun (name, value) => (name, dataType value)))

def dataRecordField? : List (String × Ty) → String → Option Ty
  | [], _ => none
  | (name, ty) :: rest, target =>
      if name = target then some ty else dataRecordField? rest target

def liftScalarDataTy : Ty → Ty
  | .dataInt => .int
  | .dataBool => .bool
  | .dataString => .string
  | ty => ty

partial def refineKnownDataTy (ty : Ty) : Ty :=
  match ty with
  | .dataInt => .int
  | .dataBool => .bool
  | .dataString => .string
  | .dataNull => .data
  | .dataArray itemTy => .dataArray (refineKnownDataTy itemTy)
  | .dataRecord fields =>
      .dataRecord (fields.map (fun (name, fieldTy) => (name, refineKnownDataTy fieldTy)))
  | other => other

def dataArrayGet? : List Data → Nat → Option Data
  | [], _ => none
  | x :: _, 0 => some x
  | _ :: xs, n + 1 => dataArrayGet? xs n

def runtimeErrorToValue : RuntimeError → Value
  | .unboundVariable name => .error s!"unbound variable: {name}"
  | .typeError msg => .error s!"type error: {msg}"
  | .divisionByZero => .error "division by zero"
  | .ioError msg => .error s!"io error: {msg}"

def runtimeErrorToMessage : RuntimeError → String
  | .unboundVariable name => s!"unbound variable: {name}"
  | .typeError msg => s!"type error: {msg}"
  | .divisionByZero => "division by zero"
  | .ioError msg => s!"io error: {msg}"

partial def Data.ofNickel : Nickel.Term → Data
  | .null => .null
  | .bool b => .bool b
  | .int n => .int n
  | .string s => .string s
  | .array xs => .array (xs.map Data.ofNickel)
  | .record fields => .record (fields.map (fun (k, v) => (k, Data.ofNickel v)))

def parseRenderedData (rendered : String) : EvalM Data := do
  match Nickel.parse rendered with
  | .ok term => pure (Data.ofNickel term)
  | .error err => throw (.typeError s!"host Nickel render parse error: {err}")

def encodeMapResult : Except RuntimeError Value → Value
  | .ok value => .resultOk value
  | .error err => .resultErr (runtimeErrorToMessage err)

def escapeJsonChar : Char → String
  | '"' => "\\\""
  | '\\' => "\\\\"
  | '\n' => "\\n"
  | '\r' => "\\r"
  | '\t' => "\\t"
  | c => c.toString

def renderJsonString (s : String) : String :=
  "\"" ++ String.join (s.toList.map escapeJsonChar) ++ "\""

def isNickelKeyStart (c : Char) : Bool :=
  c.isAlpha || c = '_'

def isNickelKeyContinue (c : Char) : Bool :=
  c.isAlpha || c.isDigit || c = '_' || c = '-'

def renderNickelKey (s : String) : String :=
  match s.toList with
  | [] => renderJsonString s
  | c :: cs =>
      if isNickelKeyStart c && cs.all isNickelKeyContinue then
        s
      else
        renderJsonString s

partial def renderNickelData : Data → String
  | .null => "null"
  | .bool b => toString b
  | .int n => toString n
  | .string s => renderJsonString s
  | .array xs => "[" ++ String.intercalate ", " (xs.map renderNickelData) ++ "]"
  | .record fields =>
      let rendered := fields.map (fun (k, v) => s!"{renderNickelKey k} = {renderNickelData v}")
      "{ " ++ String.intercalate ", " rendered ++ " }"

partial def renderJsonData : Data → String
  | .null => "null"
  | .bool b => if b then "true" else "false"
  | .int n => toString n
  | .string s => renderJsonString s
  | .array xs => "[" ++ String.intercalate ", " (xs.map renderJsonData) ++ "]"
  | .record fields =>
      let rendered := fields.map (fun (k, v) => s!"{renderJsonString k}: {renderJsonData v}")
      "{" ++ String.intercalate ", " rendered ++ "}"

def builtinType? (name : Name) : Option Ty :=
  match name with
  | "readFile" => some (.funTy .string [.error, .io] .string)
  | "httpGet" => some (.funTy .string [.error, .io] .string)
  | "readFileNickel" => some (.funTy .string [.error, .io] .data)
  | "parseNickel" => some (.funTy .string [.error] .data)
  | "parseJson" => some (.funTy .string [.error] .data)
  | "get" => some (.funTy .data [] (.funTy .string [.error] .data))
  | "at" => some (.funTy .data [] (.funTy .int [.error] .data))
  | "asInt" => some (.funTy .data [.error] .int)
  | "asString" => some (.funTy .data [.error] .string)
  | "asBool" => some (.funTy .data [.error] .bool)
  | "toJson" => some (.funTy .data [] .string)
  | "toNickel" => some (.funTy .data [] .string)
  | _ => none

def builtinValue? (name : Name) : Option Value :=
  match name with
  | "readFile" => some .builtinReadFile
  | "httpGet" => some .builtinHttpGet
  | "readFileNickel" => some .builtinReadFileNickel
  | "parseNickel" => some .builtinParseNickel
  | "parseJson" => some .builtinParseJson
  | "get" => some .builtinDataGet
  | "at" => some .builtinDataAt
  | "asInt" => some .builtinDataAsInt
  | "asString" => some .builtinDataAsString
  | "asBool" => some .builtinDataAsBool
  | "toJson" => some .builtinDataToJson
  | "toNickel" => some .builtinDataToNickel
  | _ => none

mutual
  partial def applyValue (fnVal argVal : Value) : EvalM Value := do
    match fnVal with
    | .closure param body closureEnv =>
        eval ((param, argVal) :: closureEnv) body
    | .builtinReadFile => do
        let path ← expectString argVal
        let contents ← IO.toEIO (fun err => .ioError (toString err)) (IO.FS.readFile path)
        pure (.string contents)
    | .builtinHttpGet => do
        let url ← expectString argVal
        let body ← IO.toEIO (fun err => .ioError (toString err)) (Http.httpGet url)
        pure (.string body)
    | .builtinReadFileNickel => do
        let path ← expectString argVal
        let rendered ← IO.toEIO (fun err => .ioError (toString err)) (NickelHost.evalFile path)
        pure (.data (← parseRenderedData rendered))
    | .builtinParseNickel => do
        let source ← expectString argVal
        let rendered ← IO.toEIO (fun err => .ioError (toString err)) (NickelHost.evalString source)
        pure (.data (← parseRenderedData rendered))
    | .builtinParseJson => do
        let source ← expectString argVal
        let rendered ← IO.toEIO (fun err => .ioError (toString err)) (NickelHost.evalJsonString source)
        pure (.data (← parseRenderedData rendered))
    | .builtinDataGet => do
        let term ← expectData argVal
        pure (.builtinDataGetField term)
    | .builtinDataGetField term => do
        let key ← expectString argVal
        match term with
        | .record fields =>
            match fields.find? (fun (name, _) => name = key) with
            | some (_, value) => pure (.data value)
            | none => throw (.typeError s!"missing field: {key}")
        | _ => throw (.typeError "get expects a record")
    | .builtinDataAt => do
        let term ← expectData argVal
        pure (.builtinDataAtIndex term)
    | .builtinDataAtIndex term => do
        let idx ← expectInt argVal
        if idx < 0 then
          throw (.typeError s!"negative index: {idx}")
        else
          match term with
          | .array items =>
              match dataArrayGet? items idx.toNat with
              | some value => pure (.data value)
              | none => throw (RuntimeError.typeError s!"index out of bounds: {idx}")
          | _ => throw (.typeError "at expects an array")
    | .builtinDataAsInt => do
        let term ← expectData argVal
        match term with
        | .int n => pure (.int n)
        | _ => throw (.typeError "asInt expects an integer")
    | .builtinDataAsString => do
        let term ← expectData argVal
        match term with
        | .string s => pure (.string s)
        | _ => throw (.typeError "asString expects a string")
    | .builtinDataAsBool => do
        let term ← expectData argVal
        match term with
        | .bool b => pure (.bool b)
        | _ => throw (.typeError "asBool expects a boolean")
    | .builtinDataToJson => do
        let term ← expectData argVal
        pure (.string (renderJsonData term))
    | .builtinDataToNickel => do
        let term ← expectData argVal
        pure (.string (renderNickelData term))
    | _ =>
        match fnVal, argVal with
        | .data (.record fields), .string key =>
            match fields.find? (fun (name, _) => name = key) with
            | some (_, value) => pure (.data value)
            | none => throw (.typeError s!"missing field: {key}")
        | .data (.array items), .int idx =>
            if idx < 0 then
              throw (.typeError s!"negative index: {idx}")
            else
              match dataArrayGet? items idx.toNat with
              | some value => pure (.data value)
              | none => throw (RuntimeError.typeError s!"index out of bounds: {idx}")
        | .data _, .string _ =>
            throw (.typeError "string dispatch expects a record")
        | .data _, .int _ =>
            throw (.typeError "int dispatch expects an array")
        | _, _ =>
            throw (.typeError "expected a function")

  partial def eval (env : Env) : Expr → EvalM Value
  | .int n => pure (.int n)
  | .bool b => pure (.bool b)
  | .string s => pure (.string s)
  | .null => pure (.data .null)
  | .var name =>
      match env.lookup name with
      | some v => pure v
      | none =>
          match builtinValue? name with
          | some v => pure v
          | none => throw (.unboundVariable name)
  | .arrayE items => do
      let values ← items.mapM (eval env)
      pure (.data (.array (← values.mapM valueToData)))
  | .recordE fields => do
      let fields' ← fields.mapM (fun (name, value) => do
        let value' ← eval env value
        pure (name, (← valueToData value')))
      pure (.data (.record fields'))
  | .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
        throw .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 errName handler => do
      try
        eval env body
      catch err =>
        eval ((errName, runtimeErrorToValue err) :: env) handler
  | .parMapE itemName collection body => do
      let source ← expectData (← eval env collection)
      match source with
      | .array items =>
          let tasks ← items.mapM (fun item =>
            EIO.asTask (prio := .dedicated) (eval ((itemName, .data item) :: env) body))
          let results : List (Except RuntimeError Value) := List.map Task.get tasks
          pure (.list (results.map encodeMapResult))
      | _ =>
          throw (.typeError "pmap expects an array value")
  | .lam param _ body =>
      pure (.closure param body env)
  | .app fn arg => do
      let fnVal ← eval env fn
      let argVal ← eval env arg
      applyValue fnVal argVal
end

def renderData : Data → String := renderNickelData

def renderValue : Value → String
  | .int n => toString n
  | .bool b => toString b
  | .string s => s!"\"{s}\""
  | .data data => renderData data
  | .list items => "[" ++ String.intercalate ", " (items.map renderValue) ++ "]"
  | .resultOk value => s!"ok({renderValue value})"
  | .resultErr msg => s!"err(\"{msg}\")"
  | .error msg => s!"<error:{msg}>"
  | .builtinReadFile => "<builtin:readFile>"
  | .builtinHttpGet => "<builtin:httpGet>"
  | .builtinReadFileNickel => "<builtin:readFileNickel>"
  | .builtinParseNickel => "<builtin:parseNickel>"
  | .builtinParseJson => "<builtin:parseJson>"
  | .builtinDataGet => "<builtin:get>"
  | .builtinDataGetField _ => "<builtin:getField>"
  | .builtinDataAt => "<builtin:at>"
  | .builtinDataAtIndex _ => "<builtin:atIndex>"
  | .builtinDataAsInt => "<builtin:asInt>"
  | .builtinDataAsString => "<builtin:asString>"
  | .builtinDataAsBool => "<builtin:asBool>"
  | .builtinDataToJson => "<builtin:toJson>"
  | .builtinDataToNickel => "<builtin:toNickel>"
  | .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 := [] }

def eraseDataRefinement : Ty → Ty
  | .dataInt | .dataBool | .dataString | .dataNull | .dataArray _ | .dataRecord _ => .data
  | ty => ty

def canFlowTo (actual expected : Ty) : Bool :=
  if actual == expected then
    true
  else
    match actual, expected with
    | .dataInt, .data => true
    | .dataBool, .data => true
    | .dataString, .data => true
    | .dataNull, .data => true
    | .dataArray _, .data => true
    | .dataRecord _, .data => true
    | .dataArray _, .dataArray .data => true
    | _, _ => false

partial def constString? : Expr → Option String
  | .string s => some s
  | _ => none

unsafe def evalConstData? (env : List (Name × Data)) : Expr → Option Data
  | .null => some .null
  | .string s => some (.string s)
  | .int n => some (.int n)
  | .bool b => some (.bool b)
  | .var name => env.lookup name
  | .arrayE items => do
      some (.array (← items.mapM (evalConstData? env)))
  | .recordE fields => do
      some (.record (← fields.mapM (fun (name, value) => do
        let value' ← evalConstData? env value
        pure (name, value'))))
  | .letE name value body => do
      let value' ← evalConstData? env value
      evalConstData? ((name, value') :: env) body
  | .ifE cond thenBranch elseBranch => do
      match (← evalConstData? env cond) with
      | .bool true => evalConstData? env thenBranch
      | .bool false => evalConstData? env elseBranch
      | _ => none
  | .app fn arg => do
      match fn with
      | .var "parseNickel" =>
          let source ← evalConstData? env arg
          match source with
          | .string s =>
              match unsafeIO (NickelHost.evalString s) with
              | .ok rendered =>
                  match Nickel.parse rendered with
                  | .ok term => some (Data.ofNickel term)
                  | .error _ => none
              | .error _ => none
          | _ => none
      | .var "parseJson" =>
          let source ← evalConstData? env arg
          match source with
          | .string s =>
              match unsafeIO (NickelHost.evalJsonString s) with
              | .ok rendered =>
                  match Nickel.parse rendered with
                  | .ok term => some (Data.ofNickel term)
                  | .error _ => none
              | .error _ => none
          | _ => none
      | .var "httpGet" =>
          let url ← evalConstData? env arg
          match url with
          | .string s =>
              match unsafeIO (Http.httpGet s) with
              | .ok body => some (.string body)
              | .error _ => none
          | _ => none
      | .var "readFileNickel" =>
          let path ← evalConstData? env arg
          match path with
          | .string p =>
              match unsafeIO (NickelHost.evalFile p) with
              | .ok rendered =>
                  match Nickel.parse rendered with
                  | .ok term => some (Data.ofNickel term)
                  | .error _ => none
              | .error _ => none
          | _ => none
      | .app (.var "get") base => do
          let container ← evalConstData? env base
          let key ← evalConstData? env arg
          match container, key with
          | .record fields, .string field =>
              match fields.find? (fun (name, _) => name = field) with
              | some (_, value) => some value
              | none => none
          | _, _ => none
      | .app (.var "at") base => do
          let container ← evalConstData? env base
          let index ← evalConstData? env arg
          match container, index with
          | .array items, .int idx =>
              if idx < 0 then none else dataArrayGet? items idx.toNat
          | _, _ => none
      | _ => none
  | _ => none

unsafe def inferType (env : TyEnv) : Expr → CheckM Judgment
  | .int _ => pure (pureJudgment .int)
  | .bool _ => pure (pureJudgment .bool)
  | .string _ => pure (pureJudgment .string)
  | .null => pure (pureJudgment .dataNull)
  | .var name =>
      match env.lookup name with
      | some ty => pure (pureJudgment ty)
      | none =>
          match builtinType? name with
          | some ty => pure (pureJudgment ty)
          | none => throw (.unboundVariable name)
  | .arrayE items => do
      let itemJs ← items.mapM (inferType env)
      let dataTys ← itemJs.mapM (fun j =>
        match j.ty with
        | .int => pure .dataInt
        | .bool => pure .dataBool
        | .string => pure .dataString
        | .data => pure .data
        | .dataInt => pure .dataInt
        | .dataBool => pure .dataBool
        | .dataString => pure .dataString
        | .dataNull => pure .dataNull
        | .dataArray ty => pure (.dataArray ty)
        | .dataRecord fields => pure (.dataRecord fields)
        | ty => throw (.mismatch .data ty))
      let itemTy :=
        match dataTys with
        | [] => .data
        | first :: rest => if rest.all (· == first) then first else .data
      pure {
        ty := .dataArray itemTy
        effects := itemJs.foldl (fun acc j => acc.union j.effects) []
      }
  | .recordE fields => do
      let fieldJs ← fields.mapM (fun (name, value) => do pure (name, ← inferType env value))
      let fieldTys ← fieldJs.mapM (fun (name, j) => do
        let ty ← match j.ty with
          | .int => pure .dataInt
          | .bool => pure .dataBool
          | .string => pure .dataString
          | .data => pure .data
          | .dataInt => pure .dataInt
          | .dataBool => pure .dataBool
          | .dataString => pure .dataString
          | .dataNull => pure .dataNull
          | .dataArray ty => pure (.dataArray ty)
          | .dataRecord fields => pure (.dataRecord fields)
          | ty => throw (.mismatch .data ty)
        pure (name, ty, j.effects))
      pure {
        ty := .dataRecord (fieldTys.map (fun (name, ty, _) => (name, ty)))
        effects := fieldTys.foldl (fun acc (_, _, effects) => acc.union effects) []
      }
  | .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 errName handler => do
      let bodyJ ← inferType env body
      let handlerJ ← inferType ((errName, .error) :: env) handler
      ensureType bodyJ.ty handlerJ.ty
      pure {
        ty := bodyJ.ty
        effects := (bodyJ.effects.erase .error).union handlerJ.effects
      }
  | .parMapE itemName collection body => do
      let collectionJ ← inferType env collection
      match collectionJ.ty with
      | .data
      | .dataArray _ =>
          let itemTy :=
            match collectionJ.ty with
            | .dataArray ty => ty
            | _ => .data
          let bodyJ ← inferType ((itemName, itemTy) :: env) body
          pure {
            ty := .list
            effects := collectionJ.effects.union bodyJ.effects
          }
      | _ =>
          throw (.mismatch .data collectionJ.ty)
  | .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 =>
          if canFlowTo argJ.ty paramTy then
            let refinedResultTy :=
              match fn, arg, resultTy with
              | .app (.var "get") base, .string key, .data =>
                  match inferType env base with
                  | .ok baseJ =>
                      match baseJ.ty with
                      | .dataRecord fields =>
                          match dataRecordField? fields key with
                          | some ty => liftScalarDataTy ty
                          | none => .data
                      | _ =>
                          match evalConstData? [] (.app fn arg) with
                          | some data => liftScalarDataTy (dataType data)
                          | none => .data
                  | .error _ => .data
              | .app (.var "at") base, _, .data =>
                  match evalConstData? [] (.app fn arg) with
                  | some data => liftScalarDataTy (dataType data)
                  | none =>
                      match inferType env base with
                      | .ok baseJ =>
                          match baseJ.ty with
                          | .dataArray itemTy => liftScalarDataTy itemTy
                          | _ => .data
                      | .error _ => .data
              | .var "parseNickel", _, .data =>
                  match evalConstData? [] (.app fn arg) with
                  | some data => refineKnownDataTy (dataType data)
                  | none => .data
              | .var "parseJson", _, .data =>
                  match evalConstData? [] (.app fn arg) with
                  | some data => refineKnownDataTy (dataType data)
                  | none => .data
              | .var "readFileNickel", _, .data =>
                  match evalConstData? [] (.app fn arg) with
                  | some data => refineKnownDataTy (dataType data)
                  | none => .data
              | .var "asInt", _, .int =>
                  match argJ.ty with
                  | .dataInt => .int
                  | _ => .int
              | .var "asString", _, .string =>
                  match argJ.ty with
                  | .dataString => .string
                  | _ => .string
              | .var "asBool", _, .bool =>
                  match argJ.ty with
                  | .dataBool => .bool
                  | _ => .bool
              | _, _, _ => resultTy
            pure {
            ty := refinedResultTy
            effects := fnJ.effects.union (argJ.effects.union latentEffects)
          }
          else
            throw (.mismatch paramTy argJ.ty)
      | _ =>
          throw (.expectedFunction fnJ.ty)

partial def renderType : Ty → String
  | .int => "Int"
  | .bool => "Bool"
  | .string => "String"
  | .data => "Data"
  | .dataInt => "Data:Int"
  | .dataBool => "Data:Bool"
  | .dataString => "Data:String"
  | .dataNull => "Data:Null"
  | .dataArray itemTy => s!"Data:[{renderType itemTy}]"
  | .dataRecord fields =>
      let rendered := fields.map (fun (name, ty) => s!"{name}: {renderType ty}")
      "Data:{" ++ String.intercalate ", " rendered ++ "}"
  | .list => "List"
  | .result => "Result"
  | .error => "Error"
  | .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