plum

module std

import std/Option
import std/Result
import std/List
import std/Map
import std/Number
import std/Str
import std/Buffer
import std/Err
import std/Bool

# A parsed JSON value. Numbers are always stored as `Float` (JSON doesn't
# distinguish integers from floats the way Plum does), and `JsonList`/
# `JsonMap` nest `Json` itself — the same self-referential shape `List[T]`
# already uses internally for `Node[T].next: Option[Node[T]]`.
enum Json =
  | JsonNull
  | JsonBool(Bool)
  | JsonFloat(Float)
  | JsonStr(Str)
  | JsonList(List[Json])
  | JsonMap(Map[Str, Json])

  # Parses `s` as a single JSON document, requiring the whole string (aside
  # from surrounding whitespace) to be consumed. Called as `Json.fromStr(s)`,
  # matching `Int.fromStr`/`Float.fromStr`'s naming.
  fun fromStr(s: Str) -> Result[Json, JsonParseError] =
    parser := JsonParser(src: s, pos: 0)
    parser.skipSpace()
    r := parser.parseValue()
    match r
      Err(e) =>
        return Err(e)
      Ok(v) =>
        parser.skipSpace()
        if parser.pos < parser.src.length()
          return Err(parser.fail("unexpected trailing characters"))
        return Ok(v)

  # Renders `self` back to its JSON text form. A literal `{`/`}` is written
  # doubled (`{{`/`}}`) — a single, un-doubled brace in a Plum string starts
  # or ends a `{expr}` interpolation instead of standing for itself.
  fun toStr(self) -> Str =
    match self
      JsonNull => "null"
      JsonBool(b) => b ? "true" : "false"
      JsonFloat(f) => f.toStr()
      JsonStr(s) => "\"" + escapeJsonStr(s) + "\""
      JsonList(items) => "[" + items.join(",") + "]"
      JsonMap(pairs) => "{{" + jsonMapToStr(pairs) + "}}"

  fun isNull(self) -> Bool =
    match self
      JsonNull => True
      _ => False

  fun asBool(self) -> Option[Bool] =
    match self
      JsonBool(b) => Some(b)
      _ => None

  fun asFloat(self) -> Option[Float] =
    match self
      JsonFloat(f) => Some(f)
      _ => None

  fun asStr(self) -> Option[Str] =
    match self
      JsonStr(s) => Some(s)
      _ => None

  fun asList(self) -> Option[List[Json]] =
    match self
      JsonList(items) => Some(items)
      _ => None

  fun asMap(self) -> Option[Map[Str, Json]] =
    match self
      JsonMap(pairs) => Some(pairs)
      _ => None

# Carries where in the source a parse failure happened, alongside the
# message — implements the shared `Err` trait (`libs/std/err.plum`) rather
# than being a plain `Str`, so a caller that wants the position can get it
# without re-parsing the message.
enum JsonParseError(Err) =
  | JsonParseError(pos: Int, text: Str)

  fun code(self) -> Int =
    1

  fun msg(self) -> Str =
    "json parse error at position {self.pos}: {self.text}"

  # A bare payload-free `None` here can't be disambiguated to one specific
  # `Option` instantiation from this expression alone (see the README's
  # "Known gaps" note on this residual generics limitation — unrelated to the
  # 3 points fixed above) — an explicit `Option[JsonParseError]` return type
  # sidesteps it.
  fun cause(self) -> Option[JsonParseError] =
    None

