import {
    Float,
    Reader,
    Record,
    Value as PreservesValue,
    embed,
    is,
    isEmbedded,
    stringify,
} from '@syndicate-lang/core';

export type Input = PreservesValue<any>[] | string;

export type Token = number | boolean | string | symbol | Token[];

export type Value<V extends VM<V>> = PreservesValue<Closure<V> | Primitive<V>>;
export type Closure<V extends VM<V>> = { env: EnvironmentChain<V>, code: Token[] };
export type Primitive<V extends VM<V>> = (this: V, ... args: Value<V>[]) => Value<V>[];
export type EnvironmentChain<V extends VM<V>> = null | { rib: Environment<V>, next: EnvironmentChain<V> };
export type Environment<V extends VM<V>> = { [key: string]: Value<V> };
export type Frame<V extends VM<V>> = Closure<V> & { ip: number };

export class FuelCell {
    fuel = 10000;
}

export class RuntimeError extends Error {}
export class FuelExhausted extends RuntimeError {}
export class StackUnderflow extends RuntimeError {}
export class TypeError extends RuntimeError {}

export class SyntaxError extends Error {}

export function asToken(x: PreservesValue<unknown>): Token {
    switch (typeof x) {
        case 'number':
        case 'boolean':
        case 'symbol':
        case 'string':
            return x;
        case 'object':
            if (Float.isFloat(x)) {
                return x.value;
            }
            if (Array.isArray(x) && !Record.isRecord(x)) {
                return x.map(asToken);
            }
            /* fall through */
        default:
            throw new SyntaxError('Invalid program token: ' + stringify(x));
    }
}

export class VM<Self extends VM<Self>> {
    stack: Value<Self>[] = [];
    rstack: Frame<Self>[] = [];
    fuel = new FuelCell();
    debug = false;

    static parse(program: string): Token[] {
        return new Reader(program).readToEnd().map(asToken);
    }

    constructor(program: Input, primitives: Environment<Self> = Primitives) {
        const code = typeof program === 'string' ? VM.parse(program) : program.map(asToken);
        this.invoke({ env: { rib: primitives, next: null }, code });
    }

    push(... vs: Value<Self>[]) {
        for (const v of vs) {
            if (v === void 0 || v === null) {
                throw new TypeError("Unexpected null/undefined");
            }
            this.stack.push(v);
        }
    }

    pop(): Value<Self> {
        const v = this.stack.pop();
        if (v === void 0) throw new StackUnderflow("Stack underflow");
        return v;
    }

    take(n: number): Value<Self>[] {
        if (this.stack.length < n) {
            throw new StackUnderflow("Stack underflow: need " +n+", have "+this.stack.length);
        }
        return n > 0 ? this.stack.splice(-n) : [];
    }

    invoke(c: Closure<Self> | Primitive<Self>) {
        if (typeof c === 'function') {
            this.push(... c.call(this as unknown as Self, ... this.take(c.length)));
        } else {
            this.popToPending(); // tail calls
            this.rstack.push({ env: { rib: {}, next: c.env }, code: c.code, ip: 0 });
        }
    }

    apply(who: string, v: Value<Self>) {
        if (isEmbedded(v)) {
            return this.invoke(v.embeddedValue);
        } else {
            throw new TypeError('Got non-callable in `'+who+'`: ' + stringify(v));
        }
    }

    popToPending(): boolean {
        while (this.rstack.length > 0) {
            if (this.frame.ip < this.frame.code.length) {
                return true;
            }
            this.rstack.pop();
        }
        return false;
    }

    get frame(): Frame<Self> {
        return this.rstack[this.rstack.length - 1];
    }

    lookup(name: string): Value<Self> {
        let env = this.frame.env;
        while (env !== null) {
            const v = env.rib[name];
            if (v !== void 0) return v;
            env = env.next;
        }
        throw new Error("Unknown operator: " + name);
    }

    step(): boolean {
        if (!this.popToPending()) return false;

        if (this.fuel.fuel <= 0) throw new FuelExhausted("Fuel exhausted");
        this.fuel.fuel--;

        const op = this.frame.code[this.frame.ip++];

        switch (typeof op) {
            case 'number':
            case 'boolean':
            case 'string':
                this.push(op);
                break;

            case 'symbol':
                const v = this.lookup(op.description!);
                if (isEmbedded(v)) {
                    this.invoke(v.embeddedValue);
                } else {
                    this.push(v);
                }
                break;

            case 'object':
                this.push(closure(this.frame.env, op));
                break;

            default:
                ((_: never) => { throw new Error("Unhandled token: " + _); })(op);
        }
        return true;
    }

    nextToken<T extends Token>(who: string, f: ((v: Token) => v is T) = ((_v): _v is T => true)): T {
        const t = this.frame.code[this.frame.ip++];
        if (typeof t === 'undefined') throw new SyntaxError("Missing token after `"+who+"`");
        if (!f(t)) throw new SyntaxError("Syntax error after `"+who+"`");
        return t;
    }

    nextLiteralSymbol(who: string): symbol {
        return this.nextToken(who, (v): v is symbol => typeof v === 'symbol');
    }

    exec() {
        if (this.debug) {
            while (this.step()) {
                console.log(this._summary());
                this.popToPending();
                console.log(this.frame?.code[this.frame.ip]);
            }
        } else {
            while (this.step()) {}
        }
    }

