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

import std/[hashes, lists, options, sets, tables]
import preserves
import ./actors, ./bags, ./patterns

type
  Value = Preserve[Ref]
  Path = seq[Value]
  ClassKind = enum classNone, classRecord, classSequence, classDictionary
  Class = object
    kind: ClassKind
    label: Value
    arity: int

func classOf(v: Value): Class =
  case v.kind
  of pkRecord: Class(kind: classRecord, label: v.label, arity: v.arity)
  of pkSequence: Class(kind: classSequence, arity: v.len)
  of pkDictionary: Class(kind: classDictionary)
  else: Class(kind: classNone)

proc classOf(p: Pattern): Class =
  if p.orKind == PatternKind.DCompound:
    case p.dcompound.orKind
    of DCompoundKind.rec:
      Class(kind: classRecord,
        label: p.dcompound.rec.label,
        arity: p.dcompound.rec.fields.len)
    of DCompoundKind.arr:
      Class(kind: classSequence, arity: p.dcompound.arr.items.len)
    of DCompoundKind.dict:
      Class(kind: classDictionary)
  else:
    Class(kind: classNone)

type
  EventKind = enum addedEvent, removedEvent, messageEvent

  AssertionCache = HashSet[Value]

  Paths = seq[Path]

  ObserverGroup = ref object # Endpoints
    cachedCaptures: Bag[seq[Value]]
    observers: Table[Ref, TableRef[seq[Value], Handle]]

  Leaf = ref object
    cachedAssertions: AssertionCache
    observerGroups: Table[Paths, ObserverGroup]

  Continuation = ref object
    cachedAssertions: AssertionCache
    leafMap: Table[Paths, TableRef[seq[Value], Leaf]]

  Selector = tuple[popCount: int; index: Value]

  Node = ref object
    edges: Table[Selector, TableRef[Class, Node]]
    continuation: Continuation

func isEmpty(leaf: Leaf): bool =
  leaf.cachedAssertions.len == 0 and leaf.observerGroups.len == 0

type
  ContinuationProc = proc (c: Continuation; v: Value) {.gcsafe.}
  LeafProc = proc (l: Leaf; v: Value) {.gcsafe.}
  ObserverProc = proc (turn: var Turn; group: ObserverGroup; vs: seq[Value]) {.gcsafe.}

type TermStack = SinglyLinkedList[Value]

proc push(stack: TermStack; val: Value): Termstack =
  result = stack
  prepend(result, val)

proc pop(stack: TermStack; n: int): TermStack =
  result = stack
  var n = n
  while n > 0:
    result.remove(result.head)
    assert not stack.head.isNil, "popped too far"
    dec n

proc top(stack: TermStack): Value =
  assert not stack.head.isNil, "stack is empty"
  stack.head.value

proc modify(node: Node; turn: var Turn; outerValue: Value; event: EventKind;
    modCont: ContinuationProc; modLeaf: LeafProc; modObs: ObserverProc) =

  proc walk(turn: var Turn; cont: Continuation; outerValue: Value; event: EventKind) =
    modCont(cont, outerValue)
    for constPaths, constValMap in cont.leafMap.pairs:
      let constVals = projectPaths(outerValue, constPaths)
      var leaf = constValMap.getOrDefault(constVals)
      if leaf.isNil and event == addedEvent:
        new leaf
        constValMap[constVals] = leaf
      if not leaf.isNil:
        modLeaf(leaf, outerValue)
        for capturePaths, observerGroup in leaf.observerGroups.pairs:
          modObs(turn, observerGroup, projectPaths(outerValue, capturePaths))
    # TODO: cleanup dead leaves

  proc walk(node: Node; turn: var Turn; outerValue: Value; event: EventKind; termStack: TermStack) =
    walk(turn, node.continuation, outerValue, event)
    for selector, table in node.edges:
      let
        nextStack = pop(termStack, selector.popCount)
        nextValue = step(nextStack.top, selector.index)
      if nextValue.isSome:
        let nextClass = classOf(get nextValue)
        if nextClass.kind != classNone:
          let nextNode = table.getOrDefault(nextClass)
          if not nextNode.isNil:
            walk(nextNode, turn, outerValue, event, push(nextStack, get nextValue))

  var stack: TermStack
  walk(node, turn, outerValue, event, push(stack, @[outerValue].toPreserve(Ref)))

proc getOrNew[A, B, C](t: var Table[A, TableRef[B, C]], k: A): TableRef[B, C] =
  result = t.getOrDefault(k)
  if result.isNil:
    result = newTable[B, C]()
    t[k] = result

iterator pairs(dc: DCompound): (Value, Pattern) =
  case dc.orKind
  of DCompoundKind.rec:
    for i, p in dc.rec.fields:
      yield (toPreserve(i, Ref), p,)
  of DCompoundKind.arr:
    for i, p in dc.arr.items:
      yield (toPreserve(i, Ref), p,)
  of DCompoundKind.dict:
    for pair in dc.dict.entries.pairs:
      yield pair

