# SPDX-FileCopyrightText: ☭ Emery Hemingway
# SPDX-License-Identifier: Unlicense

import std/[algorithm, hashes, math, options, sets, sequtils, tables]

import bigints

type
  PreserveKind* = enum
    pkBoolean, pkFloat, pkRegister, pkBigInt, pkString, pkByteString, pkSymbol,
        pkRecord, pkSequence, pkSet, pkDictionary, pkEmbedded

const
  atomKinds* = {pkBoolean, pkFloat, pkRegister, pkBigInt, pkString, pkByteString, pkSymbol}
  compoundKinds* = {pkRecord, pkSequence, pkSet, pkDictionary}

type Symbol* = distinct string
proc `<`*(x, y: Symbol): bool {.borrow.}
proc `==`*(x, y: Symbol): bool {.borrow.}
proc hash*(s: Symbol): Hash {.borrow.}
proc len*(s: Symbol): int {.borrow.}

type
  Atom* = object
    ## Atomic Preserves value.
    ## Useful when a `const Value` is required.
    case kind*: PreserveKind
    of pkBoolean:
      bool*: bool
    of pkFloat:
      float*: float
    of pkRegister:
      register*: int
    of pkBigInt:
      bigint*: BigInt
    of pkString:
      string*: string
    of pkByteString:
      bytes*: seq[byte]
    of pkSymbol:
      symbol*: Symbol
    else:
      discard

  Value* = object
    case kind*: PreserveKind
    of pkBoolean:
      bool*: bool
    of pkFloat:
      float*: float
    of pkRegister:
      register*: int
    of pkBigInt:
      bigint*: BigInt
    of pkString:
      string*: string
    of pkByteString:
      bytes*: seq[byte]
    of pkSymbol:
      symbol*: Symbol
    of pkRecord:
      record*: seq[Value] # label is last
    of pkSequence:
      sequence*: seq[Value]
    of pkSet:
      set*: seq[Value]
        # TODO: HashSet
    of pkDictionary:
      dict*: seq[DictEntry]
        # TODO: Tables
    of pkEmbedded:
      embeddedRef*: EmbeddedRef
    embedded*: bool
      ## Flag to mark embedded Preserves value

  DictEntry* = tuple[key: Value, val: Value]

  EmbeddedRef* = ref RootObj
  EmbeddedObj* = RootObj
    ## Object refs embedded in Preserves `Value`s must inherit from `EmbeddedObj`.
    ## At the moment this is just an alias to `RootObj` but this may change in the future.

func `===`[T: SomeFloat](a, b: T): bool =
  ## Compare where Nan == NaN.
  let class = a.classify
  (class == b.classify) and ((class notin {fcNormal,fcSubnormal}) or (a == b))

func `==`*(x, y: Value): bool =
  ## Check `x` and `y` for equivalence.
  if x.kind == y.kind and x.embedded == y.embedded:
    case x.kind
    of pkBoolean:
      result = x.bool == y.bool
    of pkFloat:
      result = x.float === y.float
    of pkRegister:
      result = x.register == y.register
    of pkBigInt:
      result = x.bigint == y.bigint
    of pkString:
      result = x.string == y.string
    of pkByteString:
      result = x.bytes == y.bytes
    of pkSymbol:
      result = x.symbol == y.symbol
    of pkRecord:
      result = x.record.len == y.record.len
      for i in 0..x.record.high:
        if not result: break
        result = result and (x.record[i] == y.record[i])
    of pkSequence:
      for i, val in x.sequence:
        if y.sequence[i] != val: return false
      result = true
    of pkSet:
      result = x.set.len == y.set.len
      for i in 0..x.set.high:
        if not result: break
        result = result and (x.set[i] == y.set[i])
    of pkDictionary:
      result = x.dict.len == y.dict.len
      for i in 0..x.dict.high:
        if not result: break
        result = result and
          (x.dict[i].key == y.dict[i].key) and
          (x.dict[i].val == y.dict[i].val)
    of pkEmbedded:
      result = x.embeddedRef == y.embeddedRef

proc `<`(x, y: string | seq[byte]): bool =
  for i in 0 .. min(x.high, y.high):
    if x[i] < y[i]: return true
    if x[i] != y[i]: return false
  x.len < y.len