# A recursive-descent parser over a `Str` source, tracking its read position
# in `pos`. (The original design threaded a `Readable` stream through here,
# but `Readable` is an undefined trait with no methods anywhere in `libs/std`
# — there is nothing to call on it — so this reads directly from a `Str`
# instead, the same way every other stdlib parser, e.g. `Int.fromStr`, does.)
enum JsonParser =
  | JsonParser(src: Str, pos: Int)

  fun fail(self, text: Str) -> JsonParseError =
    JsonParseError(pos: self.pos, text: text)

  # The raw byte at the current position, or 0 (never a valid JSON byte) past
  # the end of `src` — lets every caller test for "end of input" the same way
  # it tests for any other byte, with no separate bounds check.
  fun peek(self) -> Int =
    self.pos >= self.src.length() ? 0 : self.src.byteAt(self.pos)

  fun peekAt(self, offset: Int) -> Int =
    i := self.pos + offset
    i >= self.src.length() ? 0 : self.src.byteAt(i)

  fun advance(self) -> Unit =
    self.pos = self.pos + 1

  fun skipSpace(self) -> Unit =
    while isSpace(self.peek())
      self.advance()

  fun parseValue(self) -> Result[Json, JsonParseError] =
    self.skipSpace()
    c := self.peek()
    if c == 34
      return self.parseString()
    if c == 91
      return self.parseArray()
    if c == 123
      return self.parseObject()
    if c == 116
      return self.parseLiteral("true", JsonBool(True))
    if c == 102
      return self.parseLiteral("false", JsonBool(False))
    if c == 110
      return self.parseLiteral("null", JsonNull)
    if c == 45 || isDigit(c)
      return self.parseNumber()
    return Err(self.fail("unexpected character"))

  fun parseLiteral(self, lit: Str, value: Json) -> Result[Json, JsonParseError] =
    n := lit.length()
    if !matchesAt(self.src, lit, self.pos, n)
      return Err(self.fail("invalid literal, expected '{lit}'"))
    self.pos = self.pos + n
    return Ok(value)

  fun parseNumber(self) -> Result[Json, JsonParseError] =
    start := self.pos
    if self.peek() == 45
      self.advance()
    if !isDigit(self.peek())
      return Err(self.fail("invalid number"))
    while isDigit(self.peek())
      self.advance()
    if self.peek() == 46
      self.advance()
      if !isDigit(self.peek())
        return Err(self.fail("invalid number"))
      while isDigit(self.peek())
        self.advance()
    if self.peek() == 101 || self.peek() == 69
      self.advance()
      if self.peek() == 43 || self.peek() == 45
        self.advance()
      if !isDigit(self.peek())
        return Err(self.fail("invalid number"))
      while isDigit(self.peek())
        self.advance()
    return self.parseNumberText(self.src.sub(start, self.pos))

  # `Float.fromStr` (`libs/std/float.plum`) doesn't understand exponent
  # notation, so a number that has one is split into its mantissa and
  # exponent around the `e`/`E` and recombined here instead.
  fun parseNumberText(self, text: Str) -> Result[Json, JsonParseError] =
    e_index := exponentIndex(text)
    if e_index < 0
      r := parseFloat(text)
      match r
        Ok(f) =>
          return Ok(JsonFloat(f))
        Err(m) =>
          return Err(self.fail(m))
    mantissa := parseFloat(text.sub(0, e_index))
    exponent := parseInt(text.sub(e_index + 1, text.length()))
    match mantissa
      Err(m) =>
        return Err(self.fail(m))
      Ok(mv) =>
        match exponent
          Err(m) =>
            return Err(self.fail(m))
          Ok(ev) =>
            return Ok(JsonFloat(mv * 10.pow(Float(ev))))

  fun parseString(self) -> Result[Json, JsonParseError] =
    self.advance()
    return self.parseStringBody(Buffer())

  fun parseStringBody(self, buf: Buffer) -> Result[Json, JsonParseError] =
    c := self.peek()
    if c == 0
      return Err(self.fail("unterminated string"))
    if c == 34
      self.advance()
      return Ok(JsonStr(buf.toStr()))
    if c == 92
      return self.parseEscape(buf)
    n := self.src.codePointByteLength(self.pos)
    buf.write(self.src.sub(self.pos, self.pos + n))
    self.pos = self.pos + n
    return self.parseStringBody(buf)

  fun parseEscape(self, buf: Buffer) -> Result[Json, JsonParseError] =
    self.advance()
    esc := self.peek()
    if esc == 34
      buf.write("\"")
      self.advance()
      return self.parseStringBody(buf)
    if esc == 92
      buf.write("\\")
      self.advance()
      return self.parseStringBody(buf)
    if esc == 47
      buf.write("/")
      self.advance()
      return self.parseStringBody(buf)
    if esc == 98
      buf.write("\b")
      self.advance()
      return self.parseStringBody(buf)
    if esc == 102
      buf.write("\f")
      self.advance()
      return self.parseStringBody(buf)
    if esc == 110
      buf.write("\n")
      self.advance()
      return self.parseStringBody(buf)
    if esc == 114
      buf.write("\r")
      self.advance()
      return self.parseStringBody(buf)
    if esc == 116
      buf.write("\t")
      self.advance()
      return self.parseStringBody(buf)
    if esc == 117
      self.advance()
      return self.parseUnicodeEscape(buf)
    return Err(self.fail("invalid escape character"))

  fun parseHex4(self) -> Result[Int, JsonParseError] =
    return self.parseHex4From(0, 0)

  fun parseHex4From(self, value: Int, count: Int) -> Result[Int, JsonParseError] =
    if count == 4
      return Ok(value)
    d := hexDigitValue(self.peek())
    if d < 0
      return Err(self.fail("invalid unicode escape"))
    self.advance()
    return self.parseHex4From(value * 16 + d, count + 1)

  fun parseUnicodeEscape(self, buf: Buffer) -> Result[Json, JsonParseError] =
    hr := self.parseHex4()
    match hr
      Err(e) =>
        return Err(e)
      Ok(cp) =>
        return self.parseUnicodeEscapeCont(buf, cp)

  # A high surrogate (`0xD800`..`0xDBFF`) must be followed immediately by a
  # low surrogate (`0xDC00`..`0xDFFF`) `\u` escape; the pair together encodes
  # one codepoint above `0xFFFF`. Any other codepoint is used as-is.
  fun parseUnicodeEscapeCont(self, buf: Buffer, cp: Int) -> Result[Json, JsonParseError] =
    if cp >= 0xD800 && cp <= 0xDBFF && self.peek() == 92 && self.peekAt(1) == 117
      self.advance()
      self.advance()
      hr := self.parseHex4()
      match hr
        Err(e) =>
          return Err(e)
        Ok(low) =>
          if low < 0xDC00 || low > 0xDFFF
            return Err(self.fail("invalid low surrogate"))
          combined := 0x10000 + {{cp - 0xD800} << 10} + {low - 0xDC00}
          buf.writeRune(combined)
          return self.parseStringBody(buf)
    buf.writeRune(cp)
    return self.parseStringBody(buf)

  fun parseArray(self) -> Result[Json, JsonParseError] =
    self.advance()
    self.skipSpace()
    items := List[Json](head: None, tail: None, size: 0)
    if self.peek() == 93
      self.advance()
      return Ok(JsonList(items))
    return self.parseArrayItems(items)

  fun parseArrayItems(self, items: List[Json]) -> Result[Json, JsonParseError] =
    self.skipSpace()
    r := self.parseValue()
    match r
      Err(e) =>
        return Err(e)
      Ok(v) =>
        items.add(v)
        self.skipSpace()
        c := self.peek()
        if c == 44
          self.advance()
          return self.parseArrayItems(items)
        if c == 93
          self.advance()
          return Ok(JsonList(items))
        return Err(self.fail("expected ',' or ']'"))

  fun parseObject(self) -> Result[Json, JsonParseError] =
    self.advance()
    self.skipSpace()
    pairs := Map[Str, Json]()
    if self.peek() == 125
      self.advance()
      return Ok(JsonMap(pairs))
    return self.parseObjectItems(pairs)

  fun parseObjectItems(self, pairs: Map[Str, Json]) -> Result[Json, JsonParseError] =
    self.skipSpace()
    if self.peek() != 34
      return Err(self.fail("expected string key"))
    kr := self.parseString()
    match kr
      Err(e) =>
        return Err(e)
      Ok(keyJson) =>
        return self.parseObjectValue(pairs, jsonStrValue(keyJson))

  fun parseObjectValue(self, pairs: Map[Str, Json], key: Str) -> Result[Json, JsonParseError] =
    self.skipSpace()
    if self.peek() != 58
      return Err(self.fail("expected ':'"))
    self.advance()
    self.skipSpace()
    vr := self.parseValue()
    match vr
      Err(e) =>
        return Err(e)
      Ok(v) =>
        pairs.set(key, v)
        self.skipSpace()
        c := self.peek()
        if c == 44
          self.advance()
          return self.parseObjectItems(pairs)
        if c == 125
          self.advance()
          return Ok(JsonMap(pairs))
        return Err(self.fail("expected ',' or '}}'"))

