import Mlang

open Mlang

def refineJudgmentFromValue (value : Value) (judgment : Judgment) : Judgment :=
  match value with
  | .data term =>
      let refinedTy := refineKnownDataTy (dataType term)
      match judgment.ty with
      | .data
      | .dataInt
      | .dataBool
      | .dataString
      | .dataNull
      | .dataArray _
      | .dataRecord _ =>
          { judgment with ty := refinedTy }
      | _ => judgment
  | _ => judgment

def renderValueForJudgment (value : Value) (judgment : Judgment) : String :=
  match judgment.ty, value with
  | .decimal, .number (.exact q) => renderDecimal q
  | .int, .number (.exact q) =>
      if q.den = 1 then toString q.num else renderValue value
  | _, _ => renderValue value

def renderNumericStability (value : Value) (judgment : Judgment) : String :=
  match judgment.ty with
  | .number =>
      s!" [rational={rationalStabilityOfValue value}, decimal={decimalStabilityOfValue value}]"
  | _ => ""

def renderResult (value : Value) (judgment : Judgment) : String :=
  let refined := refineJudgmentFromValue value judgment
  s!"{renderValueForJudgment value refined} : {renderJudgment refined}{renderNumericStability value refined}"

unsafe def checkExprsWith (tyEnv : TyEnv) (env : Env) (exprs : List Expr) : IO (Except String (List (Value × Judgment))) := do
  let rec loop (pending : List Expr) (acc : List (Value × Judgment)) : IO (Except String (List (Value × Judgment))) := do
    match pending with
    | [] => pure (.ok acc.reverse)
    | expr :: rest =>
        match inferType tyEnv expr with
        | .error err =>
            pure (.error s!"type error: {renderTypeError err}")
        | .ok judgment =>
            match (← (eval env expr).toIO') with
            | .ok value =>
                loop rest ((value, judgment) :: acc)
            | .error err =>
                pure (.error s!"runtime error: {reprStr err}")
  loop exprs []

unsafe def checkAndEvalWith (tyEnv : TyEnv) (env : Env) (source : String) : IO (Except String (List (Value × Judgment))) := do
  match parseProgram source with
  | .error err =>
      pure (.error s!"parse error: {reprStr err}")
  | .ok exprs => do
      checkExprsWith tyEnv env exprs

unsafe def checkAndEval (source : String) : IO (Except String (List (Value × Judgment))) :=
  checkAndEvalWith [] [] source

unsafe def runFile (path : String) : IO UInt32 := do
  let source ← IO.FS.readFile path
  match (← checkAndEval source) with
  | .ok results =>
      for (value, judgment) in results do
        IO.println (renderResult value judgment)
      pure 0
  | .error msg =>
      IO.eprintln s!"{path}: {msg}"
      pure 1

unsafe def runEval (source : String) : IO UInt32 := do
  match (← checkAndEval source) with
  | .ok results =>
      for (value, judgment) in results do
        IO.println (renderResult value judgment)
      pure 0
  | .error msg =>
      IO.eprintln msg
      pure 1

unsafe def repl (tyEnv : TyEnv := []) (env : Env := []) : IO UInt32 := do
  let stdin ← IO.getStdin
  let line ← stdin.getLine
  let input := (String.trimAscii line).toString
  if input.isEmpty then
    repl tyEnv env
  else if input = ":quit" || input = ":q" then
    pure 0
  else do
    match parseReplCommand input with
    | .error err =>
        IO.eprintln s!"parse error: {reprStr err}"
        repl tyEnv env
    | .ok command =>
        match command with
        | .bind name value =>
            match inferType tyEnv value with
            | .error err =>
                IO.eprintln s!"type error: {renderTypeError err}"
                repl tyEnv env
            | .ok judgment =>
                match (← (eval env value).toIO') with
                | .error err =>
                    IO.eprintln s!"runtime error: {reprStr err}"
                    repl tyEnv env
                | .ok boundValue =>
                    IO.println (renderResult boundValue judgment)
                    repl ((name, judgment.ty) :: tyEnv) ((name, boundValue) :: env)
        | .evalProgram exprs =>
            match (← checkExprsWith tyEnv env exprs) with
            | .ok results =>
                for (value, judgment) in results do
                  IO.println (renderResult value judgment)
            | .error msg =>
                IO.eprintln msg
            repl tyEnv env

def printUsage : IO Unit := do
  IO.println "usage:"
  IO.println "  mlang               # start REPL"
  IO.println "  mlang --eval <src>  # evaluate source from the command line"
  IO.println "  mlang <file>        # run a semicolon-separated source file"
  IO.println "  mlang :quit         # not valid; use inside REPL"

unsafe def main (args : List String) : IO UInt32 := do
  match args with
  | [] =>
      IO.println "mlang REPL. enter :quit to exit."
      repl
  | ["-h"] | ["--help"] =>
      printUsage
      pure 0
  | ["--eval", source] =>
      runEval source
  | [path] =>
      runFile path
  | _ =>
      printUsage
      pure 1