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

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

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

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

type
  Value = Preserve[Cap]
  AnyAtom = dataspacePatterns.AnyAtom[Cap]
  DBind = dataspacePatterns.DBind[Cap]
  DCompound = dataspacePatterns.DCompound[Cap]
  DCompoundArr = dataspacePatterns.DCompoundArr[Cap]
  DCompoundDict = dataspacePatterns.DCompoundDict[Cap]
  DCompoundRec = dataspacePatterns.DCompoundRec[Cap]
  DLit = dataspacePatterns.DLit[Cap]
  Pattern* = dataspacePatterns.Pattern[Cap]

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*[T](pr: Preserve[T]): 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>}> <_>]>"""

  assert not pr.embedded
  case pr.kind
  of pkBoolean:
    AnyAtom(orKind: AnyAtomKind.`bool`, bool: pr.bool).toPattern
  of pkFloat:
    AnyAtom(orKind: AnyAtomKind.`float`, float: pr.float).toPattern
  of pkDouble:
    AnyAtom(orKind: AnyAtomKind.`double`, double: pr.double).toPattern
  of pkSignedInteger:
    AnyAtom(orKind: AnyAtomKind.`int`, int: pr.int).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: cast[Preserve[Cap]](pr.label), # TODO: don't cast like this
        fields: map[Preserve[T], 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[cast[Preserve[Cap]](key)] = grab val
    dict.toPattern
  of pkEmbedded:
    # TODO: can patterns be constructed over embedded literals?
    drop()

proc grab*[T](val: 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 (toPreserve(val, Cap))

proc patternOfType(typ: static typedesc; `bind`: static bool): Pattern =
  when typ is ref:
    patternOfType(pointerBase(typ), `bind`)
  elif typ.hasPreservesRecordPragma:
    var rec = DCompoundRec(label: typ.recordLabel.tosymbol(Cap))
    for _, f in fieldPairs(default typ):
      add(rec.fields, patternOfType(typeof f, `bind`))
    result = rec.toPattern
  elif typ.hasPreservesDictionaryPragma:
    var dict = DCompoundDict()
    for key, val in fieldPairs(default typ):
      dict.entries[toSymbol(key, Cap)] = patternOfType(typeof val, `bind`)
    dict.toPattern
  elif typ is tuple:
    var arr = DCompoundArr()
    for _, f in fieldPairs(default typ):
      add(arr.items, patternOfType(typeof f, `bind`))
    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:
    if `bind`: grab()
    else: drop()

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 <_>>]>]>}>"
  patternOfType(typ, true)

proc dropType*(typ: static typedesc): Pattern =
  ## Derive a `Pattern` from type `typ` without any bindings.
  patternOfType(typ, false)

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

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

proc grab*(typ: static typedesc; bindings: sink openArray[(int, Pattern)]): Pattern =
  ## Construct a `Pattern` from type `typ` that selectively captures fields.
  when typ is ptr | ref:
    grab(pointerBase(typ), bindings)
  elif typ.hasPreservesRecordPragma:
    var rec = DCompoundRec(label: typ.recordLabel.tosymbol(Cap))
    rec.fields.setLen(fieldCount typ)
    var i: int
    for _, f in fieldPairs(default typ):
      rec.fields[i] = lookup(bindings, i, f)
      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, f)
      inc i
    result = arr.toPattern
  else:
    {.error: "grab with bindings not implemented for " & $typ.}

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

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

proc unpackLiterals*[E](pr: Preserve[E]): Preserve[E] =
  result = pr
  apply(result) do (pr: var Preserve[E]):
    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[(Preserve[Cap], 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: Preserve[Cap], fields: varargs[Pattern]): Pattern =
  runnableExamples:
    from std/unittest import check
    import syndicate/actors, preserves
    check:
      $grabRecord("Says".toSymbol(Cap), 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(Cap), 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(Cap)] = val

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.toPreserve(Cap))
  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.toPreserve(Cap))

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

func projectPath*(v: Value; path: Path): Option[Value] =
  result = some(v)
  for index in path:
    result = preserves.step(result.get, index)
    if result.isNone: break

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

func 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 = projectPath(pr, path)
    if v.isNone : return false
    if analysis.constValues[i] != v.get: return false
  for path in analysis.capturePaths:
    if isNone projectPath(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 = projectPath(pr, path)
    if v.isNone : return @[]
    if analysis.constValues[i] != v.get: return @[]
  for path in analysis.capturePaths:
    let v = projectPath(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)