import { EncodableDictionary, KeyedDictionary, Dictionary, Value, is, Record, Float, Bytes, isEmbedded, isSequence, Set, Atom, Embedded, merge as plainMerge, Preservable, PreserveWritable, _iterMap, stringify, fromJS } from '@preserves/core';
import { SchemaDefinition } from './reflection';
import * as M from './meta';
import * as H from './host';

export const UNIT: true = true;

export type Parsed<V> = Atom | V | Parsed<V>[] | DictOf<V> | Bindings<V>;
export type TopParsed<V> = Atom | V | Parsed<V>[] | DictOf<V> | TopBindings<V>;

export type Top<V> = Preservable<V> & PreserveWritable<V> & { __as_preserve__(): Value<V> };

export type DictOf<V> = EncodableDictionary<Parsed<V>, Parsed<V>, V>;

export type BindingName = string;
export type Bindings<V> = { [key: BindingName]: Parsed<V> };
export type TopBindings<V> = Bindings<V> & Top<V>;

export type SingleConstructor<V> = ((input: any) => Parsed<V>) & { schema(): SchemaDefinition };
export type MultipleConstructors<V> = { [key: string]: SingleConstructor<V> };
export type DefinitionConstructors<V> = SingleConstructor<V> | MultipleConstructors<V>;

export namespace Bindings {
    export function empty<V>(): Bindings<V> {
        return {};
    }
    export function single<V>(k: BindingName, v: Parsed<V>): Bindings<V> {
        const bs = empty<V>();
        bs[k] = v;
        return bs;
    }
    export function merge<V>(... vs: Bindings<V>[]): Bindings<V> {
        const acc = empty<V>();
        for (const v of vs) {
            Object.entries(v).forEach(([kw, vw]) => acc[kw] = vw);
        }
        return acc;
    }
}

export type DynField<V> =
    | { type: 'simple', value: Parsed<V> }
    | { type: 'compound', values: Bindings<V> }
;
export namespace DynField {
    export function unwrap<V>(f: DynField<V>): Parsed<V> {
        if (f.type === 'simple') return f.value;
        return f.values;
    }
    export function unwrap_compound<V>(f: DynField<V>): Bindings<V> {
        if (f.type === 'simple') throw new Error("Cannot unwrap DynField.simple to compound fields");
        return f.values;
    }
    export function simple<V>(value: Parsed<V>): DynField<V> {
        return { type: 'simple', value };
    }
    export function maybeSimple<V>(value: Parsed<V> | null): DynField<V> {
        return value === null ? compound(Bindings.empty()) : simple(value);
    }
    export function compound<V>(values: Bindings<V>): DynField<V> {
        return { type: 'compound', values };
    }
    export function promote<V>(f: DynField<V>, key?: symbol): Bindings<V> {
        if (f.type === 'compound') return f.values;
        return key ? Bindings.single(M.jsId(key.description!), f.value) : Bindings.empty();
    }
}

function optmap<A,B>(a: A | undefined, f: (a: A) => B): B | undefined {
    if (a === void 0) return void 0;
    return f(a);
}

export type Unparseable<V> = TopParsed<V>;
export type Unparser<V> = (v: Parsed<V>) => Value<V>;
export type UnparserCompound<V> = (v: Bindings<V>) => Value<V>;

export class SchemaInterpreter<V> {
    activeModule: M.ModulePath = [];
    unparserCache: { [key: string]: [Unparser<V>] } = {};

    constructor (
        public env: M.Modules<V> = new KeyedDictionary(),
        public mergeEmbeddeds: (a: V, b: V) => V | undefined = (_a, _b) => void 0,
    ) {}

    _withModule<R>(modulePath: M.ModulePath, f: () => R): R {
        const saved = this.activeModule;
        if (modulePath.length > 0) this.activeModule = modulePath;
        try {
            return f();
        } finally {
            if (modulePath.length > 0) this.activeModule = saved;
        }
    }

