import { Dictionary, KeyedSet, FlexSet, Position, stringify } from "@preserves/core";
import { refPosition } from "../reader";
import * as M from "../meta";
import { anglebrackets, block, braces, commas, formatItems, Item, keyvalue, seq, opseq } from "./block";
import { ANY_TYPE, RefType, Type } from "./type";
import { renderType, variantInitFor } from "./rendertype";
import { typeForDefinition } from "./gentype";
import { SchemaSyntaxError } from "../error";

export interface CompilerOptions {
    preservesModule?: string;
    defaultEmbeddedType?: M.Ref;
    warn?(message: string, pos: Position | null): void;
}

export interface Capture {
    fieldName: string;
    sourceExpr: string;
}

export const RECURSION_LIMIT = 128;

export class ModuleContext {
    readonly env: M.Environment;
    readonly modulePath: M.ModulePath;
    readonly schema: M.Schema;
    readonly options: CompilerOptions;
    readonly embeddedType: Item;

    readonly literals = new Dictionary<M.InputEmbedded, string>();
    readonly preamble: Item[] = [];
    readonly typedefs: Item[] = [];
    readonly functiondefs: Item[] = [];
    readonly imports = new KeyedSet<[M.ModulePath, string, string, string]>();

    constructor(
        env: M.Environment,
        modulePath: M.ModulePath,
        schema: M.Schema,
        options: CompilerOptions,
    ) {
        this.env = env;
        this.modulePath = modulePath;
        this.schema = schema;
        this.options = options;
        switch (schema.embeddedType._variant) {
            case 'false':
                this.embeddedType = '_.GenericEmbedded';
                break;
            case 'Ref': {
                const t = this.resolver()(schema.embeddedType.value);
                this.embeddedType = t.typeName;
                break;
            }
        }
    }

    literal(v: M.Input): Item {
        let varname = this.literals.get(v);
        if (varname === void 0) {
            varname = M.jsId('$' + stringify(v), () => '__lit' + this.literals.size);
            this.literals.set(v, varname);
        }
        return varname;
    }

    derefPattern(p: M.Definition, refCount = 0): M.Definition {
        if (refCount > RECURSION_LIMIT) {
            throw new Error('Recursion limit exceeded');
        }
        if (p._variant === 'Pattern' &&
            p.value._variant === 'SimplePattern' &&
            p.value.value._variant === 'Ref')
        {
            return this.lookup(p.value.value.value,
                               (p, _t) => this.derefPattern(p, refCount + 1),
                               ((_modPath, _modId, _modFile, _modExpr, pp, _tt) =>
                                   this.derefPattern(pp ?? p, refCount + 1)));
        } else {
            return p;
        }
    }

    definePreamble(i: Item): void {
        this.preamble.push(i);
    }

    defineType(f: Item): void {
        this.typedefs.push(f);
    }

    defineFunctions(definitionName: string, f: (ctx: FunctionContext) => Item[]): void {
        this.functiondefs.push(... f(new FunctionContext(this, definitionName)));
    }

    resolver(modulePath?: M.ModulePath): (ref: M.Ref) => RefType {
        return (ref) => this.lookup(
            ref,
            (_p, _t) => Type.ref(ref.name.description!, ref),
            (modPath, modId, modFile, modExpr, _p, t) => {
                this.imports.add([modPath, modId, modFile, modExpr]);
                return Type.ref(`${modId}${modExpr}.${ref.name.description!}`, ref);
            },
            modulePath);
    }

    lookupType(name: M.Ref, modulePath?: M.ModulePath): Type | null {
        const t = this.lookup(
            name,
            (_p, t) => t,
            (_modPath, _modId, _modFile, _modExpr, _p, t) => t,
            modulePath);
        return t ? t() : null;
    }

    lookup<R>(name: M.Ref,
              kLocal: (p: M.Definition, t: () => Type) => R,
              kOther: (modPath: M.ModulePath,
                       modId: string,
                       modFile: string,
                       modExpr: string,
                       p: M.Definition | null,
                       t: (() => Type) | null) => R,
              modulePath?: M.ModulePath)
    : R {
        const soughtModule = name.module.length ? name.module : (modulePath ?? this.modulePath);

        const e = M.envLookup(this.env, soughtModule);
        if (e !== null) {
            const expr = (e.typescriptModuleExpr === null) ? '' : '.' + e.typescriptModuleExpr;
            if (e.schema === null) {
                // It's an artificial module, not from a schema. Assume the identifier is present.
                return kOther(soughtModule,
                              M.modsymFor(e),
                              e.typescriptModulePath,
                              expr,
                              null,
                              null);
            } else {
                const p = e.schema.definitions.get(name.name);
                if (p !== void 0) {
                    let t = () => typeForDefinition(this.resolver(soughtModule), p);
                    if (name.module.length) {
                        return kOther(soughtModule,
                                      M.modsymFor(e),
                                      e.typescriptModulePath,
                                      expr,
                                      p,
                                      t);
                    } else {
                        return kLocal(p, t);
                    }
                }
            }
        }

        throw new SchemaSyntaxError(`Undefined reference: ${M.formatRef(name)}`, refPosition(name));
    }

