import Mlang.Interpreter

namespace Mlang

inductive Token where
  | int : Int → Token
  | bool : Bool → Token
  | string : String → Token
  | ident : String → Token
  | lparen : Token
  | rparen : Token
  | lbracket : Token
  | rbracket : Token
  | lbrace : Token
  | rbrace : Token
  | semicolon : Token
  | comma : 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
  | kwPmap : Token
  | kwNull : 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 _ | .string _ | .ident _ | .lparen | .lbracket | .lbrace | .kwNull => 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
    | "pmap" => .kwPmap
    | "null" => .kwNull
    | "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 takeStringChars : List Char → ParseM (List Char × List Char)
  | [] => throw (.lexError "unterminated string literal")
  | '"' :: rest => pure ([], rest)
  | '\\' :: '"' :: rest => do
      let (chars, tail) ← takeStringChars rest
      pure ('"' :: chars, tail)
  | '\\' :: 'n' :: rest => do
      let (chars, tail) ← takeStringChars rest
      pure ('\n' :: chars, tail)
  | c :: rest => do
      let (chars, tail) ← takeStringChars rest
      pure (c :: chars, tail)

def lexString (chars : List Char) : ParseM (Token × List Char) := do
  let (content, rest) ← takeStringChars chars
  pure (.string (charsToString content), rest)

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
            let (tok, rest) ← lexString cs
            let toks ← lexChars rest
            pure (tok :: toks)
        | '(', _ => do pure (.lparen :: (← lexChars cs))
        | ')', _ => do pure (.rparen :: (← lexChars cs))
        | '[', _ => do pure (.lbracket :: (← lexChars cs))
        | ']', _ => do pure (.rbracket :: (← lexChars cs))
        | '{', _ => do pure (.lbrace :: (← lexChars cs))
        | '}', _ => do pure (.rbrace :: (← lexChars cs))
        | ';', _ => do pure (.semicolon :: (← lexChars cs))
        | ',', _ => do pure (.comma :: (← 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

inductive ReplCommand where
  | bind : Name → Expr → ReplCommand
  | evalProgram : List Expr → ReplCommand
  deriving Repr, Inhabited

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}")

def threadStep (current : Expr) (stepName : Name) : Expr :=
  .app (.var stepName) current

def threadTemplateStep (current : Expr) (step : Expr) : Expr :=
  match step with
  | .app fn arg => .app (.app fn current) arg
  | .var _ => .app step current
  | _ => .app step current

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)
    | .ident "String" :: rest => pure (.string, rest)
    | .ident "Data" :: rest => pure (.data, rest)
    | .ident "List" :: rest => pure (.list, rest)
    | .ident "Result" :: rest => pure (.result, rest)
    | .ident "Error" :: rest => pure (.error, 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
    | .thinArrow :: rest => do
        let (start, rest) ← parseThreadSeed rest
        parseThreadTail start rest
    | .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
        match rest with
        | .ident errName :: .fatArrow :: rest => do
            let (handler, rest) ← parseExpr rest
            pure (.tryE body errName handler, rest)
        | tok :: _ =>
            throw (.unexpectedToken s!"expected handler binder, got {reprStr tok}")
        | [] =>
            throw (.unexpectedEof "expected handler binder")
    | .kwPmap :: .ident itemName :: .kwIn :: rest => do
        let (collection, rest) ← parseExpr rest
        let rest ← expectToken .fatArrow rest
        let (body, rest) ← parseExpr rest
        pure (.parMapE itemName collection body, 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 parseThreadTail (current : Expr) (tokens : ParserState) : ParseM (Expr × ParserState) := do
    match tokens with
    | .ident stepName :: rest =>
        parseThreadTail (threadStep current stepName) rest
    | .lparen :: _ => do
        let (step, rest) ← parseAtom tokens
        parseThreadTail (threadTemplateStep current step) rest
    | _ =>
        pure (current, tokens)

  partial def parseThreadSeed (tokens : ParserState) : ParseM (Expr × ParserState) := do
    parseAtom 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 parseArrayItems (items : List Expr) : ParserState → ParseM (Expr × ParserState)
    | .rbracket :: rest => pure (.arrayE items.reverse, rest)
    | .comma :: rest => do
        let (item, rest) ← parseExpr rest
        parseArrayItems (item :: items) rest
    | tok :: _ => throw (.unexpectedToken s!"expected ',' or ']', got {reprStr tok}")
    | [] => throw (.unexpectedEof "unterminated array")

  partial def parseArray : ParserState → ParseM (Expr × ParserState)
    | .rbracket :: rest => pure (.arrayE [], rest)
    | tokens => do
        let (first, rest) ← parseExpr tokens
        parseArrayItems [first] rest

  partial def parseFieldKey : ParserState → ParseM (String × ParserState)
    | .ident name :: rest => pure (name, rest)
    | .string name :: rest => pure (name, rest)
    | tok :: _ => throw (.unexpectedToken s!"expected field key, got {reprStr tok}")
    | [] => throw (.unexpectedEof "expected field key")

  partial def parseRecordFields (fields : List (String × Expr)) : ParserState → ParseM (Expr × ParserState)
    | .rbrace :: rest => pure (.recordE fields.reverse, rest)
    | .comma :: rest => do
        let (name, rest) ← parseFieldKey rest
        let rest ← expectToken .eq rest
        let (value, rest) ← parseExpr rest
        parseRecordFields ((name, value) :: fields) rest
    | tok :: _ => throw (.unexpectedToken s!"expected ',' or '}}', got {reprStr tok}")
    | [] => throw (.unexpectedEof "unterminated record")

  partial def parseRecord : ParserState → ParseM (Expr × ParserState)
    | .rbrace :: rest => pure (.recordE [], rest)
    | tokens => do
        let (name, rest) ← parseFieldKey tokens
        let rest ← expectToken .eq rest
        let (value, rest) ← parseExpr rest
        parseRecordFields [(name, value)] rest

  partial def parseAtom : ParserState → ParseM (Expr × ParserState)
    | .int n :: rest => pure (.int n, rest)
    | .bool b :: rest => pure (.bool b, rest)
    | .string s :: rest => pure (.string s, rest)
    | .kwNull :: rest => pure (.null, rest)
    | .ident name :: rest => pure (.var name, rest)
    | .lbracket :: rest => parseArray rest
    | .lbrace :: rest => parseRecord 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

partial def parseProgramTokens : ParserState → ParseM (List Expr)
  | [] => pure []
  | .semicolon :: rest => parseProgramTokens rest
  | tokens => do
      let (expr, rest) ← parseExpr tokens
      match rest with
      | [] => pure [expr]
      | .semicolon :: more =>
          let tail ← parseProgramTokens more
          pure (expr :: tail)
      | _ => throw (.trailingTokens rest)

def parseProgram (input : String) : ParseM (List Expr) := do
  let tokens ← lex input
  parseProgramTokens tokens

def parseReplCommand (input : String) : ParseM ReplCommand := do
  let tokens ← lex input
  match tokens with
  | .kwLet :: .ident name :: .eq :: rest =>
      match parseExpr rest with
      | .ok (value, []) => pure (.bind name value)
      | _ =>
          let program ← parseProgramTokens tokens
          pure (.evalProgram program)
  | _ =>
      let program ← parseProgramTokens tokens
      pure (.evalProgram program)

def parse (input : String) : ParseM Expr := do
  let program ← parseProgram input
  match program with
  | [expr] => pure expr
  | [] => throw (ParseError.unexpectedEof "expected expression")
  | _ => throw (ParseError.unexpectedToken "expected a single expression")

end Mlang