fun jsonStrValue(j: Json) -> Str =
  match j
    JsonStr(s) => s
    _ => ""

# `Map` has no built-in `toStr`/`join` (unlike `List`, which uses `join` for
# `Json.toStr`'s `JsonList` case above) — a plain free function taking the
# `Buffer` as a captured-by-reference parameter instead of a closure sidesteps
# `libs/std/list.plum`'s `join` note on closures: reassigning a CAPTURED
# variable inside a closure body doesn't persist across that closure's own
# repeated calls, but mutating a heap object's fields THROUGH a captured
# reference (as `appendJsonPair` does to `buf`) does.
fun jsonMapToStr(pairs: Map[Str, Json]) -> Str =
  buf := Buffer()
  pairs.each(|k, v| appendJsonPair(buf, k, v))
  return buf.toStr()

fun appendJsonPair(buf: Buffer, k: Str, v: Json) -> Unit =
  if !buf.isEmpty()
    buf.write(",")
  buf.write("\"")
  buf.write(escapeJsonStr(k))
  buf.write("\":")
  buf.write(v.toStr())

fun exponentIndex(text: Str) -> Int =
  e := text.indexOf("e")
  e >= 0 ? e : text.indexOf("E")

fun escapeJsonStr(s: Str) -> Str =
  escapeJsonFrom(s, 0, s.length())