    _summary(): string {
        return `${this.rstack.length}: ${stringify(this.stack)}`;
    }
}

export function primitiveEnvironment<V extends VM<V>>(
    input: { [key: string]: Primitive<V> },
): Environment<V> {
    const output: Environment<V> = {};
    Object.entries(input).forEach(([k, f]) => {
        f.toString = () => `#${k}`;
        output[k] = embed(f);
    });
    return output;
}

export const D2R = Math.PI / 180;

export const Primitives: Environment<any> = primitiveEnvironment({
    '+'(a, b) { return [(a as number) + (b as number)]; },
    '-'(a, b) { return [(a as number) - (b as number)]; },
    '*'(a, b) { return [(a as number) * (b as number)]; },
    '/'(a, b) { return [(a as number) / (b as number)]; },
    '%'(a, b) { return [(a as number) % (b as number)]; },
    'neg'(v) { return [-(v as number)]; },

    'cos'(n) { return [Math.cos(n as number * D2R)]; },
    'sin'(n) { return [Math.sin(n as number * D2R)]; },
    'tan'(n) { return [Math.tan(n as number * D2R)]; },

    'to'() {
        const n_or_ns = this.nextToken('to', (v: any): v is (symbol | symbol[]) => {
            return typeof v === 'symbol' ||
                (Array.isArray(v) && v.every(w => typeof w === 'symbol'));
        });
        const ns = Array.isArray(n_or_ns) ? n_or_ns : [n_or_ns];
        const vs = this.take(ns.length);
        const env = this.frame.env!;
        ns.forEach((n, i) => env.rib[n.description!] = vs[i]);
        return [];
    },
    'quote'() {
        return [this.lookup(this.nextLiteralSymbol('quote').description!)];
    },
    'apply'(v) {
        this.apply('apply', v);
        return [];
    },

    'eq'(a, b) { return [is(a, b)]; },
    'lt'(a, b) { return [(a as any) < (b as any)]; },
    'gt'(a, b) { return [(a as any) > (b as any)]; },
    'le'(a, b) { return [(a as any) <= (b as any)]; },
    'ge'(a, b) { return [(a as any) >= (b as any)]; },

    'swap'(v, w) { return [w, v]; },
    'dup'(v) { return [v, v]; },
    'over'(v, w) { return [v, w, v]; },
    'rot'(v, w, x) { return [w, x, v]; },
    '-rot'(v, w, x) { return [x, v, w]; },
    'drop'(_v) { return []; },

    'take'(n) { return [this.take(n as number)]; },
    '++'(vs, ws) {
        (vs as Value<any>[]).push(... (ws as Value<any>[]));
        return [vs];
    },
    '!'(vs, w) {
        (vs as Value<any>[]).push(w);
        return [vs];
    },

    '?'(n, vs) { return [(vs as Value<any>[])[n as number]]; },
    'length'(a) { return [(a as Value<any>[]).length]; },
    'saveStack'() { return [this.take(this.stack.length)]; },
    'restoreStack'(vs) { return (vs as Value<any>[]); },
    'untake'(vs) { (vs as Value<any>[]).push((vs as Value<any>[]).length); return (vs as Value<any>[]); },

    'not'(v) { return [!v]; },

    'if'(v, t) { if (!!v) this.apply('if', t); return []; },
    'ifelse'(v, t, f) { if (!!v) this.apply('ifelse', t); else this.apply('ifelse', f); return []; },

    '*dump*'() { console.log(this._summary()); return []; },
});

function closure<V extends VM<V>>(env: EnvironmentChain<V>, code: Token[]): Value<V> {
    return embed({ env, code, toString: () => `#<closure ${stringify(code)}>` });
}

function _code(s: string): Value<any> {
    return closure({ rib: Primitives, next: null }, VM.parse(s));
}

Object.assign(Primitives, {
    'prepend': _code('to [vs n] n take vs ++'),
    'times': _code('dup 0 le [drop drop] [to [c n] c quote c n 1 - times] ifelse'),
    'iota': _code('to [n] 0 [dup 1 +] n times take'),
    'map': _code('to [vs f] 0 [dup vs ? f swap 1 + ] vs length times take'),
    'flatMap': _code('to [vs f] 0 take 0 [dup vs ? f rot swap ++ swap 1 +] vs length times drop'),
    'filter': _code('to [vs f] vs [dup f [1 take] [drop 0 take] ifelse] flatMap'),
});

// function R(s: string) {
//     console.log('\n' + s);
//     try {
//         const vm = new VM(s);
//         vm.exec();
//         console.log(s, '-->', vm.stack);
//     } catch (e) {
//         console.error(s, '-/->', (e as Error).message);
//     }
// }
//
// R('[dup 0 eq [drop] [dup 1 - x] ifelse] to x 5 x')
// R('[dup 0 eq [drop] [dup to v 1 - x v] ifelse] to x 5 x')
// R('[1 - dup 0 lt [drop] [dup x] ifelse] to x 5 x')
// R('[1 - dup 0 lt [drop] [dup to v x v] ifelse] to x 5 x')
// R('[3] 5 times');
// R('2 3 4 saveStack [2 *] map restoreStack');
// R('20 iota [dup 3 % 0 eq [1 take] [drop 0 take] ifelse] flatMap restoreStack');
// R('20 iota [3 % 0 eq] filter restoreStack');