    genericParameters(): Item {
        return anglebrackets(seq('_embedded = ', this.embeddedType));
    }

    genericParametersFor(t: Type): Item {
        return this.hasEmbedded(t) ? this.genericParameters() : '';
    }

    genericArgs(): Item {
        return `<_embedded>`;
    }

    genericArgsFor(t: Type): Item {
        return this.hasEmbedded(t) ? this.genericArgs() : '';
    }

    hasEmbedded(t: Type): boolean {
        const self = this;
        const state = new WalkState(this.modulePath);

        function walk(t: Type): boolean {
            switch (t.kind) {
                case 'union':
                    for (const v of t.variants.values()) { if (walk(v)) return true; };
                    return false;
                case 'unit': return false;
                case 'array': return walk(t.type);
                case 'set': return true; // because ref to _embedded in renderType()
                case 'dictionary': return true; // because ref to _embedded in renderType()
                case 'ref': {
                    if (t.ref === null) {
                        switch (t.typeName) {
                            case '_embedded': return true;
                            case '_.Value': return true;
                            default: return false;
                        }
                    } else {
                        return state.cycleCheck(
                            t.ref,
                            ref => self.lookupType(ref, state.modulePath),
                            t => t ? walk(t) : false,
                            () => false);
                    }
                }
                case 'record':
                    for (const v of t.fields.values()) { if (walk(v)) return true; };
                    return false;
            }
        }

        return walk(t);
    }

    withAsPreserveMixinType(name: string, t: Type): Item {
        if (t.kind === 'unit' || t.kind === 'record' || t.kind === 'union') {
            return opseq('any', ' & ',
                         seq(name, this.genericArgsFor(t)),
                         braces(seq('__as_preserve__',
                                    this.hasEmbedded(t) ? '' : this.genericParameters(),
                                    '()',
                                    ': _.Value', this.genericArgs())));
        } else {
            return seq(name, this.genericArgsFor(t));
        }
    }
}

export class FunctionContext {
    readonly mod: ModuleContext;
    readonly definitionName: string;

    tempCounter = 0;
    temps: Map<string, { type: Item, names: string[] }> = new Map();

    captures: Capture[] = [];
    variantName: string | undefined = void 0;

    constructor(mod: ModuleContext, definitionName: string) {
        this.mod = mod;
        this.definitionName = definitionName;
    }

    gentempname(): string {
        return '_tmp' + this.tempCounter++;
    }

    gentemp(vartype: Type = ANY_TYPE): string {
        const typeitem = renderType(this.mod, vartype);
        const typestr = formatItems([typeitem], Infinity);
        const varname = this.gentempname();
        let e = this.temps.get(typestr);
        if (e === void 0) {
            e = { type: typeitem, names: [] };
            this.temps.set(typestr, e);
        }
        e.names.push(varname);
        return varname;
    }

    block(f: () => Item[]): Item {
        const oldTemps = this.temps;
        this.temps = new Map();
        const items = f();
        const ts = this.temps;
        this.temps = oldTemps;
        return block(
            ... Array.from(ts).map(([_typestr, { type, names }]) =>
                seq(`let `, commas(... names), `: (`, type, `) | undefined`)),
            ... items);
    }

    withCapture<R>(
        fieldName: string | undefined, sourceExpr: string, ks: (sourceExpr: string) => R): R
    {
        if (fieldName !== void 0) this.captures.push({ fieldName, sourceExpr });
        const result = ks(sourceExpr);
        if (fieldName !== void 0) this.captures.pop();
        return result;
    }

    convertCapture(
        fieldName: string | undefined, sourceExpr: string, ks: () => Item[]): Item
    {
        return this.withCapture(fieldName, sourceExpr, sourceExpr =>
            seq(`if (${sourceExpr} !== void 0) `, this.block(() => ks())));
    }

    buildCapturedCompound(dest: string): Item {
        const fields = [
            ... variantInitFor(this.variantName),
            ... this.captures.map(({ fieldName, sourceExpr }) =>
                keyvalue(fieldName, sourceExpr)),
            seq(`__as_preserve__() `, block(`return from${this.definitionName}(this)`))
        ];
        return seq(`${dest} = `, braces(... fields));
    }
}

export class WalkState {
    modulePath: M.ModulePath;
    readonly seen: FlexSet<M.Ref>;

    constructor(modulePath: M.ModulePath) {
        this.modulePath = modulePath;
        this.seen = new FlexSet(refCanonicalizer);
    }

    cycleCheck<E, R>(
        r0: M.Ref,
        step: (ref: M.Ref) => E,
        ks: (e: E) => R,
        kf: () => R,
    ): R {
        const r = M.Ref({
            module: r0.module.length ? r0.module : this.modulePath,
            name: r0.name
        });
        if (this.seen.has(r)) {
            return kf();
        } else {
            this.seen.add(r);
            const maybe_e = step(r);
            const saved = this.modulePath;
            this.modulePath = r.module;
            const result = ks(maybe_e);
            this.modulePath = saved;
            return result;
        }
    }
}

function refCanonicalizer(r: M.Ref): string {
    return stringify([... r.module, r.name]);
}