/// SPDX-License-Identifier: GPL-3.0-or-later
/// SPDX-FileCopyrightText: Copyright © 2016-2023 Tony Garnock-Jones <tonyg@leastfixedpoint.com>

import { Set, Dictionary, KeyedDictionary, IdentityMap, stringify, Value, embed } from '@preserves/core';
import { AnyValue, Assertion, Handle, Ref, Turn } from './actor.js';
import { Bag, ChangeDescription } from './bag.js';
import * as Stack from './stack.js';
import * as P from '../gen/dataspacePatterns.js';
import { Path, analysePattern, classOfCtor, classOfValue, step } from './pattern.js';

enum EventType {
    ADDED = +1,
    REMOVED = -1,
    MESSAGE = 0,
}

const _nop = function() {};

const INDENT = '    ';

export class Index {
    readonly allAssertions: Bag<Ref> = new Bag();
    readonly root: Node = new Node(new Continuation(new Set()));

    addObserver(pattern: P.Pattern, observer: Ref) {
        let {constPaths, constValues, capturePaths} = analysePattern(pattern);
        const continuation = this.root.extend(pattern);
        let constValMap = continuation.leafMap.get(constPaths);
        if (!constValMap) {
            constValMap = new Dictionary();
            continuation.cachedAssertions.forEach((a) => {
                const key = projectPaths(a, constPaths);
                let leaf = constValMap!.get(key);
                if (!leaf) {
                    leaf = new Leaf();
                    constValMap!.set(key, leaf);
                }
                leaf.cachedAssertions.add(a);
            });
            continuation.leafMap.set(constPaths, constValMap);
        }
        let leaf = constValMap.get(constValues);
        if (!leaf) {
            leaf = new Leaf();
            constValMap.set(constValues, leaf);
        }
        let observerGroup = leaf.observerGroups.get(capturePaths);
        if (!observerGroup) {
            const cachedCaptures = new Bag<Ref, Array<AnyValue>>();
            leaf.cachedAssertions.forEach((a) =>
                cachedCaptures._items.update(projectPaths(a, capturePaths), n => n! + 1, 0));
            observerGroup = new ObserverGroup(cachedCaptures);
            leaf.observerGroups.set(capturePaths, observerGroup);
        }
        const captureMap: KeyedDictionary<Array<AnyValue>, Handle, Ref> = new KeyedDictionary();
        observerGroup.observers.set(observer, captureMap);
        observerGroup.cachedCaptures.forEach((_count, captures) =>
            captureMap.set(captures, Turn.active.assert(observer, captures)));
    }

    removeObserver(pattern: P.Pattern, observer: Ref) {
        let {constPaths, constValues, capturePaths} = analysePattern(pattern);
        const continuation = this.root.extend(pattern);
        let constValMap = continuation.leafMap.get(constPaths);
        if (!constValMap) return;
        let leaf = constValMap.get(constValues);
        if (!leaf) return;
        let observerGroup = leaf.observerGroups.get(capturePaths);
        if (!observerGroup) return;
        const captureMap = observerGroup.observers.get(observer);
        if (captureMap) {
            captureMap.forEach((handle, _captures) => Turn.active.retract(handle));
            observerGroup.observers.delete(observer);
        }
        if (observerGroup.observers.size === 0) {
            leaf.observerGroups.delete(capturePaths);
        }
        if (leaf.isEmpty()) {
            constValMap.delete(constValues);
        }
        if (constValMap.size === 0) {
            continuation.leafMap.delete(constPaths);
        }
    }

