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

import std/[algorithm, assertions, options, sequtils, tables, typetraits]

import preserves
import ./protocols/dataspacePatterns
from ./actors import Cap

export dataspacePatterns.`$`, PatternKind, DCompoundKind, AnyAtomKind

type
  AnyAtom = dataspacePatterns.AnyAtom
  DBind = dataspacePatterns.DBind
  DCompound = dataspacePatterns.DCompound
  DCompoundArr = dataspacePatterns.DCompoundArr
  DCompoundDict = dataspacePatterns.DCompoundDict
  DCompoundRec = dataspacePatterns.DCompoundRec
  DLit = dataspacePatterns.DLit
  Pattern* = dataspacePatterns.Pattern

iterator orderedEntries*(dict: DCompoundDict): (Value, Pattern) =
  ## Iterate a `DCompoundDict` in Preserves order.
  ## Values captured from a dictionary are represented as an
  ## array of values ordered by their former key, so using an
  ## ordered iterator is sometimes essential.
  var keys = dict.entries.keys.toSeq
  sort(keys, preserves.cmp)
  for k in keys:
    yield(k, dict.entries.getOrDefault(k))
      # getOrDefault doesn't raise and we know the keys will match

proc toPattern(d: sink DBind): Pattern =
  Pattern(orKind: PatternKind.DBind, dbind: d)

proc toPattern(d: sink DLit): Pattern =
  Pattern(orKind: PatternKind.DLit, dlit: d)

proc toPattern(aa: sink AnyAtom): Pattern =
  DLit(value: aa).toPattern

proc toPattern(d: sink DCompound): Pattern =
  Pattern(orKind: PatternKind.DCompound, dcompound: d)

proc toPattern(d: sink DCompoundRec): Pattern =
  DCompound(orKind: DCompoundKind.rec, rec: d).toPattern

proc toPattern(d: sink DCompoundArr): Pattern =
  DCompound(orKind: DCompoundKind.arr, arr: d).toPattern

proc toPattern(d: sink DCompoundDict): Pattern =
  DCompound(orKind: DCompoundKind.dict, dict: d).toPattern

proc drop*(): Pattern {.inline.} = Pattern(orKind: PatternKind.DDiscard)
  ## Create a pattern to match any value without capture.

proc grab*(): Pattern {.inline.} = DBind(pattern: drop()).toPattern
  ## Create a pattern to capture any value.

proc grab*(pr: Value): Pattern =
  ## Convert a `Preserve` value to a `Pattern`.
  runnableExamples:
    from std/unittest import check
    import preserves
    check:
      $(grab parsePreserves"""<foo "bar" #"00" [0 1 2.0] {maybe: #t} <_>>""") ==
        """<rec foo [<lit "bar"> <lit #"00"> <arr [<lit 0> <lit 1> <lit 2.0>]> <dict {maybe: <lit #t>}> <_>]>"""

  case pr.kind
  of pkBoolean:
    AnyAtom(orKind: AnyAtomKind.`bool`, bool: pr.bool).toPattern
  of pkFloat:
    AnyAtom(orKind: AnyAtomKind.`double`, double: pr.float).toPattern
  of pkRegister:
    AnyAtom(orKind: AnyAtomKind.`int`, int: pr.register).toPattern
  of pkString:
    AnyAtom(orKind: AnyAtomKind.`string`, string: pr.string).toPattern
  of pkByteString:
    AnyAtom(orKind: AnyAtomKind.`bytes`, bytes: pr.bytes).toPattern
  of pkSymbol:
    AnyAtom(orKind: AnyAtomKind.`symbol`, symbol: pr.symbol).toPattern
  of pkRecord:
    if (pr.isRecord("_") and pr.arity == 0) or (pr.isRecord("bind") and pr.arity == 1):
      drop()
    else:
      DCompoundRec(
        label: pr.label,
        fields: map[Value, Pattern](pr.fields, grab)).toPattern
  of pkSequence:
    DCompoundArr(items: map(pr.sequence, grab)).toPattern
  of pkSet:
    raiseAssert "cannot construct a pattern over a set literal"
  of pkDictionary:
    var dict = DCompoundDict()
    for key, val in pr.pairs: dict.entries[key] = grab val
    dict.toPattern
  of pkEmbedded:
    if pr.embeddedRef.isNil: drop()
    else:
      AnyAtom(orKind: AnyAtomKind.`embedded`, embedded: pr.embeddedRef).toPattern
  else:
    raise newException(ValueError, "cannot generate a pattern for unhandled Value type")

