import Mlang.Interpreter

namespace Mlang

inductive Token where
  | int : Int → Token
  | bool : Bool → Token
  | ident : String → Token
  | lparen : Token
  | rparen : Token
  | colon : Token
  | plus : Token
  | minus : Token
  | star : Token
  | slash : Token
  | eq : Token
  | thinArrow : Token
  | fatArrow : Token
  | kwLet : Token
  | kwIn : Token
  | kwIf : Token
  | kwThen : Token
  | kwElse : Token
  | kwFun : Token
  | kwTry : Token
  | kwWith : Token
  deriving Repr, BEq, Inhabited

inductive ParseError where
  | lexError : String → ParseError
  | unexpectedEof : String → ParseError
  | unexpectedToken : String → ParseError
  | trailingTokens : List Token → ParseError
  deriving Repr, Inhabited

abbrev ParseM := Except ParseError

def isDigit (c : Char) : Bool :=
  '0' <= c && c <= '9'

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

def isIdentContinue (c : Char) : Bool :=
  isIdentStart c || isDigit c

def isAtomStart : Token → Bool
  | .int _ | .bool _ | .ident _ | .lparen => true
  | _ => false

partial def spanChars (p : Char → Bool) : List Char → List Char × List Char
  | [] => ([], [])
  | c :: cs =>
      if p c then
        let (taken, rest) := spanChars p cs
        (c :: taken, rest)
      else
        ([], c :: cs)

def charsToString (chars : List Char) : String :=
  String.ofList chars

def lexIdent (chars : List Char) : Token × List Char :=
  let (nameChars, rest) := spanChars isIdentContinue chars
  let name := charsToString nameChars
  let token :=
    match name with
    | "let" => .kwLet
    | "in" => .kwIn
    | "if" => .kwIf
    | "then" => .kwThen
    | "else" => .kwElse
    | "fun" => .kwFun
    | "try" => .kwTry
    | "with" => .kwWith
    | "true" => .bool true
    | "false" => .bool false
    | _ => .ident name
  (token, rest)

def lexNumber (chars : List Char) : ParseM (Token × List Char) := do
  let (digits, rest) := spanChars isDigit chars
  match String.toInt? (charsToString digits) with
  | some n => pure (.int n, rest)
  | none => throw (.lexError "invalid integer literal")

partial def lexChars : List Char → ParseM (List Token)
  | [] => pure []
  | c :: cs =>
      if c.isWhitespace then
        lexChars cs
      else if isDigit c then do
        let (tok, rest) ← lexNumber (c :: cs)
        let toks ← lexChars rest
        pure (tok :: toks)
      else if isIdentStart c then do
        let (tok, rest) := lexIdent (c :: cs)
        let toks ← lexChars rest
        pure (tok :: toks)
      else
        match c, cs with
        | '(', _ => do pure (.lparen :: (← lexChars cs))
        | ')', _ => do pure (.rparen :: (← lexChars cs))
        | ':', _ => do pure (.colon :: (← lexChars cs))
        | '+', _ => do pure (.plus :: (← lexChars cs))
        | '-', '>' :: rest => do pure (.thinArrow :: (← lexChars rest))
        | '-', _ => do pure (.minus :: (← lexChars cs))
        | '*', _ => do pure (.star :: (← lexChars cs))
        | '/', _ => do pure (.slash :: (← lexChars cs))
        | '=', '>' :: rest => do pure (.fatArrow :: (← lexChars rest))
        | '=', _ => do pure (.eq :: (← lexChars cs))
        | _, _ => throw (.lexError s!"unexpected character '{c}'")

def lex (input : String) : ParseM (List Token) :=
  lexChars input.toList

abbrev ParserState := List Token

def expectToken (expected : Token) : ParserState → ParseM ParserState
  | tok :: rest =>
      if tok == expected then
        pure rest
      else
        throw (.unexpectedToken s!"expected {reprStr expected}, got {reprStr tok}")
  | [] => throw (.unexpectedEof s!"expected {reprStr expected}")

mutual
  partial def parseTy (tokens : ParserState) : ParseM (Ty × ParserState) := do
    let (lhs, rest) ← parseTyAtom tokens
    match rest with
    | .thinArrow :: rest =>
        let (rhs, rest) ← parseTy rest
        pure (.funTy lhs [] rhs, rest)
    | _ => pure (lhs, rest)

  partial def parseTyAtom : ParserState → ParseM (Ty × ParserState)
    | .ident "Int" :: rest => pure (.int, rest)
    | .ident "Bool" :: rest => pure (.bool, rest)
    | .lparen :: rest => do
        let (ty, rest) ← parseTy rest
        let rest ← expectToken .rparen rest
        pure (ty, rest)
    | tok :: _ => throw (.unexpectedToken s!"expected type, got {reprStr tok}")
    | [] => throw (.unexpectedEof "expected type")
