# SPDX-License-Identifier: ISC

import ./assertions, ./bags, ./events
import preserves, preserves/records
import lists, options, sets, tables

type
  NonEmptySkeleton*[Shape] = object
    shape: Shape
    members: seq[Skeleton[Shape]]
  Skeleton*[Shape] = Option[NonEmptySkeleton[Shape]]

  Shape = string

  Value = Preserve

  HandlerCallback* = proc (event: EventKind; bindings: seq[Value]) {.gcsafe.}

  Path = seq[Natural]

  Analysis* = object
    skeleton: Skeleton[Shape]
    constPaths: seq[Path]
    constVals: seq[Value]
    capturePaths: seq[Path]

proc projectPath(v: Value; path: Path): Value =
  result = v
  for index in path:
    result = result[index]

proc projectPaths(v: Value; paths: seq[Path]): seq[Value] =
  result.setLen(paths.len)
  for i, path in paths: result[i] = projectPath(v, path)

proc analyzeAssertion*(a: Value): Analysis =
  var path: Path
  proc walk(analysis: var Analysis; a: Value): Skeleton[Shape] =
    if Discard.isClassOf a:
      discard
    elif Capture.isClassOf a:
      analysis.capturePaths.add(path)
      result = walk(analysis, a.fields[0])
    else:
      if a.kind == pkRecord:
        let class = classOf(a)
        result = some NonEmptySkeleton[Shape](shape: $class)
        path.add(0)
        var i: int
        for field in a.fields:
          path[path.high] = i
          result.get.members.add(walk(analysis, field))
          inc(i)
        discard path.pop
      else:
        analysis.constPaths.add(path)
        analysis.constVals.add(a)
  result.skeleton = walk(result, a)

type
  Handler = ref object
    cachedCaptures: Bag[seq[Value]]
    callbacks: HashSet[HandlerCallback]

  AssertionCache = HashSet[Value]

  Leaf = ref object
    cachedAssertions: AssertionCache
    handlerMap: Table[seq[Path], Handler]

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

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

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

using
  continuation: Continuation
  leaf: Leaf
  node: Node

proc isEmpty(leaf): bool =
  leaf.cachedAssertions.len == 0 and leaf.handlerMap.len == 0

type
  ContinuationProc = proc (c: Continuation; v: Value) {.gcsafe.}
  LeafProc = proc (l: Leaf; v: Value) {.gcsafe.}
  HandlerProc = proc (h: Handler; vs: seq[Value]) {.gcsafe.}

proc modify(node; operation: EventKind; outerValue: Value;
    mCont: ContinuationProc; mLeaf: LeafProc; mHandler: HandlerProc) =

  proc walkContinuation(continuation) {.gcsafe.}

  proc walkNode(node; termStack: SinglyLinkedList[seq[Value]]) =
    # TODO: use a seq for the stack?
    walkContinuation(node.continuation)
    for (selector, table) in node.edges.pairs:
      var nextStack = termStack
      for _ in 1..selector.popCount:
        nextStack.head = nextStack.head.next
      let nextValue = nextStack.head.value[selector.index]
      if nextValue.isRecord:
        let nextClass = classOf(nextValue)
        let nextNode = table.getOrDefault($nextClass)
        if not nextNode.isNil:
          nextStack.prepend(nextValue.record)
          walkNode(nextNode, nextStack)

  proc walkContinuation(continuation: Continuation) =
    mCont(continuation, outerValue)
    for (constPaths, constValMap) in continuation.leafMap.pairs:
      let constVals = projectPaths(outerValue, constPaths)
      let leaf = constValMap.getOrDefault(constVals)
      if leaf.isNil:
        if operation == addedEvent:
          constValMap[constVals] = Leaf()
      else:
        mLeaf(leaf, outerValue)
        for (capturePaths, handler) in leaf.handlerMap.pairs:
          mHandler(handler, projectPaths(outerValue, capturePaths))
        if operation == removedEvent and leaf.isEmpty:
          constValMap.del(constVals)
          if constValMap.len == 0:
            continuation.leafMap.del(constPaths)
  var stack: SinglyLinkedList[seq[Value]]
  stack.prepend(@[outerValue])
  walkNode(node, stack)

