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

import { Actor, Assertion, Entity, Facet, Handle, Ref, Turn } from '../runtime/actor.js';
import { BytesLike, Decoder, Dictionary, embed, encode, IdentityMap, mapEmbeddeds, stringify, underlying, Value } from '@preserves/core';
import * as IO from '../gen/protocol.js';
import { wireRefEmbeddedType } from './protocol.js';
import { queueTask } from '../runtime/task.js';
import { attenuate } from '../runtime/rewrite.js';
import { fromAttenuation, WireRef } from '../gen/sturdy.js';

export class WireSymbol {
    count = 0;

    constructor (
        public side: Membrane,
        public oid: IO.Oid,
        public ref: Ref,
    ) {}

    drop(): void {
        this.count--;
        if (this.count === 0) {
            this.side.byOid.delete(this.oid);
            this.side.byRef.delete(this.ref);
        }
    }
};

export class SyncPeerEntity implements Entity {
    readonly relay: Relay;
    readonly peer: Ref;
    readonly handleMap = new IdentityMap<Handle, Handle>();
    pins: Array<WireSymbol> = [];

    constructor(relay: Relay, peer: Ref) {
        this.relay = relay;
        this.peer = peer;
    }

    assert(assertion: Assertion, handle: Handle): void {
        this.handleMap.set(handle, Turn.active.assert(this.peer, assertion));
    }

    retract(handle: Handle): void {
        Turn.active.retract(this.handleMap.get(handle)!);
        this.handleMap.delete(handle);
    }

    message(body: Assertion): void {
        // We get to vanish from the indexes now
        this.pins.forEach(e => e.drop());
        Turn.active.message(this.peer, body);
    }

    sync(peer: Ref): void {
        Turn.active._sync(this.peer, peer);
    }
}

export class RelayEntity implements Entity {
    readonly relay: Relay;
    readonly oid: IO.Oid;

    constructor(relay: Relay, oid: IO.Oid) {
        this.relay = relay;
        this.oid = oid;
    }

    send(m: IO.Event<WireRef>): void {
        this.relay.send(this.oid, m);
    }

    assert(assertion: Assertion, handle: Handle): void {
        this.send(IO.Event.Assert(IO.Assert({
            assertion: this.relay.register(this.oid, assertion, handle),
            handle
        })))
    }

    retract(handle: Handle): void {
        this.relay.deregister(handle);
        this.send(IO.Event.Retract(IO.Retract(handle)));
    }

    message(body: Assertion): void {
        this.send(IO.Event.Message(IO.Message(this.relay.register(null, body, null))));
    }

    sync(peer: Ref): void {
        const peerEntity = new SyncPeerEntity(this.relay, peer);
        this.relay.grabImportedOid(this.oid, peerEntity.pins);
        const ior = this.relay.rewriteRefOut(Turn.ref(peerEntity), false, peerEntity.pins);
        this.send(IO.Event.Sync(IO.Sync(ior)));
    }
}

export type WhichTable = "byOid" | "byRef";

export class Membrane {
    readonly byOid = new IdentityMap<IO.Oid, WireSymbol>();
    readonly byRef = new IdentityMap<Ref, WireSymbol>();

    grab<Table extends WhichTable>(table: Table,
                                   key: Parameters<Membrane[Table]['get']>[0],
                                   transient: boolean,
                                   f: () => WireSymbol): WireSymbol;
    grab<Table extends WhichTable>(table: Table,
                                   key: Parameters<Membrane[Table]['get']>[0],
                                   transient: boolean,
                                   f: null): WireSymbol | null;
    grab<Table extends WhichTable>(table: Table,
                                   key: Parameters<Membrane[Table]['get']>[0],
                                   transient: boolean,
                                   f: (() => WireSymbol) | null): WireSymbol | null {
        let e = this[table].get(key as any);
        if (e === void 0) {
            if (f === null) return null;
            e = f();
            this.byRef.set(e.ref, e);
            this.byOid.set(e.oid, e);
        }
        if (!transient) e.count++;
        return e;
    }
}

export const INERT_REF: Ref = {
    relay: Actor.boot(() => Turn.active.stop()).root,
    target: {},
};