proc grab*[T](x: T): Pattern =
  ## Construct a `Pattern` from value of type `T`.
  runnableExamples:
    from std/unittest import check
    check:
      $grab(true) == "<lit #t>"
      $grab(3.14) == "<lit 3.14>"
      $grab([0, 1, 2, 3]) == "<arr [<lit 0> <lit 1> <lit 2> <lit 3>]>"
  grab(x.toPreserves)

proc grabType*(typ: static typedesc): Pattern =
  ## Derive a `Pattern` from type `typ`.
  ## This works for `tuple` and `object` types but in the
  ## general case will return a wildcard binding.
  runnableExamples:
    import preserves
    from std/unittest import check
    check:
      $grabType(array[3, int]) ==
        """<arr [<bind <_>> <bind <_>> <bind <_>>]>"""
    type
      Point = tuple[x: int; y: int]
      Rect {.preservesRecord: "rect".} = tuple[a: Point; B: Point]
      ColoredRect {.preservesDictionary.} = tuple[color: string; rect: Rect]
    check:
      $(grabType Point) ==
        "<arr [<bind <_>> <bind <_>>]>"
      $(grabType Rect) ==
        "<rec rect [<arr [<bind <_>> <bind <_>>]> <arr [<bind <_>> <bind <_>>]>]>"
      $(grabType ColoredRect) ==
        "<dict {color: <bind <_>> rect: <rec rect [<arr [<bind <_>> <bind <_>>]> <arr [<bind <_>> <bind <_>>]>]>}>"
  when typ is ref:
    grabType(pointerBase(typ))
  elif typ.hasPreservesRecordPragma:
    var rec = DCompoundRec(label: typ.recordLabel.toSymbol)
    for _, f in fieldPairs(default typ):
      add(rec.fields, grabType(typeof f))
    result = rec.toPattern
  elif typ.hasPreservesDictionaryPragma:
    var dict = DCompoundDict()
    for key, val in fieldPairs(default typ):
      dict.entries[key.toSymbol] = grabType(typeof val)
    dict.toPattern
  elif typ is tuple:
    var arr = DCompoundArr()
    for _, f in fieldPairs(default typ):
      add(arr.items, grabType(typeof f))
    arr.toPattern
  elif typ is array:
    var arr = DCompoundArr()
    arr.items.setLen(len(typ))
    for e in arr.items.mitems: e = grab()
    arr.toPattern
  else:
    grab()

proc fieldCount(T: typedesc): int =
  for _, _ in fieldPairs(default T):
    inc result

proc dropType*(typ: static typedesc): Pattern =
  ## Derive a `Pattern` from type `typ` without any bindings.
  when typ is ref:
    dropType(pointerBase(typ))
  elif typ.hasPreservesRecordPragma:
    var rec = DCompoundRec(label: typ.recordLabel.toSymbol)
    rec.fields.setLen(fieldCount typ)
    for i, _ in rec.fields:
      rec.fields[i] = drop()
    result = rec.toPattern
  elif typ.hasPreservesDictionaryPragma:
    DCompoundDict().toPattern
  elif typ is tuple:
    var arr = DCompoundArr()
    arr.items.setLen(len typ)
    for i, _ in arr.items:
      arr.items[i] = drop()
    arr.toPattern
  elif typ is array:
    var arr = DCompoundArr()
    arr.items.setLen(len(typ))
    for e in arr.items.mitems: e = drop()
    arr.toPattern
  else:
    drop()

proc lookup(bindings: openArray[(int, Pattern)]; i: int): Pattern =
  for (j, b) in bindings:
    if i == j: return b
  return drop()

proc grab*(typ: static typedesc; bindings: sink openArray[(int, Pattern)]): Pattern =
  ## Construct a `Pattern` from type `typ` with pattern `bindings` by integer offset.
  when typ is ptr | ref:
    grab(pointerBase(typ), bindings)
  elif typ.hasPreservesRecordPragma:
    var rec = DCompoundRec(label: typ.recordLabel.toSymbol)
    rec.fields.setLen(fieldCount typ)
    var i: int
    for _, f in fieldPairs(default typ):
      rec.fields[i] = lookup(bindings, i)
      inc i
    result = rec.toPattern
  elif typ is tuple:
    var arr = DCompoundArr()
    arr.items.setLen(fieldCount typ)
    var i: int
    for _, f in fieldPairs(default typ):
      arr.items[i] = lookup(bindings, i)
      inc i
    result = arr.toPattern
  else:
    {.error: "grab with indexed bindings not implemented for " & $typ.}