proc extendWalk(node: Node; popCount: Natural; stepIndex: Value; pat: Pattern; path: var Path): tuple[popCount: Natural, nextNode: Node] =
  case pat.orKind
  of PatternKind.DDiscard, PatternKind.DLit:
    result = (popCount, node)
  of PatternKind.DBind:
    result = extendWalk(node, popCount, stepIndex, pat.dbind.pattern, path)
  of PatternKind.DCompound:
    let
      class = classOf pat
      selector: Selector = (popCount, stepIndex,)
      table = node.edges.getOrNew(selector)
    result.nextNode = table.getOrDefault(class)
    if result.nextNode.isNil:
      new result.nextNode
      table[class] = result.nextNode
      new result.nextNode.continuation
      for a in node.continuation.cachedAssertions:
        var v = projectPath(a, path)
        if v.isSome and class == classOf(get v):
          result.nextNode.continuation.cachedAssertions.incl a
    for i, p in pat.dcompound.pairs:
      add(path, i)
      result = extendWalk(result.nextNode, result.popCount, i, p, path)
      discard pop(path)
    inc(result.popCount)

proc extend(node: var Node; pat: Pattern): Continuation =
  var path: Path
  extendWalk(node, 0, toPreserve(0, Ref), pat, path).nextNode.continuation

type
  Index* = object
    allAssertions: Bag[Value]
    root: Node
    observerCount: int

proc initIndex*(): Index =
  Index(root: Node(continuation: Continuation()))

proc add*(index: var Index; turn: var Turn; pattern: Pattern; observer: Ref) =
  let
    analysis = analyse pattern
    continuation = index.root.extend pattern
  var constValMap = continuation.leafMap.getOrDefault(analysis.constPaths)
  if constValMap.isNil:
    new constValMap
    for a in continuation.cachedAssertions:
      let key = projectPaths(a, analysis.constPaths)
      var leaf = constValMap.getOrDefault(key)
      if leaf.isNil:
        new leaf
        constValMap[key] = leaf
      leaf.cachedAssertions.incl(a)
    continuation.leafMap[analysis.constPaths] = constValMap
  var leaf = constValMap.getOrDefault(analysis.constValues)
  if leaf.isNil:
    new leaf
    constValMap[analysis.constValues] = leaf
  var observerGroup = leaf.observerGroups.getOrDefault(analysis.capturePaths)
  if observerGroup.isNil:
    new observerGroup
    for a in leaf.cachedAssertions:
      discard observerGroup.cachedCaptures.change(projectPaths(a, analysis.capturePaths), +1)
    leaf.observerGroups[analysis.capturePaths] = observerGroup
  var captureMap = newTable[seq[Value], Handle]()
  for (count, captures) in observerGroup.cachedCaptures:
    captureMap[captures] = publish(turn, observer, captures)
  observerGroup.observers[observer] = captureMap

proc remove*(index: var Index; turn: var Turn; pattern: Pattern; observer: Ref) =
  var
    analysis = analyse pattern
    continuation = index.root.extend pattern
  let constValMap = continuation.leafMap.getOrDefault(analysis.constPaths)
  if not constValMap.isNil:
    let leaf = constValMap.getOrDefault(analysis.constValues)
    if not leaf.isNil:
      let observerGroup = leaf.observerGroups.getOrDefault(analysis.capturePaths)
      if not observerGroup.isNil:
        let captureMap = observerGroup.observers.getOrDefault(observer)
        if not captureMap.isNil:
          for handle in captureMap.values: retract(observer.target, turn, handle)
          observerGroup.observers.del(observer)
        if observerGroup.observers.len == 0:
          leaf.observerGroups.del(analysis.capturePaths)
        if leaf.isEmpty:
          constValMap.del(analysis.constValues)
        if constValMap.len == 0:
          continuation.leafMap.del(analysis.constPaths)

proc adjustAssertion*(index: var Index; turn: var Turn; outerValue: Value; delta: int): bool =
  case index.allAssertions.change(outerValue, delta)
  of cdAbsentToPresent:
    result = true
    proc modContinuation(c: Continuation; v: Value) =
      c.cachedAssertions.incl(v)
    proc modLeaf(l: Leaf; v: Value) =
      l.cachedAssertions.incl(v)
    proc modObserver(turn: var Turn; group: ObserverGroup; vs: seq[Value]) =
      if group.cachedCaptures.change(vs, +1) == cdAbsentToPresent:
        for (observer, captureMap) in group.observers.pairs:
          let a = vs.toPreserve(Ref)
          captureMap[vs] = publish(turn, observer, a)
          # TODO: this handle is coming from the facet?
    modify(index.root, turn, outerValue, addedEvent, modContinuation, modLeaf, modObserver)
  of cdPresentToAbsent:
    result = true
    proc modContinuation(c: Continuation; v: Value) =
      c.cachedAssertions.excl(v)
    proc modLeaf(l: Leaf; v: Value) =
      l.cachedAssertions.excl(v)
    proc modObserver(turn: var Turn; group: ObserverGroup; vs: seq[Value]) =
      if group.cachedCaptures.change(vs, -1) == cdPresentToAbsent:
        for (observer, captureMap) in group.observers.pairs:
          retract(observer.target, turn, captureMap[vs])
          captureMap.del(vs)
    modify(index.root, turn, outerValue, removedEvent, modContinuation, modLeaf, modObserver)
  else: discard

proc continuationNoop(c: Continuation; v: Value) = discard
proc leafNoop(l: Leaf; v: Value) = discard

proc add*(index: var Index; turn: var Turn; v: Assertion): bool =
  adjustAssertion(index, turn, v, +1)
proc remove*(index: var Index; turn: var Turn; v: Assertion): bool =
  adjustAssertion(index, turn, v, -1)

proc deliverMessage*(index: var Index; turn: var Turn; v: Value) =
  proc observersCb(turn: var Turn; group: ObserverGroup; vs: seq[Value]) =
    for observer in group.observers.keys: message(turn, observer, vs)
  index.root.modify(turn, v, messageEvent, continuationNoop, leafNoop, observersCb)