    adjustAssertion(outerValue: Assertion, delta: number): boolean {
        switch (this.allAssertions.change(outerValue, delta)) {
            case ChangeDescription.ABSENT_TO_PRESENT:
                this.root.modify(
                    EventType.ADDED,
                    outerValue,
                    (c, v) => c.cachedAssertions.add(v),
                    (l, v) => l.cachedAssertions.add(v),
                    (h, vs) => {
                        if (h.cachedCaptures.change(vs, +1) === ChangeDescription.ABSENT_TO_PRESENT)
                            h.observers.forEach((captureMap, observer) =>
                                captureMap.set(vs, Turn.active.assert(observer, vs)));
                    });
                return true;

            case ChangeDescription.PRESENT_TO_ABSENT:
                this.root.modify(
                    EventType.REMOVED,
                    outerValue,
                    (c, v) => c.cachedAssertions.delete(v),
                    (l, v) => l.cachedAssertions.delete(v),
                    (h, vs) => {
                        if (h.cachedCaptures.change(vs, -1) === ChangeDescription.PRESENT_TO_ABSENT)
                            h.observers.forEach((captureMap, _observer) => {
                                Turn.active.retract(captureMap.get(vs));
                                captureMap.delete(vs);
                            });
                    });
                return true;

            default:
                return false;
        }
    }

    addAssertion(v: Assertion): boolean {
        return this.adjustAssertion(v, +1);
    }

    removeAssertion(v: Assertion): boolean {
        return this.adjustAssertion(v, -1);
    }

    deliverMessage(v: Assertion, leafCallback: (l: Leaf, v: Assertion) => void = _nop) {
        this.root.modify(EventType.MESSAGE, v, _nop, leafCallback, (h, vs) =>
            h.observers.forEach((_captureMap, observer) => Turn.active.message(observer, vs)));
    }

    dump() {
        console.log('INDEX');
        // console.log('allAssertions:');
        // dumpBag(this.allAssertions, 4);
        console.log('tree:');
        this.root.dump(INDENT);
        console.log();
    }
}

function dumpBag<T, V extends Value<T>>(b: Bag<T, V>, indent: string) {
    for (const [v, count] of b.entries()) {
        console.log(indent + stringify(v) + ' = ' + count);
    }
}

function dumpSet<V>(s: Set<V>, indent: string) {
    s.forEach(v => console.log(indent + stringify(v)));
}

type Selector = [number, AnyValue];

class Node {
    readonly continuation: Continuation;
    readonly edges: KeyedDictionary<Selector, { [shape: string]: Node }, Ref> = new KeyedDictionary();

    constructor(continuation: Continuation) {
        this.continuation = continuation;
    }

    extend(p: P.Pattern): Continuation {
        const path: Path = [];

        function walkNode(node: Node, popCount: number, stepIndex: AnyValue, p: P.Pattern): [number, Node]
        {
            switch (p._variant) {
                case 'DDiscard':
                case 'DLit':
                    return [popCount, node];
                case 'DBind':
                    return walkNode(node, popCount, stepIndex, p.value.pattern);
                case 'DCompound': {
                    const selector: Selector = [popCount, stepIndex];
                    let table = node.edges.get(selector);
                    if (!table) {
                        table = {};
                        node.edges.set(selector, table);
                    }
                    let cls = classOfCtor(p.value);
                    let nextNode = table[cls];
                    if (!nextNode) {
                        nextNode = new Node(new Continuation(
                            node.continuation.cachedAssertions.filter(
                                (a) => classOfValue(projectPath(a, path)) === cls)));
                        table[cls] = nextNode;
                    }
                    popCount = 0;
                    function walkKey(pp: P.Pattern, stepIndex: AnyValue) {
                        path.push(stepIndex);
                        [popCount, nextNode] = walkNode(nextNode, popCount, stepIndex, pp);
                        path.pop();
                    }
                    switch (p.value._variant) {
                        case 'rec': p.value.fields.forEach(walkKey); break;
                        case 'arr': p.value.items.forEach(walkKey); break;
                        case 'dict': p.value.entries.forEach(walkKey); break;
                    }
                    return [popCount + 1, nextNode];
                }
            }
        }

        return walkNode(this, 0, 0, p)[1].continuation;
    }

