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

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

import preserves
import ./protocols/[dataspacePatterns, dataspace]
from ./actors import Cap

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

type
  Pattern* = dataspacePatterns.Pattern

proc toPattern(b: sink PatternBind): Pattern =
  Pattern(orKind: PatternKind.`bind`, `bind`: b)

proc toPattern(l: sink PatternLit): Pattern =
  Pattern(orKind: PatternKind.`lit`, lit: l)

proc toPattern(g: sink PatternGroup): Pattern =
  Pattern(orKind: PatternKind.`group`, group: g)

proc toPattern(a: sink AnyAtom): Pattern =
  PatternLit(value: a).toPattern

proc grab*(p: sink Pattern): Pattern =
  PatternBind(pattern: p).toPattern

proc drop*(): Pattern = Pattern(orKind: PatternKind.`discard`)
  ## Create a pattern to match any value without capture.

proc grab*(): Pattern = drop().grab()
  ## Create a pattern to capture any value.

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

  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 pkBigInt:
    raiseAssert "cannot make a pattern over a big integer"
  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("_", 0):
      drop()
    elif pr.isRecord("bind", 1):
      pr.fields[0].drop
    else:
      var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.rec))
      group.`type`.rec.label = pr.label
      var i: int
      for v in pr.fields:
        group.entries[toPreserves i] = drop v
        inc i
      group.toPattern

  of pkSequence:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.arr))
    for i, v in pr.sequence:
      group.entries[toPreserves i] = drop v
    group.toPattern

  of pkSet:
    raiseAssert "cannot construct a pattern over a set literal"

  of pkDictionary:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.dict))
    for key, val in pr.pairs:
      group.entries[key] = drop val
    group.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 drop*[T](x: T): Pattern =
  ## Construct a `Pattern` from value of type `T`.
  ## This proc is called `drop` because the value `x` is matched but discarded.
  runnableExamples:
    from std/unittest import check
    check:
      $drop(true) == "<lit #t>"
      $drop(3.14) == "<lit 3.14>"
      $drop([0, 1, 2, 3]) == "<group <arr> {0: <lit 0> 1: <lit 1> 2: <lit 2> 3: <lit 3>}>"
  drop(x.toPreserves)