    _findModule(modulePath: M.ModulePath): { resolved: M.ModulePath, schema: M.Schema<V> } {
        const prefix = this.activeModule.slice();
        while (true) {
            const probe = [... prefix, ... modulePath];
            const schema = this.env.get(probe);
            if (schema !== void 0) {
                return { resolved: probe, schema };
            }
            if (prefix.length === 0) {
                throw new Error(`No such preserves-schema module: ${M.formatModulePath(modulePath)}, referred to in module ${M.formatModulePath(this.activeModule)}`);
            }
            prefix.pop();
        }
    }

    _lookup<R>(
        modulePath: M.ModulePath,
        name: symbol,
        f: (d: M.Definition<V>, schema: M.Schema<V>) => R,
    ): R {
        const { resolved, schema } = this._findModule(modulePath);
        return this._withModule(resolved, () => {
            const definition = schema.definitions.get(name);
            if (definition === void 0) {
                throw new Error(`No such preserves-schema definition: ${[... modulePath, name].map(s => s.description!).join('.')}`);
            }
            return f(definition, schema);
        });
    }

    makeTop(modulePath: M.ModulePath, name: symbol, fields: Bindings<V>): TopBindings<V> {
        const result = fields as TopBindings<V>;
        result.__as_preserve__ = () => this.unparser(modulePath, name)(result);
        result.__preserve_on__ = function (e) { e.push(this.__as_preserve__()); };
        result.__preserve_text_on__ = function (w) { w.push(this.__as_preserve__()); };
        return result;
    }

    definitionConstructor(
        modulePath: M.ModulePath,
        name: symbol,
    ): DefinitionConstructors<V> {
        return this._lookup(modulePath, name, (definition, schema): DefinitionConstructors<V> => {
            function attachSchema(
                f: (input: any) => Parsed<V>,
                variant?: symbol,
            ): SingleConstructor<V> {
                const g = f as SingleConstructor<V>;
                g.schema = () => ({
                    schema: fromJS(schema),
                    imports: {}, // TODO
                    definitionName: name,
                    variant,
                });
                return g;
            }
            const ty = H.definitionType(definition);
            if (ty._variant === 'union') {
                const multiple: MultipleConstructors<V> = {};
                ty.variants.forEach(v => {
                    const _variant = v.label.description!;
                    switch (v.type._variant) {
                        case 'Field':
                            multiple[_variant] = attachSchema(
                                (value: any) => this.makeTop(modulePath, name, { _variant, value }),
                                v.label);
                            break;
                        case 'Record':
                            multiple[_variant] = attachSchema(
                                (fields: object) => this.makeTop(modulePath, name, { _variant, ... fields }),
                                v.label);
                            break;
                    }
                });
                return multiple;
            } else {
                const flatName = M.formatModulePath([... modulePath, name]);
                switch (ty.value._variant) {
                    case 'Field': {
                        const tmp = { [flatName]: (value: any) => value };
                        return attachSchema(tmp[flatName]);
                    }
                    case 'Record': {
                        const rec = ty.value.value;
                        if (rec.fields.length > 1) {
                            const tmp = { [flatName]: (fields: Bindings<V>) =>
                                this.makeTop(modulePath, name, fields) };
                            return attachSchema(tmp[flatName]);
                        } else {
                            const tmp = { [flatName]: (field: Parsed<V>) =>
                                this.makeTop(modulePath, name, {
                                    [M.jsId(rec.fields[0].name.description!)]: field,
                                }) };
                            return attachSchema(tmp[flatName]);
                        }
                    }
                }
            }
        });
    }

    parse(
        modulePath: M.ModulePath,
        name: symbol,
        input: Value<V>,
    ): Unparseable<V> {
        const v = this.tryParse(modulePath, name, input);
        if (v === void 0) {
            throw new TypeError(
                `Invalid ${M.formatModulePath([... modulePath, name])}: ${stringify(input)}`)
        }
        return v;
    }

    tryParse(
        modulePath: M.ModulePath,
        name: symbol,
        input: Value<V>,
    ): Unparseable<V> | undefined {
        return this._lookup(modulePath, name, definition =>
            optmap(this.parseDefinition(definition, input), result0 => {
                const ty = H.definitionType(definition);
                if (ty._variant === 'union' || ty.value._variant === 'Record') {
                    return this.makeTop(modulePath, name, result0 as Bindings<V>);
                } else {
                    return result0 as Exclude<Parsed<V>, Bindings<V>>;
                }
            }));
    }