    modify(operation: EventType,
           outerValue: Assertion,
           m_cont: (c: Continuation, v: Assertion) => void,
           m_leaf: (l: Leaf, v: Assertion) => void,
           m_observerGroup: (h: ObserverGroup, vs: Array<Assertion>) => void)
    {
        function walkNode(node: Node, termStack: Stack.NonEmptyStack<Assertion>) {
            walkContinuation(node.continuation);
            node.edges.forEach((table, [popCount, stepIndex]) => {
                const nextStack = Stack.dropNonEmpty(termStack, popCount);
                const nextValue = step(nextStack.item, stepIndex);
                const nextClass = classOfValue(nextValue);
                const nextNode = nextClass && table[nextClass];
                if (nextNode) walkNode(nextNode, Stack.push(nextValue!, nextStack));
            });
        }

        function walkContinuation(continuation: Continuation) {
            m_cont(continuation, outerValue);
            continuation.leafMap.forEach((constValMap, constPaths) => {
                let constValues = projectPaths(outerValue, constPaths);
                let leaf = constValMap.get(constValues);
                if (!leaf && operation === EventType.ADDED) {
                    leaf = new Leaf();
                    constValMap.set(constValues, leaf);
                }
                if (leaf) {
                    m_leaf(leaf, outerValue);
                    leaf.observerGroups.forEach((observerGroup, capturePaths) => {
                        m_observerGroup(observerGroup, projectPaths(outerValue, capturePaths));
                    });
                    if (operation === EventType.REMOVED && leaf.isEmpty()) {
                        constValMap.delete(constValues);
                        if (constValMap.size === 0) {
                            continuation.leafMap.delete(constPaths);
                        }
                    }
                }
                return true;
            });
        }

        walkNode(this, Stack.push([outerValue], Stack.empty()));
    }

    dump(indent: string) {
        this.continuation.dump(indent);
        this.edges.forEach((shapeMap, selector) => {
            const [popCount, stepIndex] = selector;
            console.log(indent + `popCount ${popCount} stepIndex ${stringify(stepIndex)} -->`);
            for (const shape in shapeMap) {
                console.log(indent + INDENT + stringify(shape));
                shapeMap[shape].dump(indent + INDENT + INDENT);
            }
        });
    }
}

class Continuation {
    readonly cachedAssertions: Set<Ref>;
    readonly leafMap: KeyedDictionary<Array<Path>, Dictionary<Ref, Leaf>, Ref> = new KeyedDictionary();

    constructor(cachedAssertions: Set<Ref>) {
        this.cachedAssertions = cachedAssertions;
    }

    dump(indent: string) {
        dumpSet(this.cachedAssertions, indent);
        this.leafMap.forEach((valueMap, paths) => {
            valueMap.forEach((leaf, values) => {
                console.log(indent + `when ${stringify(paths)} == ${stringify(values)} -->`);
                leaf.dump(indent + INDENT);
            });
        });
    }
}

class Leaf {
    readonly cachedAssertions: Set<Ref> = new Set();
    readonly observerGroups: KeyedDictionary<Array<Path>, ObserverGroup, Ref> = new KeyedDictionary();

    isEmpty(): boolean {
        return this.cachedAssertions.size === 0 && this.observerGroups.size === 0;
    }

    dump(indent: string) {
        dumpSet(this.cachedAssertions, indent);
        this.observerGroups.forEach((observerGroup, paths) => {
            observerGroup.dump(paths, indent);
        });
    }
}

class ObserverGroup {
    readonly cachedCaptures: Bag<Ref, Array<AnyValue>>;
    readonly observers: IdentityMap<Ref, KeyedDictionary<Array<AnyValue>, Handle, Ref>> = new IdentityMap();

    constructor(cachedCaptures: Bag<Ref, Array<AnyValue>>) {
        this.cachedCaptures = cachedCaptures;
    }

    dump(paths: Path[], indent: string) {
        dumpBag(this.cachedCaptures, indent);
        this.observers.forEach((valueMap, observer) => {
            console.log(indent + `${stringify(embed(observer))} projecting ${stringify(paths)}`);
            valueMap.forEach((handle, values) => {
                console.log(indent + INDENT + `captured ${stringify(values)} handle ${handle}`);
            });
        });
    }
}

// Total by assumption that path is valid for v
function projectPath(v: AnyValue, path: Path): AnyValue {
    for (let index of path) {
        v = step(v, index)!;
    }
    return v;
}

// Total by assumption that paths are valid for v
function projectPaths(v: AnyValue, paths: Array<Path>): AnyValue[] {
  return paths.map((path) => projectPath(v, path));
}