export type PacketWriter = (bs: Uint8Array) => void;

export interface RelayOptions {
    packetWriter: PacketWriter;
    setup(r: Relay): void;
    debug?: boolean;
    trustPeer?: boolean;
    initialOid?: IO.Oid;
    initialRef?: Ref;
    nextLocalOid?: IO.Oid;
}

export class Relay {
    readonly facet: Facet;
    readonly w: PacketWriter;
    readonly inboundAssertions = new IdentityMap<Handle, {
        localHandle: Handle,
        pins: Array<WireSymbol>,
    }>();
    readonly outboundAssertions = new IdentityMap<Handle, Array<WireSymbol>>();
    readonly exported = new Membrane();
    readonly imported = new Membrane();
    readonly peer: Ref | null;
    nextLocalOid: IO.Oid = 0;
    pendingTurn: IO.Turn<WireRef> = [];
    debug: boolean;
    trustPeer: boolean;

    readonly decoder = new Decoder(void 0, {
        includeAnnotations: false,
        embeddedDecode: wireRefEmbeddedType,
    });

    constructor(options: RelayOptions) {
        this.facet = Turn.activeFacet;
        this.w = options.packetWriter;
        this.debug = options.debug ?? false;
        this.trustPeer = options.trustPeer ?? true;

        this.facet.preventInertCheck();
        options.setup(this);

        if (options.initialRef !== void 0) {
            this.rewriteRefOut(options.initialRef, false, []);
        }

        this.peer = (options.initialOid !== void 0)
            ? this.rewriteRefIn(WireRef.mine(options.initialOid), [])
            : null;

        if (options.nextLocalOid !== void 0) {
            this.nextLocalOid = (options.nextLocalOid === 0) ? 1 : options.nextLocalOid;
        }
    }

    register(targetRemoteOid: IO.Oid, assertion: Assertion, handle: Handle): Value<WireRef>;
    register(targetRemoteOid: null, assertion: Assertion, handle: null): Value<WireRef>;
    register(targetRemoteOid: IO.Oid | null, assertion: Assertion, handle: Handle | null): Value<WireRef> {
        const transient = (handle === null);
        const pins: Array<WireSymbol> = [];
        const rewritten = mapEmbeddeds(assertion, r => embed(this.rewriteRefOut(r, transient, pins)));
        if (handle !== null) {
            if (targetRemoteOid !== null /* belt and suspenders */) {
                this.grabImportedOid(targetRemoteOid, pins);
            }
            this.outboundAssertions.set(handle, pins);
        }
        return rewritten;
    }

    deregister(handle: Handle): void {
        (this.outboundAssertions.get(handle) ?? []).forEach(e => e.drop());
        this.outboundAssertions.delete(handle);
    }

    grabImportedOid(oid: IO.Oid, pins: Array<WireSymbol>) {
        const e = this.imported.grab("byOid", oid, false, null);
        if (e === null) {
            throw new Error("Internal error: import table missing entry for oid " + oid);
        }
        pins.push(e);
    }

    grabExportedOid(oid: IO.Oid, pins: Array<WireSymbol>): Ref {
        const e = this.exported.grab("byOid", oid, false, null);
        if (e === null) return INERT_REF;
        pins.push(e);
        return e.ref;
    }

    rewriteRefOut(r: Ref, transient: boolean, pins: Array<WireSymbol>): WireRef {
        if (r.target instanceof RelayEntity && r.target.relay === this) {
            if (r.attenuation === void 0 || r.attenuation.length === 0) {
                // No extra conditions on this reference since it was sent to us.
                this.grabImportedOid(r.target.oid, pins);
                return WireRef.yours({ oid: r.target.oid, attenuation: [] });
            } else {
                // This reference has been attenuated since it was sent to us.
                // Do we trust the peer to enforce such attenuation on our behalf?
                if (this.trustPeer) {
                    this.grabImportedOid(r.target.oid, pins);
                    return WireRef.yours({ oid: r.target.oid, attenuation: r.attenuation });
                } else {
                    // fall through: treat the attenuated ref as a local ref, and re-export it.
                }
            }
        }

        const e = this.exported.grab(
            "byRef", r, transient, () => {
                if (transient) throw new Error("Cannot send transient reference");
                return new WireSymbol(this.exported, this.nextLocalOid++, r);
            });
        pins.push(e);
        return WireRef.mine(e.oid);
    }