    parseDefinition(d: M.Definition<V>, input: Value<V>): Parsed<V> | undefined {
        switch (d._variant) {
            case 'or':
                return this.parseNamedAlternative(d.pattern0, input) ??
                    this.parseNamedAlternative(d.pattern1, input) ??
                    (() => {
                        for (const p of d.patternN) {
                            const r = this.parseNamedAlternative(p, input);
                            if (r !== void 0) return r;
                        }
                        return void 0;
                    })();
            case 'and': {
                const rs = [this.parseNamedPattern(d.pattern0, input),
                            this.parseNamedPattern(d.pattern1, input),
                            ... d.patternN.map(p => this.parseNamedPattern(p, input))];
                for (const r of rs) {
                    if (r === void 0) return void 0;
                }
                return Bindings.merge(... rs as Bindings<V>[]);
            }
            case 'Pattern':
                return optmap(this.parsePattern(d.value, input), DynField.unwrap);
        }
    }

    parseNamedAlternative(p: M.NamedAlternative<V>, input: Value<V>): Bindings<V> | undefined {
        return optmap(this.parsePattern(p.pattern, input), w => {
            const result = DynField.promote(w, Symbol.for('value'));
            result._variant = p.variantLabel;
            return result;
        });
    }

    parseNamedPattern(p: M.NamedPattern<V>, input: Value<V>): Bindings<V> | undefined {
        switch (p._variant) {
            case 'named':
                return optmap(this.parseSimplePattern(p.value.pattern, input),
                              w => DynField.promote(DynField.maybeSimple(w), p.value.name));
            case 'anonymous':
                return optmap(this.parsePattern(p.value, input),
                              w => DynField.promote(w));
        }
    }

    parseNamedSimplePattern(p: M.NamedSimplePattern<V>, input: Value<V>): DynField<V> | undefined {
        switch (p._variant) {
            case 'named':
                return optmap(this.parseSimplePattern(p.value.pattern, input),
                              w => DynField.compound(DynField.promote(DynField.maybeSimple(w), p.value.name)));
            case 'anonymous':
                return optmap(this.parseSimplePattern(p.value, input), DynField.maybeSimple<V>);
        }
    }

    parseSimplePattern(p: M.SimplePattern<V>, input: Value<V>): Parsed<V> | null | undefined {
        const inputIf = (b: boolean) => b ? input : void 0;
        switch (p._variant) {
            case 'any': return input;
            case 'atom': switch (p.atomKind._variant) {
                case 'Boolean': return inputIf(typeof input === 'boolean');
                case 'Double': return inputIf(Float.isDouble(input));
                case 'SignedInteger': return inputIf(typeof input === 'number' || typeof input === 'bigint');
                case 'String': return inputIf(typeof input === 'string');
                case 'ByteString': return inputIf(Bytes.isBytes(input));
                case 'Symbol': return inputIf(typeof input === 'symbol');
            }
            case 'embedded': return isEmbedded(input) ? input.embeddedValue : void 0;
            case 'lit': return is(input, p.value) ? null : void 0;
            case 'seqof': {
                if (!isSequence(input)) return void 0;
                const result: Parsed<V>[] = [];
                for (const v of input) {
                    const w = this.parseSimplePattern(p.pattern, v);
                    if (w === void 0) return void 0;
                    if (w !== null) result.push(w);
                }
                return result;
            }
            case 'setof': {
                if (!Set.isSet<V>(input)) return void 0;
                const result: Parsed<V>[] = [];
                for (const v of input) {
                    const w = this.parseSimplePattern(p.pattern, v);
                    if (w === void 0) return void 0;
                    if (w !== null) result.push(w);
                }
                return result;
            }
            case 'dictof': {
                if (!Dictionary.isDictionary<V>(input)) return void 0;
                const result: DictOf<V> = new EncodableDictionary(
                    this.unparserSimplePattern(p.key),
                    this.unparserSimplePattern(p.value));
                for (const [k, v] of input) {
                    const kw = this.parseSimplePattern(p.key, k);
                    if (kw === void 0) return void 0;
                    const vw = this.parseSimplePattern(p.value, v);
                    if (vw === void 0) return void 0;
                    result.set(kw === null ? UNIT : kw, vw === null ? UNIT : vw);
                }
                return result;
            }
            case 'Ref': return this.tryParse(p.value.module, p.value.name, input);
        }
    }