proc grab*(typ: static typedesc; bindings: sink openArray[(Value, Pattern)]): Pattern =
  ## Construct a `Pattern` from type `typ` with dictionary field `bindings`.
  when typ.hasPreservesDictionaryPragma:
    DCompoundDict(entries: bindings.toTable).toPattern
  else:
    {.error: "grab with dictionary bindings not implemented for " & $typ.}

proc grabLit*(): Pattern =
  runnableExamples:
    from std/unittest import check
    check:
      $grabLit() == """<rec lit [<bind <_>>]>"""
  grabType(dataspacePatterns.DLit)

proc grabDict*(): Pattern =
  grabType(dataspacePatterns.DCompoundDict)

proc unpackLiterals*(pr: Value): Value =
  result = pr
  apply(result) do (pr: var Value):
    if pr.isRecord("lit", 1) or pr.isRecord("dict", 1) or pr.isRecord("arr", 1) or pr.isRecord("set", 1):
      pr = pr.record[0]

proc inject*(pat: Pattern; bindings: openArray[(int, Pattern)]): Pattern =
  ## Construct a `Pattern` from `pat` with injected overrides from `bindings`.
  ## Injects are made at offsets indexed by the discard (`<_>`) patterns in `pat`.
  proc inject(pat: Pattern; bindings: openArray[(int, Pattern)]; offset: var int): Pattern =
    case pat.orKind
    of PatternKind.DDiscard:
      result = pat
      for (off, injection) in bindings:
        if off == offset:
          result = injection
          break
      inc offset
    of PatternKind.DBind:
      let bindOff = offset
      result = pat
      result.dbind.pattern = inject(pat.dbind.pattern, bindings, offset)
      if result.orKind == PatternKind.DBind:
        for (off, injection) in bindings:
          if (off == bindOff) and (result.dbind.pattern == injection):
            result = result.dbind.pattern
            break # promote the injected pattern over the bind
    of PatternKind.DLit:
      result = pat
    of PatternKind.DCompound:
      result = pat
      case pat.dcompound.orKind
      of DCompoundKind.rec:
        var fields = mapIt(pat.dcompound.rec.fields):
          inject(it, bindings, offset)
        result = DCompoundRec(label: result.dcompound.rec.label, fields: fields).toPattern
      of DCompoundKind.arr:
        var items = mapIt(pat.dcompound.arr.items):
          inject(it, bindings, offset)
        result = DCompoundArr(items: items).toPattern
      of DCompoundKind.dict:
        var dict = DCompoundDict()
        for key, val in pat.dcompound.dict.entries:
          dict.entries[key] = inject(val, bindings, offset)
        result = toPattern(dict)
  var offset = 0
  inject(pat, bindings, offset)

proc inject*(pat: Pattern; bindings: openArray[(Value, Pattern)]): Pattern =
  ## Inject `bindings` into a dictionary pattern.
  assert pat.orKind == PatternKind.DCompound
  assert pat.dcompound.orKind == DCompoundKind.dict
  result = pat
  for (key, val) in bindings:
    result.dcompound.dict.entries[key] = val

proc grabRecord*(label: Value, fields: varargs[Pattern]): Pattern =
  runnableExamples:
    from std/unittest import check
    import syndicate/actors, preserves
    check:
      $grabRecord("Says".toSymbol, grab(), grab()) ==
        """<rec Says [<bind <_>> <bind <_>>]>"""
  DCompoundRec(label: label, fields: fields.toSeq).toPattern

proc grabRecord*(label: string, fields: varargs[Pattern]): Pattern =
  ## Sugar for creating record patterns.
  ## `label` is converted to a symbol value.
  grabRecord(label.toSymbol, fields)

proc grabDictionary*(bindings: sink openArray[(Value, Pattern)]): Pattern =
  ## Construct a pattern that grabs some dictionary pairs.
  DCompoundDict(entries: bindings.toTable).toPattern

proc grabDictionary*(bindings: sink openArray[(string, Pattern)]): Pattern =
  ## Construct a pattern that grabs some dictionary pairs.
  ## Keys are converted from strings to symbols.
  result = DCompoundDict().toPattern
  for (key, val) in bindings.items:
    result.dcompound.dict.entries[key.toSymbol] = val

