namespace Mlang.Nickel

abbrev Decimal := String

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

abbrev ParseM := Except String
abbrev Input := List Char

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

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 || c = '-'

partial def skipSpace : Input → Input
  | c :: rest => if c.isWhitespace then skipSpace rest else c :: rest
  | [] => []

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

def consumeChar (expected : Char) : Input → ParseM Input
  | c :: rest =>
      if c = expected then pure rest
      else throw s!"expected '{expected}', got '{c}'"
  | [] => throw s!"expected '{expected}', got end of input"

partial def takeStringChars : Input → ParseM (List Char × Input)
  | [] => throw "unterminated string literal"
  | '"' :: rest => pure ([], rest)
  | '\\' :: '"' :: rest => do
      let (chars, tail) ← takeStringChars rest
      pure ('"' :: chars, tail)
  | '\\' :: '\\' :: 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 parseString (input : Input) : ParseM (String × Input) := do
  let input ← consumeChar '"' input
  let (chars, rest) ← takeStringChars input
  pure (charsToString chars, rest)

def parseInt (input : Input) : ParseM (Int × Input) := do
  let (sign, rest) :=
    match input with
    | '-' :: tail => ("-", tail)
    | _ => ("", input)
  let (digits, rest) := spanChars isDigit rest
  if digits.isEmpty then
    throw "expected integer literal"
  else
    match String.toInt? (sign ++ charsToString digits) with
    | some n => pure (n, rest)
    | none => throw "invalid integer literal"

def parseNumber (input : Input) : ParseM (Term × Input) := do
  let (sign, rest) :=
    match input with
    | '-' :: tail => ("-", tail)
    | _ => ("", input)
  let (wholeDigits, rest) := spanChars isDigit rest
  if wholeDigits.isEmpty then
    throw "expected numeric literal"
  else
    match rest with
    | '.' :: fracRest =>
        let (fracDigits, tail) := spanChars isDigit fracRest
        if fracDigits.isEmpty then
          throw "invalid decimal literal"
        else
          pure (.decimal (sign ++ charsToString wholeDigits ++ "." ++ charsToString fracDigits), tail)
    | _ =>
        match String.toInt? (sign ++ charsToString wholeDigits) with
        | some n => pure (.int n, rest)
        | none => throw "invalid integer literal"

def parseIdent (input : Input) : ParseM (String × Input) := do
  match input with
  | c :: rest =>
      if isIdentStart c then
        let (tailChars, tail) := spanChars isIdentContinue rest
        pure (charsToString (c :: tailChars), tail)
      else
        throw s!"expected identifier, got '{c}'"
  | [] => throw "expected identifier, got end of input"

mutual
  partial def parseTerm (input : Input) : ParseM (Term × Input) := do
    let input := skipSpace input
    match input with
    | [] => throw "expected Nickel term"
    | '"' :: _ =>
        let (s, rest) ← parseString input
        pure (.string s, rest)
    | '[' :: rest =>
        parseArray rest
    | '{' :: rest =>
        parseRecord rest
    | 't' :: _ | 'f' :: _ | 'n' :: _ =>
        let (name, rest) ← parseIdent input
        match name with
        | "true" => pure (.bool true, rest)
        | "false" => pure (.bool false, rest)
        | "null" => pure (.null, rest)
        | _ => throw s!"unexpected identifier '{name}'"
    | '-' :: _ =>
        parseNumber input
    | c :: _ =>
        if isDigit c then
          parseNumber input
        else
          throw s!"unexpected character '{c}'"

  partial def parseArray (input : Input) : ParseM (Term × Input) := do
    let input := skipSpace input
    if let ']' :: rest := input then
      pure (.array [], rest)
    else
      let (first, rest) ← parseTerm input
      parseArrayTail [first] rest

  partial def parseArrayTail (items : List Term) (input : Input) : ParseM (Term × Input) := do
    let input := skipSpace input
    match input with
    | ']' :: rest => pure (.array items.reverse, rest)
    | ',' :: rest => do
        let (next, rest) ← parseTerm rest
        parseArrayTail (next :: items) rest
    | c :: _ => throw s!"expected ',' or ']', got '{c}'"
    | [] => throw "unterminated array"

  partial def parseRecord (input : Input) : ParseM (Term × Input) := do
    let input := skipSpace input
    if let '}' :: rest := input then
      pure (.record [], rest)
    else
      let (key, rest) ← parseFieldKey input
      let rest := skipSpace rest
      let rest ← consumeChar '=' rest
      let (value, rest) ← parseTerm rest
      parseRecordTail [(key, value)] rest

  partial def parseRecordTail (fields : List (String × Term)) (input : Input) : ParseM (Term × Input) := do
    let input := skipSpace input
    match input with
    | '}' :: rest => pure (.record fields.reverse, rest)
    | ',' :: rest => do
        let (key, rest) ← parseFieldKey rest
        let rest := skipSpace rest
        let rest ← consumeChar '=' rest
        let (value, rest) ← parseTerm rest
        parseRecordTail ((key, value) :: fields) rest
    | c :: _ => throw s!"expected ',' or '}}', got '{c}'"
    | [] => throw "unterminated record"

  partial def parseFieldKey (input : Input) : ParseM (String × Input) := do
    let input := skipSpace input
    match input with
    | '"' :: _ => parseString input
    | _ => parseIdent input
end

def parse (source : String) : ParseM Term := do
  let (term, rest) ← parseTerm source.toList
  let rest := skipSpace rest
  if rest.isEmpty then
    pure term
  else
    throw s!"unexpected trailing input: {charsToString rest}"

partial def render : Term → String
  | .null => "null"
  | .bool true => "true"
  | .bool false => "false"
  | .int n => toString n
  | .decimal d => d
  | .string s => s!"\"{s}\""
  | .array xs =>
      "[" ++ String.intercalate ", " (xs.map render) ++ "]"
  | .record fields =>
      let rendered := fields.map (fun (k, v) => s!"{k} = {render v}")
      "{ " ++ String.intercalate ", " rendered ++ " }"

end Mlang.Nickel