    parseCompoundPattern(p: M.CompoundPattern<V>, input: Value<V>): Bindings<V> | undefined {
        switch (p._variant) {
            case 'rec':
                if (!Record.isRecord<Value<V>, Array<Value<V>>, V>(input)) return void 0;
                return optmap(this.parseNamedPattern(p.label, input.label),
                              lw => optmap(this.parseNamedPattern(p.fields, Array.from(input)),
                                           fsw => Bindings.merge(lw, fsw)));
            case 'tuple': {
                if (!isSequence(input)) return void 0;
                if (input.length !== p.patterns.length) return void 0;
                let results: Bindings<V>[] = [];
                for (let i = 0; i < p.patterns.length; i++) {
                    const w = this.parseNamedPattern(p.patterns[i], input[i]);
                    if (w === void 0) return void 0;
                    results.push(w);
                }
                return Bindings.merge(... results);
            }
            case 'tuplePrefix': {
                if (!isSequence(input)) return void 0;
                if (input.length < p.fixed.length) return void 0;
                let fixed_results: Bindings<V>[] = [];
                for (let i = 0; i < p.fixed.length; i++) {
                    const w = this.parseNamedPattern(p.fixed[i], input[i]);
                    if (w === void 0) return void 0;
                    fixed_results.push(w);
                }
                const remainder = input.slice(p.fixed.length);
                return optmap(this.parseNamedSimplePattern(p.variable, remainder), vw => {
                    const variable_results = DynField.unwrap_compound(vw);
                    return Bindings.merge(variable_results, ... fixed_results);
                });
            }
            case 'dict': {
                if (!Dictionary.isDictionary<V>(input)) return void 0;
                const results: Bindings<V>[] = [];
                for (const [key, vp] of p.entries) {
                    const v = input.get(key);
                    if (v === void 0) return void 0;
                    const vw = this.parseNamedSimplePattern(vp, v);
                    if (vw === void 0) return void 0;
                    results.push(DynField.unwrap_compound(vw));
                }
                return Bindings.merge(... results);
            }
        }
    }

    parsePattern(p: M.Pattern<V>, input: Value<V>): DynField<V> | undefined {
        switch (p._variant) {
            case 'SimplePattern':
                return optmap(this.parseSimplePattern(p.value, input), DynField.maybeSimple);
            case 'CompoundPattern':
                return optmap(this.parseCompoundPattern(p.value, input), DynField.compound);
        }
    }

    unparse(
        modulePath: M.ModulePath,
        name: symbol,
        v: Unparseable<V>,
    ): Value<V> {
        return this.unparser(modulePath, name)(v);
    }

    unparser(modulePath: M.ModulePath, name: symbol): Unparser<V> {
        return this._unparser(modulePath, name)[0];
    }

    _unparser(modulePath: M.ModulePath, name: symbol): [Unparser<V>] {
        const key = [... modulePath.map(n => n.description!), name.description!].join('.');
        if (!(key in this.unparserCache)) {
            const cell: [Unparser<V>] = [null!];
            this.unparserCache[key] = cell;
            cell[0] = this._lookup(modulePath, name, p => this.unparserDefinition(p));
        }
        return this.unparserCache[key];
    }

    unparserDefinition(p: M.Definition<V>): Unparser<V> {
        switch (p._variant) {
            case 'or': {
                const ups = [p.pattern0, p.pattern1, ... p.patternN].map(
                    p => this.unparserNamedAlternative(p));
                return v => {
                    const bs = v as Bindings<V>;
                    return ups.find(up => up[0] === bs._variant)![1](bs);
                };
            }
            case 'and': {
                const ups = [p.pattern0, p.pattern1, ... p.patternN].map(
                    p => this.unparserNamedPattern(p));
                return v => plainMerge(this.mergeEmbeddeds,
                                       ups[0](v), ... ups.slice(1).map(up => up(v)));
            }
            case 'Pattern':
                return this.unparserPattern(p.value);
        }
    }