proc `<`*(x, y: Value): bool =
  ## Preserves have a total order over values. Check if `x` is ordered before `y`.
  if x.embedded != y.embedded:
    result = y.embedded
  elif x.kind != y.kind:
    result = x.kind < y.kind
  else:
    case x.kind
    of pkBoolean:
      result = (not x.bool) and y.bool
    of pkFloat:
      result = x.float < y.float
    of pkRegister:
      result = x.register < y.register
    of pkBigInt:
      result = x.bigint < y.bigint
    of pkString:
      result = x.string < y.string
    of pkByteString:
      result = x.bytes < y.bytes
    of pkSymbol:
      result = x.symbol < y.symbol
    of pkRecord:
      if x.record[x.record.high] < y.record[y.record.high]: return true
      for i in 0..<min(x.record.high, y.record.high):
        if x.record[i] < y.record[i]: return true
        if x.record[i] == y.record[i]: return false
      result = x.record.len < y.record.len
    of pkSequence:
      for i in 0..min(x.sequence.high, y.sequence.high):
        if x.sequence[i] < y.sequence[i]: return true
        if x.sequence[i] != y.sequence[i]: return false
      result = x.sequence.len < y.sequence.len
    of pkSet:
      for i in 0..min(x.set.high, y.set.high):
        if x.set[i] < y.set[i]: return true
        if x.set[i] != y.set[i]: return false
      result = x.set.len < y.set.len
    of pkDictionary:
      for i in 0..min(x.dict.high, y.dict.high):
        if x.dict[i].key < y.dict[i].key: return true
        if x.dict[i].key == y.dict[i].key:
          if x.dict[i].val < y.dict[i].val: return true
          if x.dict[i].val != y.dict[i].val: return false
      result = x.dict.len < y.dict.len
    of pkEmbedded:
      result = x.embeddedRef < y.embeddedRef

func cmp*(x, y: Value): int =
  ## Compare by Preserves total ordering.
  if x == y: 0
  elif x < y: -1
  else: 1

proc sort*(pr: var Value) = sort(pr.sequence, cmp)
  ## Sort a Preserves array by total ordering.

proc hash*(pr: Value): Hash =
  ## Produce a `Hash` of `pr` for use with a `HashSet` or `Table`.
  var h = hash(pr.kind.int) !& hash(pr.embedded)
  case pr.kind
  of pkBoolean:
    h = h !& hash(pr.bool)
  of pkFloat:
    h = h !& hash(pr.float)
  of pkRegister:
    h = h !& hash(pr.register)
  of pkBigInt:
    h = h !& hash(pr.bigint)
  of pkString:
    h = h !& hash(pr.string)
  of pkByteString:
    h = h !& hash(pr.bytes)
  of pkSymbol:
    h = h !& hash(string pr.symbol)
  of pkRecord:
    for val in pr.record:
      h = h !& hash(val)
  of pkSequence:
    for val in pr.sequence:
      h = h !& hash(val)
  of pkSet:
    for val in pr.set.items:
      h = h !& hash(val)
  of pkDictionary:
    for (key, val) in pr.dict.items:
      h = h !& hash(key) !& hash(val)
  of pkEmbedded:
    h = h !& hash(cast[uint](addr pr.embeddedRef[]))
  !$h

proc `[]`*(pr: Value; i: int): Value =
  ## Select an indexed value from ``pr``.
  ## Only valid for records and sequences.
  case pr.kind
  of pkRecord: pr.record[i]
  of pkSequence: pr.sequence[i]
  else:
    raise newException(ValueError, "Preserves value is not indexable")

proc `[]=`*(pr: var Value; i: Natural; val: Value) =
  ## Assign an indexed value into ``pr``.
  ## Only valid for records and sequences.
  case pr.kind
  of pkRecord: pr.record[i] = val
  of pkSequence: pr.sequence[i] = val
  else:
    raise newException(ValueError, "Preserves value is not indexable")

proc `[]=`*(pr: var Value; key, val: Value) =
  ## Insert `val` by `key` in the Preserves dictionary `pr`.
  for i in 0..pr.dict.high:
    if key < pr.dict[i].key:
      insert(pr.dict, [(key, val, )], i)
      return
    elif key == pr.dict[i].key:
      pr.dict[i].val = val
      return
  pr.dict.add((key, val, ))

proc incl*(pr: var Value; key: Value) =
  ## Include `key` in the Preserves set `pr`.
  for i in 0..pr.set.high:
    if key < pr.set[i]:
      insert(pr.set, [key], i)
      return
  pr.set.add(key)

proc excl*(pr: var Value; key: Value) =
  ## Exclude `key` from the Preserves set `pr`.
  for i in 0..pr.set.high:
    if pr.set[i] == key:
      delete(pr.set, i..i)
      break