    rewriteRefIn(n: WireRef, pins: Array<WireSymbol>): Ref {
        switch (n._variant) {
            case 'yours': {
                const e = this.exported.grab("byOid", n.oid, false, null);
                if (e === null) {
                    return INERT_REF;
                } else {
                    pins.push(e);
                    const r = e.ref;
                    if (n.attenuation.length === 0) {
                        return r;
                    } else {
                        type AttenuatedRef = Ref & { __attenuations?: Dictionary<any, Ref> };
                        const ar = r as AttenuatedRef;
                        if (ar.__attenuations === void 0) {
                            ar.__attenuations = new Dictionary();
                        }
                        return ar.__attenuations.getOrSet(fromAttenuation(n.attenuation), () =>
                            attenuate(r, ... n.attenuation));
                    }
                }
            }
            case 'mine': {
                const e = this.imported.grab("byOid", n.oid, false, () =>
                    new WireSymbol(this.imported, n.oid, Turn.ref(new RelayEntity(this, n.oid))));
                pins.push(e);
                return e.ref;
            }
        }
    }

    send(remoteOid: IO.Oid, m: IO.Event<WireRef>): void {
        if (this.pendingTurn.length === 0) {
            queueTask(() => {
                if (this.debug) console.log('OUT', stringify(IO.fromTurn(this.pendingTurn)));
                this.w(underlying(encode(IO.fromTurn(this.pendingTurn), {
                    canonical: true,
                    embeddedEncode: wireRefEmbeddedType,
                })));
                this.pendingTurn = [];
            });
        }
        this.pendingTurn.push(IO.TurnEvent({ oid: remoteOid, event: m }));
    }

    accept(bs: BytesLike): void {
        Turn.for(this.facet, () => {
            this.decoder.write(bs);
            while (true) {
                const rawTurn = this.decoder.try_next();
                if (rawTurn === void 0) break;
                const wireTurn = IO.toTurn(rawTurn);
                if (wireTurn === void 0) throw new Error("Bad IO.Turn");
                if (this.debug) console.log('IN', stringify(rawTurn));
                wireTurn.forEach(v => this.handle(v.oid, v.event));
            }
        });
    }

    rewriteIn(a: Value<WireRef>): [Assertion, Array<WireSymbol>]
    {
        const pins: Array<WireSymbol> = [];
        const rewritten = mapEmbeddeds(a, r => embed(this.rewriteRefIn(r, pins)));
        return [rewritten, pins];
    }

    handle(localOid: IO.Oid, m: IO.Event<WireRef>) {
        switch (m._variant) {
            case 'Assert': {
                const [a, pins] = this.rewriteIn(m.value.assertion);
                const r = this.grabExportedOid(localOid, pins);
                this.inboundAssertions.set(m.value.handle, {
                    localHandle: Turn.active.assert(r, a),
                    pins,
                });
                break;
            }
            case 'Retract': {
                const remoteHandle = m.value.handle;
                const h = this.inboundAssertions.get(remoteHandle);
                if (h === void 0) throw new Error(`Peer retracted invalid handle ${remoteHandle}`);
                this.inboundAssertions.delete(remoteHandle);
                h.pins.forEach(e => e.drop());
                Turn.active.retract(h.localHandle);
                break;
            }
            case 'Message': {
                const [a, pins] = this.rewriteIn(m.value.body);
                if (pins.length > 0) throw new Error("Cannot receive transient reference");
                const r = this.exported.byOid.get(localOid)?.ref;
                if (r) Turn.active.message(r, a);
                break;
            }
            case 'Sync': {
                const pins: Array<WireSymbol> = [];
                const r = this.grabExportedOid(localOid, pins);
                const k = this.rewriteRefIn(m.value.peer, pins);
                Turn.active.sync(r).then(() => {
                    Turn.active.message(k, true);
                    pins.forEach(e => e.drop());
                });
                break;
            }
        }
    }
}