# Byte-at-a-time except for the pass-through case, which copies a whole
# UTF-8 codepoint at once — copying a lone continuation byte as its own
# "character" (as a naive byte-by-byte pass-through would) would corrupt any
# non-ASCII text, the same hazard documented on `Str`'s own `escapeChar` in
# `libs/std/str.plum`.
fun escapeJsonFrom(s: Str, i: Int, len: Int) -> Str =
  if i >= len
    return ""
  b := s.byteAt(i)
  if b == 34
    return "\\\"" + escapeJsonFrom(s, i + 1, len)
  if b == 92
    return "\\\\" + escapeJsonFrom(s, i + 1, len)
  if b == 10
    return "\\n" + escapeJsonFrom(s, i + 1, len)
  if b == 13
    return "\\r" + escapeJsonFrom(s, i + 1, len)
  if b == 9
    return "\\t" + escapeJsonFrom(s, i + 1, len)
  if b == 8
    return "\\b" + escapeJsonFrom(s, i + 1, len)
  if b == 12
    return "\\f" + escapeJsonFrom(s, i + 1, len)
  if b < 0x20
    return "\\u00" + hexByteStr(b) + escapeJsonFrom(s, i + 1, len)
  n := s.codePointByteLength(i)
  return s.sub(i, i + n) + escapeJsonFrom(s, i + n, len)

fun hexByteStr(b: Int) -> Str =
  hexDigitChar(b / 16) + hexDigitChar(b % 16)

fun hexDigitChar(d: Int) -> Str =
  match d
    0 => "0"
    1 => "1"
    2 => "2"
    3 => "3"
    4 => "4"
    5 => "5"
    6 => "6"
    7 => "7"
    8 => "8"
    9 => "9"
    10 => "a"
    11 => "b"
    12 => "c"
    13 => "d"
    14 => "e"
    15 => "f"
    _ => "0"

fun hexDigitValue(b: Int) -> Int =
  if b >= 48 && b <= 57
    return b - 48
  if b >= 97 && b <= 102
    return b - 97 + 10
  if b >= 65 && b <= 70
    return b - 65 + 10
  return -1

fun isSpace(c: Int) -> Bool =
  c == 32 || c >= 9 && c <= 13

fun isDigit(c: Int) -> Bool =
  c >= 48 && c <= 57

# ---- regression tests ----

fun parseFloatValue(s: Str) -> Float =
  r := Json.fromStr(s)
  match r
    Ok(v) => v.asFloat().unwrapOr(-1.0f)
    Err(_) => -1.0f

test "parses null, true, and false"
  assert Json.fromStr("null").unwrap().isNull()
  assert Json.fromStr("true").unwrap().asBool().unwrap()
  assert !Json.fromStr("false").unwrap().asBool().unwrap()

test "parses integers, decimals, negatives, and exponents"
  assert parseFloatValue("42") == 42.0f
  assert parseFloatValue("-7") == -7.0f
  assert parseFloatValue("3.5") == 3.5f
  assert parseFloatValue("1e2") == 100.0f
  assert parseFloatValue("1.5e-2") == 0.015f

test "parses a plain string, with escapes"
  assert Json.fromStr("\"hello\"").unwrap().asStr().unwrap() == "hello"
  assert Json.fromStr("\"a\\nb\"").unwrap().asStr().unwrap() == "a\nb"
  assert Json.fromStr("\"\\u0041\"").unwrap().asStr().unwrap() == "A"

test "parses an array of mixed values"
  v := Json.fromStr("[1, \"two\", true, null]").unwrap()
  items := v.asList().unwrap()
  assert items.length() == 4
  assert items.get(0).unwrap().asFloat().unwrap() == 1.0f
  assert items.get(1).unwrap().asStr().unwrap() == "two"
  assert items.get(2).unwrap().asBool().unwrap()
  assert items.get(3).unwrap().isNull()

test "parses a nested object"
  v := Json.fromStr("{{\"a\": 1, \"b\": {{\"c\": [1, 2, 3]}}}}").unwrap()
  pairs := v.asMap().unwrap()
  assert pairs.get("a").unwrap().asFloat().unwrap() == 1.0f
  nested := pairs.get("b").unwrap().asMap().unwrap()
  inner := nested.get("c").unwrap().asList().unwrap()
  assert inner.length() == 3
  assert inner.get(2).unwrap().asFloat().unwrap() == 3.0f

test "reports an error on malformed input"
  r := Json.fromStr("{{\"a\": }}")
  assert r.isErr()

test "round-trips a value through toStr and back"
  original := "[1,\"two\",true,null,{{\"k\":3.5}}]"
  v := Json.fromStr(original).unwrap()
  again := Json.fromStr(v.toStr()).unwrap()
  assert again.toStr() == v.toStr()

test "toStr escapes special characters in strings"
  v := JsonStr("a\"b\\c\nd")
  assert v.toStr() == "\"a\\\"b\\\\c\\nd\""