    unparserNamedAlternative(p: M.NamedAlternative<V>): [string, UnparserCompound<V>] {
        const up = this.unparserPattern(p.pattern);
        const ty = H.patternType(p.pattern);
        switch (ty._variant) {
            case 'Field': return [p.variantLabel, bs => up(bs['value'])];
            case 'Record': return [p.variantLabel, up];
        }
    }

    unparserNamedPattern(p: M.NamedPattern<V>): Unparser<V> {
        switch (p._variant) {
            case 'named': {
                const up = this.unparserSimplePattern(p.value.pattern);
                const key = M.jsId(p.value.name.description!);
                return v => up((v as Bindings<V>)[key]);
            }
            case 'anonymous':
                return this.unparserPattern(p.value);
        }
    }

    unparserPattern(p: M.Pattern<V>): Unparser<V> {
        switch (p._variant) {
            case 'CompoundPattern': {
                const up = this.unparserCompoundPattern(p.value);
                return v => up(v as Bindings<V>);
            }
            case 'SimplePattern':
                return this.unparserSimplePattern(p.value);
        }
    }

    unparserSimplePattern(p: M.SimplePattern<V>): Unparser<V> {
        switch (p._variant) {
            case 'any': return v => v as Value<V>; // ?!
            case 'atom': return v => v as Atom;
            case 'embedded': return v => new Embedded(v as V);
            case 'lit': return _v => p.value;
            case 'seqof': {
                const up = this.unparserSimplePattern(p.pattern);
                return vs => (vs as Parsed<V>[]).map(up);
            }
            case 'setof': {
                const up = this.unparserSimplePattern(p.pattern);
                return vs => new Set<V>((vs as Parsed<V>[]).map(up));
            }
            case 'dictof': {
                const kp = this.unparserSimplePattern(p.key);
                const vp = this.unparserSimplePattern(p.value);
                return vs => new Dictionary<V>(_iterMap(
                    (vs as DictOf<V>).entries(), ([k, v]) => [kp(k), vp(v)]));
            }
            case 'Ref': {
                const up = this._unparser(p.value.module, p.value.name);
                return v => up[0](v as Bindings<V>);
            }
        }
    }

    unparserCompoundPattern(p: M.CompoundPattern<V>): UnparserCompound<V> {
        switch (p._variant) {
            case 'rec': {
                const lp = this.unparserNamedPattern(p.label);
                const fp = this.unparserNamedPattern(p.fields);
                return bs => Record(lp(bs), fp(bs) as Value<V>[]);
            }
            case 'tuple': {
                const ups = p.patterns.map(p => this.unparserNamedPattern(p));
                return bs => ups.map(up => up(bs));
            }
            case 'tuplePrefix': {
                const fixed = p.fixed.map(p => this.unparserNamedPattern(p));
                const variable = this.unparserNamedSimplePattern(p.variable);
                return bs => [... fixed.map(up => up(bs)), ... variable(bs) as Value<V>[]];
            }
            case 'dict': {
                const ups: [Value<V>, Unparser<V>][] = Array.from(p.entries.entries()).map(
                    ([key, vp]) => [key, this.unparserNamedSimplePattern(vp)]);
                return bs => {
                    const result = new Dictionary<V>();
                    for (const [key, up] of ups) {
                        result.set(key, up(bs));
                    }
                    return result;
                };
            }
        }
    }

    unparserNamedSimplePattern(p: M.NamedSimplePattern<V>): Unparser<V> {
        switch (p._variant) {
            case 'named': {
                const up = this.unparserSimplePattern(p.value.pattern);
                const key = M.jsId(p.value.name.description!);
                return v => up((v as Bindings<V>)[key]);
            }
            case 'anonymous':
                return this.unparserSimplePattern(p.value);
        }
    }
}