proc extend[Shape](node; skeleton: Skeleton[Shape]): Continuation =
  var path: Path

  proc walkNode(node; popCount, index: int; skeleton: Skeleton[Shape]): tuple[popCount: int, node: Node] =
    assert(not node.isNil)
    if skeleton.isNone:
      return (popCount, node)
    else:
      let selector: Selector = (popCount, index)
      var
        cls = skeleton.get.shape
      var
        table = node.edges.getOrDefault(selector)
      if table.isNil:
        table = newTable[string, Node]()
        node.edges[selector] = table
      var nextNode = table.getOrDefault(cls)
      if nextNode.isNil:
        nextNode = Node(continuation: Continuation())
        table[cls] = nextNode
        for a in node.continuation.cachedAssertions:
          if $classOf(projectPath(a, path)) == cls:
            nextNode.continuation.cachedAssertions.incl(a)
      block:
        var popCount, index: int
        path.add(index)
        for member in skeleton.get.members:
          (popCount, nextNode) = walkNode(nextNode, result.popCount, index, member)
          inc(index)
          discard path.pop()
          path.add(index)
        discard path.pop()
        result = (popCount.succ, nextNode)
  walkNode(node, 0, 0, skeleton).node.continuation

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

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

using index: Index

proc addHandler*(index; res: Analysis; callback: HandlerCallback) =
  assert(not index.root.isNil)
  let
    constPaths = res.constPaths
    constVals = res.constVals
    capturePaths = res.capturePaths
    continuation = index.root.extend(res.skeleton)
  var constValMap = continuation.leafMap.getOrDefault(constPaths)
  if constValMap.isNil:
    constValMap = newTable[seq[Value], Leaf]()
    continuation.leafMap[constPaths] = constValMap
    for a in continuation.cachedAssertions:
      let key = projectPaths(a, constPaths)
      var leaf = constValMap.getOrDefault(key)
      if leaf.isNil:
        new leaf
        constValMap[key] = leaf
      leaf.cachedAssertions.incl(a)
  var leaf = constValMap.getOrDefault(constVals)
  if leaf.isNil:
    new leaf
    constValMap[constVals] = leaf
  var handler = leaf.handlerMap.getOrDefault(capturePaths)
  if handler.isNil:
    new handler
    leaf.handlerMap[capturePaths] = handler
    for a in leaf.cachedAssertions:
      let a = projectPaths(a, capturePaths)
      if handler.cachedCaptures.contains(a):
        discard handler.cachedCaptures.change(a, +1)
  handler.callbacks.incl(callback)
  for captures, count in handler.cachedCaptures.pairs:
    callback(addedEvent, captures)

proc removeHandler*(index; res: Analysis; callback: HandlerCallback) =
  let continuation = index.root.extend(res.skeleton)
  try:
    let
      constValMap = continuation.leafMap[res.constPaths]
      leaf = constValMap[res.constVals]
      handler = leaf.handlerMap[res.capturePaths]
    handler.callbacks.excl(callback)
    if handler.callbacks.len == 0:
      leaf.handlerMap.del(res.capturePaths)
    if leaf.isEmpty:
      constValMap.del(res.constVals)
    if constValMap.len == 0:
      continuation.leafMap.del(res.constPaths)
  except KeyError: discard

proc adjustAssertion*(index: var Index; outerValue: Value; delta: int): ChangeDescription =
  result = index.allAssertions.change(outerValue, delta)
  case result
  of cdAbsentToPresent:
    index.root.modify(
      addedEvent,
      outerValue,
      (proc (c: Continuation; v: Value) = c.cachedAssertions.incl(v)),
      (proc (l: Leaf; v: Value) = l.cachedAssertions.incl(v)),
      (proc (h: Handler; vs: seq[Value]) =
        if h.cachedCaptures.change(vs, +1) == cdAbsentToPresent:
          #debugEcho " assertion of ", outerValue
          for cb in h.callbacks: cb(addedEvent, vs)))
  of cdPresentToAbsent:
    index.root.modify(
      removedEvent,
      outerValue,
      (proc (c: Continuation; v: Value) = c.cachedAssertions.excl(v)),
      (proc (l: Leaf; v: Value) = l.cachedAssertions.excl(v)),
      (proc (h: Handler; vs: seq[Value]) =
        if h.cachedCaptures.change(vs, -1) == cdPresentToAbsent:
          #debugEcho "retraction of ", outerValue
          for cb in h.callbacks: cb(removedEvent, vs)))
  else:
    discard

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

proc deliverMessage*(index; v: Value; leafCb: proc (l: Leaf; v: Value) {.gcsafe.}) =
  proc handlerCb(h: Handler; vs: seq[Value]) =
    for cb in h.callbacks: cb(messageEvent, vs)
  index.root.modify(messageEvent, v, continuationNoop, leafCb, handlerCb)

proc deliverMessage*(index; v: Value) =
  proc handlerCb(h: Handler; vs: seq[Value]) =
    for cb in h.callbacks: cb(messageEvent, vs)
  index.root.modify(messageEvent, v, continuationNoop, leafNoop, handlerCb)