proc depattern(comp: DCompound; values: var seq[Value]; index: var int): Value

proc depattern(pat: Pattern; values: var seq[Value]; index: var int): Value =
  case pat.orKind
  of PatternKind.DDiscard:
    discard
  of PatternKind.DBind:
    if index < values.len:
      result = move values[index]
      inc index
  of PatternKind.DLit:
    result = pat.dlit.value.toPreserves
  of PatternKind.DCompound:
    result = depattern(pat.dcompound, values, index)

proc depattern(comp: DCompound; values: var seq[Value]; index: var int): Value =
  case comp.orKind
  of DCompoundKind.rec:
    result = initRecord(comp.rec.label, comp.rec.fields.len)
    for i, f in comp.rec.fields:
      result[i] = depattern(f, values, index)
  of DCompoundKind.arr:
    result = initSequence(comp.arr.items.len)
    for i, e in comp.arr.items:
      result[i] = depattern(e, values, index)
  of DCompoundKind.dict:
    result = initDictionary(Cap)
    for key, val in comp.dict.entries:
      result[key] = depattern(val, values, index)

proc depattern*(pat: Pattern; values: sink seq[Value]): Value =
  ## Convert a `Pattern` to a `Value` while replacing binds with `values`.
  var index: int
  depattern(pat, values, index)

type Literal*[T] = object
  ## A wrapper type to deserialize patterns to native values.
  value*: T

proc fromPreservesHook*[T](lit: var Literal[T]; pr: Value): bool =
  var pat: Pattern
  pat.fromPreserves(pr) and lit.value.fromPreserves(depattern(pat, @[]))

proc toPreservesHook*[T](lit: Literal[T]): Value =
  lit.value.grab.toPreserves

type
  Path* = seq[Value]
  Paths* = seq[Path]
  Captures* = seq[Value]
  Analysis* = tuple
    constPaths: Paths
    constValues: seq[Value]
    capturePaths: Paths

func walk(result: var Analysis; path: var Path; p: Pattern)

func walk(result: var Analysis; path: var Path; key: int|Value; pat: Pattern) =
  path.add(key.toPreserves)
  walk(result, path, pat)
  discard path.pop

func walk(result: var Analysis; path: var Path; p: Pattern) =
  case p.orKind
  of PatternKind.DCompound:
    case p.dcompound.orKind
    of DCompoundKind.rec:
      for k, e in p.dcompound.rec.fields: walk(result, path, k, e)
    of DCompoundKind.arr:
       for k, e in p.dcompound.arr.items: walk(result, path, k, e)
    of DCompoundKind.dict:
      for k, e in p.dcompound.dict.orderedEntries: walk(result, path, k, e)
  of PatternKind.DBind:
    result.capturePaths.add(path)
    walk(result, path, p.dbind.pattern)
  of PatternKind.DDiscard: discard
  of PatternKind.DLit:
    result.constPaths.add(path)
    result.constValues.add(p.dlit.value.toPreserves)

func analyse*(p: Pattern): Analysis =
  var path: Path
  walk(result, path, p)

func projectPaths*(v: Value; paths: Paths): Option[Captures] =
  var res = newSeq[Value](paths.len)
  for i, path in paths:
    var vv = step(v, path)
    if vv.isSome: res[i] = get(vv)
    else: return
  some res

proc matches*(pat: Pattern; pr: Value): bool =
  let analysis = analyse(pat)
  assert analysis.constPaths.len == analysis.constValues.len
  for i, path in analysis.constPaths:
    let v = step(pr, path)
    if v.isNone: return false
    if analysis.constValues[i] != v.get: return false
  for path in analysis.capturePaths:
    if isNone step(pr, path): return false
  true

func capture*(pat: Pattern; pr: Value): seq[Value] =
  let analysis = analyse(pat)
  assert analysis.constPaths.len == analysis.constValues.len
  for i, path in analysis.constPaths:
    let v = step(pr, path)
    if v.isNone : return @[]
    if analysis.constValues[i] != v.get: return @[]
  for path in analysis.capturePaths:
    let v = step(pr, path)
    if v.isNone: return @[]
    result.add(get v)

when isMainModule:
  let txt = readAll stdin
  if txt != "":
    let
      v = parsePreserves(txt)
      pat = grab v
    stdout.writeLine(pat)