# 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]

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: seq[Path]): seq[Value] =
  result = newSeq[Value](paths.len)
  for i, path in paths:
    var vv = projectPath(v, path)
    if vv.isSome: result[i] = get(vv)

type Class = distinct string

proc `$`(cls: Class): string {.borrow.}
proc hash(cls: Class): Hash {.borrow.}
proc `==`(x, y: Class): bool {.borrow.}

proc classOf*(v: Value): Class =
  if v.isRecord:
    result = Class $v.label & "/" & $v.arity
  elif v.isSequence:
    result = Class $v.len
  elif v.isDictionary:
    result = Class "{}"

proc classOf*(p: Pattern): Class =
  if p.orKind == PatternKind.DCompound:
    case p.dcompound.orKind
    of DCompoundKind.rec:
      result = Class $p.dcompound.rec.label & "/" & $p.dcompound.rec.fields.len
    of DCompoundKind.arr:
      result = Class $p.dcompound.arr.items.len
    of DCompoundKind.dict:
      result = Class "{}"

proc step(value, index: Value): Value =
  result = value[index]

type
  EventKind = enum addedEvent, removedEvent, messageEvent

  AssertionCache = HashSet[Value]

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

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

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

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

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

from strutils import toHex
proc `$`(node: Node): string =
  toHex(cast[ByteAddress](unsafeAddr node[]), 5)

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

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

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)
        nextClass = classOf nextValue
      if nextClass != Class"":
        let nextNode = table.getOrDefault(nextClass)
        if not nextNode.isNil:
          walk(nextNode, turn, outerValue, event, push(nextStack, nextValue))

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

proc extend(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 = extend(node, popCount, stepIndex, pat.dbind.pattern, path)
  of PatternKind.DCompound:
    let selector: Selector = (popCount, stepIndex,)
    var table = node.edges.getOrDefault(selector)
    if table.isNil:
      table = newTable[Class, Node]()
      node.edges[selector] = table
    let class = classOf pat
    result.nextNode = table.getOrDefault(class)
    if result.nextNode.isNil:
      new result.nextNode
      new result.nextNode.continuation
      table[class] = result.nextNode
      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
    result.popCount = 0
    template walkKey(pat: Pattern; stepIndex: Value) =
      path.add(stepIndex)
      result = extend(result.nextNode, popCount, stepIndex, pat, path)
      discard path.pop()
    case pat.dcompound.orKind
    of DCompoundKind.rec:
      for k, e in pat.dcompound.rec.fields: walkKey(e, k.toPreserve(Ref))
    of DCompoundKind.arr:
       for k, e in pat.dcompound.arr.items: walkKey(e, k.toPreserve(Ref))
    of DCompoundKind.dict:
      for k, e in pat.dcompound.dict.entries: walkKey(e, k)
    result.popCount.inc
    when not defined(release):
      assert not node.edges[selector][classOf pat].isNil

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

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

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)