end

mutual
  partial def parseExpr (tokens : ParserState) : ParseM (Expr × ParserState) := do
    match tokens with
    | .kwLet :: .ident name :: .eq :: rest => do
        let (value, rest) ← parseExpr rest
        let rest ← expectToken .kwIn rest
        let (body, rest) ← parseExpr rest
        pure (.letE name value body, rest)
    | .kwIf :: rest => do
        let (cond, rest) ← parseExpr rest
        let rest ← expectToken .kwThen rest
        let (thenBranch, rest) ← parseExpr rest
        let rest ← expectToken .kwElse rest
        let (elseBranch, rest) ← parseExpr rest
        pure (.ifE cond thenBranch elseBranch, rest)
    | .kwTry :: rest => do
        let (body, rest) ← parseExpr rest
        let rest ← expectToken .kwWith rest
        let (handler, rest) ← parseExpr rest
        pure (.tryE body handler, rest)
    | .kwFun :: .lparen :: .ident name :: .colon :: rest => do
        let (paramTy, rest) ← parseTy rest
        let rest ← expectToken .rparen rest
        let rest ← expectToken .fatArrow rest
        let (body, rest) ← parseExpr rest
        pure (.lam name paramTy body, rest)
    | _ =>
        parseEq tokens

  partial def parseEq (tokens : ParserState) : ParseM (Expr × ParserState) := do
    let (lhs, rest) ← parseAddSub tokens
    match rest with
    | .eq :: rest =>
        let (rhs, rest) ← parseEq rest
        pure (Expr.eq lhs rhs, rest)
    | _ => pure (lhs, rest)

  partial def parseAddSub (tokens : ParserState) : ParseM (Expr × ParserState) := do
    let (first, rest) ← parseMul tokens
    parseAddSubTail first rest

  partial def parseAddSubTail (lhs : Expr) (tokens : ParserState) : ParseM (Expr × ParserState) := do
    match tokens with
    | .plus :: rest => do
        let (rhs, rest) ← parseMul rest
        parseAddSubTail (.add lhs rhs) rest
    | .minus :: rest => do
        let (rhs, rest) ← parseMul rest
        parseAddSubTail (.sub lhs rhs) rest
    | _ => pure (lhs, tokens)

  partial def parseMul (tokens : ParserState) : ParseM (Expr × ParserState) := do
    let (first, rest) ← parseApp tokens
    parseMulTail first rest

  partial def parseMulTail (lhs : Expr) (tokens : ParserState) : ParseM (Expr × ParserState) := do
    match tokens with
    | .star :: rest => do
        let (rhs, rest) ← parseApp rest
        parseMulTail (.mul lhs rhs) rest
    | .slash :: rest => do
        let (rhs, rest) ← parseApp rest
        parseMulTail (.div lhs rhs) rest
    | _ => pure (lhs, tokens)

  partial def parseApp (tokens : ParserState) : ParseM (Expr × ParserState) := do
    let (fn, rest) ← parseAtom tokens
    parseAppTail fn rest

  partial def parseAppTail (fn : Expr) (tokens : ParserState) : ParseM (Expr × ParserState) := do
    match tokens with
    | tok :: _ =>
        if isAtomStart tok then do
          let (arg, rest) ← parseAtom tokens
          parseAppTail (.app fn arg) rest
        else
          pure (fn, tokens)
    | [] => pure (fn, [])

  partial def parseAtom : ParserState → ParseM (Expr × ParserState)
    | .int n :: rest => pure (.int n, rest)
    | .bool b :: rest => pure (.bool b, rest)
    | .ident name :: rest => pure (.var name, rest)
    | .lparen :: rest => do
        let (expr, rest) ← parseExpr rest
        let rest ← expectToken .rparen rest
        pure (expr, rest)
    | tok :: _ => throw (.unexpectedToken s!"expected expression, got {reprStr tok}")
    | [] => throw (.unexpectedEof "expected expression")
end

def parse (input : String) : ParseM Expr := do
  let tokens ← lex input
  let (expr, rest) ← parseExpr tokens
  match rest with
  | [] => pure expr
  | _ => throw (.trailingTokens rest)

end Mlang