proc grab*[T](x: T): Pattern {.
    deprecated: "use drop unless you wish to capture the provided value".} =
  PatternBind(pattern: drop x).toPattern

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]) ==
        """<group <arr> {0: <bind <_>> 1: <bind <_>> 2: <bind <_>> 3: <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) ==
        "<group <arr> {0: <bind <_>> 1: <bind <_>>}>"
      $(grabType Rect) ==
        "<group <rec rect> {0: <group <arr> {0: <bind <_>> 1: <bind <_>>}> 1: <group <arr> {0: <bind <_>> 1: <bind <_>>}>}>"
      $(grabType ColoredRect) ==
        "<group <dict> {color: <bind <_>> rect: <group <rec rect> {0: <group <arr> {0: <bind <_>> 1: <bind <_>>}> 1: <group <arr> {0: <bind <_>> 1: <bind <_>>}>}>}>"
  when typ is ref:
    grabType(pointerBase(typ))
  elif typ.hasPreservesRecordPragma:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`rec`))
    group.`type`.rec.label = typ.recordLabel.toSymbol
    for _, f in fieldPairs(default typ):
      group.entries[group.entries.len.toPreserves] = grabType(typeof f)
    group.toPattern
  elif typ.hasPreservesDictionaryPragma:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`dict`))
    for key, val in fieldPairs(default typ):
      group.entries[key.toSymbol] = grabType(typeof val)
    group.toPattern
  elif typ is tuple:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`arr`))
    for _, f in fieldPairs(default typ):
      group.entries[group.entries.len.toPreserves] = grabType(typeof f)
    group.toPattern
  elif typ is array:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`arr`))
    for i in 0..len(typ):
      group.entries[toPreserves i] = grab()
    group.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 group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`rec`))
    group.`type`.rec.label = typ.recordLabel.toSymbol
    let high = typ.fieldCount.pred
    if high >= 0: group.entries[high.toPreserves] = drop()
    group.toPattern
  elif typ.hasPreservesDictionaryPragma:
    PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`dict`)).toPattern
  elif typ is tuple or typ is array:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`arr`))
    let high = typ.fieldCount.pred
    if high >= 0: group.entries[high.toPreserves] = drop()
    group.toPattern
  else:
    drop()

proc bindEntries(group: var PatternGroup; bindings: openArray[(int, Pattern)]) =
  ## Set `bindings` for a `group`.
  for (i, pat) in bindings: group.entries[toPreserves i] = pat

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 group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`rec`))
    group.`type`.rec.label = typ.recordLabel.toSymbol
    bindEntries(group, bindings)
    group.toPattern
  elif typ is tuple:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`arr`))
    bindEntries(group, bindings)
    group.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:
    var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`dict`))
    for key, val in bindinds: group.entries[key] = val
    group.toPattern
  else:
    {.error: "grab with dictionary bindings not implemented for " & $typ.}

proc grabLit*(): Pattern =
  runnableExamples:
    from std/unittest import check
    check:
      $grabLit() == """<group <rec lit> {0: <bind <_>>}>"""
  grabType(dataspacePatterns.PatternLit)

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

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*(pattern: sink Pattern; p: Pattern; path: varargs[Value, toPreserves]): Pattern =
  ## Inject `p` inside `pattern` at `path`.
  ## Injects are made at offsets indexed by the discard (`<_>`) patterns in `pat`.
  proc inject(pat: var Pattern; path: openarray[Value]) =
    if len(path) == 0:
      pat = p
    elif pat.orKind != PatternKind.`group`:
      raise newException(ValueError, "cannot inject along specified path")
    else:
      inject(pat.group.entries[path[0]], path[1..path.high])
  result = pattern
  inject(result, path)

proc grabRecord*(label: Value, fields: varargs[Pattern]): Pattern =
  runnableExamples:
    from std/unittest import check
    import preserves
    check:
      $grabRecord("Says".toSymbol, grab(), grab()) ==
        """<group <rec Says> {0: <bind <_>> 1: <bind <_>>}>"""
  var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`rec`))
  group.`type`.rec.label = label
  for i, f in fields: group.entries[toPreserves i] = f
  group.toPattern

proc grabRecord*(label: Value, fields: sink openArray[(int, Pattern)]): Pattern =
  runnableExamples:
    from std/unittest import check
    import preserves
    check:
      $grabRecord("Says".toSymbol, {3: grab(), 4: grab()}) ==
        """<group <rec Says> {3: <bind <_>> 4: <bind <_>>}>"""
  var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`rec`))
  group.`type`.rec.label = label
  for (i, p) in fields: group.entries[toPreserves i] = p
  group.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.
  var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`dict`))
  for (key, val) in bindings: group.entries[key] = val
  group.toPattern

proc grabDictionary*(bindings: sink openArray[(string, Pattern)]): Pattern =
  ## Construct a pattern that grabs some dictionary pairs.
  ## Keys are converted from strings to symbols.
  var group = PatternGroup(`type`: GroupType(orKind: GroupTypeKind.`dict`))
  for (key, val) in bindings: group.entries[toSymbol key] = val
  group.toPattern

proc depattern(group: PatternGroup; 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.`discard`:
    discard
  of PatternKind.`bind`:
    if index < values.len:
      result = move values[index]
      inc index
  of PatternKind.`lit`:
    result = pat.`lit`.value.toPreserves
  of PatternKind.`group`:
    result = depattern(pat.group, values, index)

proc depattern(group: PatternGroup; values: var seq[Value]; index: var int): Value =
  case group.`type`.orKind
  of GroupTypeKind.rec:
    result = initRecord(group.`type`.rec.label, group.entries.len)
    var i: int
    for key, val in group.entries:
      if i.fromPreserves key:
        result[i] = depattern(val, values, index)
  of GroupTypeKind.arr:
    result = initSequence(group.entries.len)
    var i: int
    for key, val in group.entries:
      if i.fromPreserves key:
        result[i] = depattern(val, values, index)
  of GroupTypeKind.dict:
    result = initDictionary(Cap)
    for key, val in group.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`.
  runnableExamples:
    from std/unittest import check
    import preserves
    type Foo {.preservesRecord: "foo".} = object
      a, b: int
    let pat = grabType Foo
    let val = depattern(pat, @[1.toPreserves, 5.toPreserves])
    check $val == "<foo 1 5>"
  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

func isGroup(pat: Pattern): bool =
  pat.orKind == PatternKind.`group`

func isMetaDict(pat: Pattern): bool =
  pat.orKind == PatternKind.`group` and
    pat.group.type.orKind == GroupTypeKind.dict

proc metaApply(result: var Pattern; pat: Pattern; path: openarray[Value], offset: int) =
  if offset == path.len:
    result = pat
  elif result.isGroup and result.group.entries[1.toPreserves].isMetaDict:
    if offset == path.high:
      result.group.entries[1.toPreserves].group.entries[path[offset]] = pat
    else:
      metaApply(result.group.entries[1.toPreserves].group.entries[path[offset]], pat, path, succ offset)
  else:
    assert result.isGroup, "non-group: " & $result
    assert result.group.entries[1.toPreserves].isMetaDict, "non-meta-dict: " & $result.group.entries[1.toPreserves]
    raise newException(ValueError, "cannot inject into non-group pattern " & $result)

proc observePattern*(pat: Pattern; injects: openarray[(seq[Value], Pattern)]): Pattern =
  result = dropType Observe
  var meta = pat.toPreserves.drop
  for (path, pat) in injects:
    metaApply(meta, pat, path, 0)
  result.group.entries[0.toPreserves] = meta

type
  Path* = seq[Value]
  Paths* = seq[Path]
  Captures* = seq[Value]
  Analysis* = tuple
    presentPaths: Paths
    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: Value; pat: Pattern) =
  path.add(key)
  walk(result, path, pat)
  discard path.pop

func walk(result: var Analysis; path: var Path; p: Pattern) =
  case p.orKind
  of PatternKind.group:
    for k, v in p.group.entries: walk(result, path, k, v)
  of PatternKind.`bind`:
    result.capturePaths.add(path)
    walk(result, path, p.`bind`.pattern)
  of PatternKind.`discard`:
    result.presentPaths.add(path)
  of PatternKind.`lit`:
    result.constPaths.add(path)
    result.constValues.add(p.`lit`.value.toPreserves)

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

func checkPresence*(v: Value; present: Paths): bool =
  result = true
  for path in present:
    if not result: break
    result = step(v, path).isSome

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
  result = checkPresence(pr, analysis.presentPaths)
  if result:
    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 step(pr, path).isNone: return false

proc capture*(pat: Pattern; pr: Value): seq[Value] =
  let analysis = analyse(pat)
  assert analysis.constPaths.len == analysis.constValues.len
  if checkPresence(pr, analysis.presentPaths):
    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:
  stdout.writeLine stdin.readAll.parsePreserves.grab