import { Def, Expr, Stmt } from "./ast.ts";

function appended<T>(vs: T[], v: T): T[] {
    vs.push(v);
    return vs;
}

export type Tok = {
    ty: string;
    text: string;
    line: number;
};

export function tokenize(text: string): Tok[] {
    return text
        .replace(/\/\/[^\n]*/g, "")
        .replace(/[\(\)\n]/g, " $& ")
        .split(/[ \t\r]+/)
        .filter((tok) => tok !== "")
        .reduce<[{ tok: string; line: number }[], number]>(
            ([toks, line], tok) =>
                tok != "\n"
                    ? [appended(toks, { tok, line }), line]
                    : [toks, line + 1],
            [[], 1],
        )[0]
        .map(({ tok, line }) => ({ ty: tok, text: tok, line }))
        .map((tok) => {
            if (/^[a-zA-Z_0-9]+$/.test(tok.text)) {
                tok.ty = "ident";
            }
            return tok;
        });
}

export type SExpr = {
    line: number;
    ty: "Ident" | "List";
    ident?: string;
    exprs?: SExpr[];
};

function parseSExpr(toks: Tok[]): [Tok[], SExpr | null] {
    if (toks.length == 0) {
        return [toks, null];
    }
    const line = toks[0].line;
    if (toks[0].ty === "ident") {
        const ident = toks[0].text;
        return [toks.slice(1), { line, ty: "Ident", ident }];
    } else if (toks[0].ty === "(") {
        toks = toks.slice(1);
        const exprs: SExpr[] = [];
        while (toks.length != 0 && toks[0].ty !== ")") {
            let expr: SExpr | null;
            [toks, expr] = parseSExpr(toks);
            if (!expr) {
                throw new Error(`expected expression on line ${line}`);
            }
            exprs.push(expr);
        }
        if (toks.length == 0 || toks[0].ty != ")") {
            throw new Error(
                `expected ')' on line ${
                    toks.at(0)?.line ?? "(eof)"
                }, '(' on line ${line}`,
            );
        }
        return [toks.slice(1), { line, ty: "List", exprs }];
    } else {
        throw new Error(`malformed on line ${line}`);
    }
}

export function parseSExprs(toks: Tok[]): SExpr[] {
    const exprs: SExpr[] = [];
    while (toks.length != 0) {
        let expr: SExpr | null;
        [toks, expr] = parseSExpr(toks);
        exprs.push(expr!);
    }
    return exprs;
}

class SExprMatcher {
    private failed = false;
    private captures = new Map<string, SExpr>();

    constructor(private s: SExpr) {}

    match(): Record<string, SExpr> | null {
        return this.failed ? null : Object.fromEntries(this.captures.entries());
    }

    ident(val?: string): this {
        if (this.s.ty !== "Ident" || val && this.s.ident! !== val) {
            this.failed = true;
            return this;
        }
        return this;
    }

    list(length?: number): this {
        if (
            this.s.ty !== "List" || (length && this.s.exprs!.length !== length)
        ) {
            this.failed = true;
            return this;
        }
        return this;
    }

    at(idx: number, func?: (s: SExprMatcher) => SExprMatcher): this {
        if (this.s.ty !== "List" || idx >= this.s.exprs!.length) {
            this.failed = true;
            return this;
        }
        if (func) {
            const inner = new SExprMatcher(this.s.exprs![idx]);
            func(inner);
            for (const [key, val] of inner.captures.entries()) {
                this.captures.set(key, val);
            }
            this.failed = this.failed || inner.failed;
        }
        return this;
    }

    all(func: (s: SExprMatcher) => SExprMatcher): this {
        if (this.s.ty !== "List") {
            this.failed = true;
            return this;
        }
        for (const s1 of this.s.exprs!) {
            const inner = new SExprMatcher(s1);
            func(inner);
            for (const [key, val] of inner.captures.entries()) {
                this.captures.set(key, val);
            }
        }
        return this;
    }

    capture(id: string): this {
        this.captures.set(id, this.s);
        return this;
    }
}

export function parseAst(ss: SExpr[]): Def[] {
    const defs: Def[] = [];
    for (const s of ss) {
        defs.push(parseDef(s));
    }
    return defs;
}

function parseDef(s: SExpr): Def {
    const m = new SExprMatcher(s)
        .list(5)
        .at(0, (s) => s.ident("def"))
        .at(1, (s) => s.ident().capture("ident"))
        .at(2, (s) => s.all((s) => s.ident()).capture("inputs"))
        .at(3, (s) => s.all((s) => s.ident()).capture("outputs"))
        .at(4, (s) => s.list().capture("stmts"))
        .match();
    if (!m) {
        throw new Error(`malformed def on line ${s.line}`);
    }
    return new Def(
        s.line,
        m.ident.ident!,
        m.inputs.exprs!.map((s) => s.ident!),
        m.outputs.exprs!.map((s) => s.ident!),
        m.stmts.exprs!.map((s) => parseStmt(s)),
    );
}

function parseStmt(s: SExpr): Stmt {
    const set_match = new SExprMatcher(s)
        .list(3)
        .at(0, (s) => s.ident("set"))
        .at(1, (s) => s.ident().capture("subject"))
        .at(2, (s) => s.capture("expr"))
        .match();
    if (set_match) {
        const m = set_match;
        return new Stmt(
            s.line,
            {
                tag: "Set",
                subject: parseExpr(m.subject),
                expr: parseExpr(m.expr),
            },
        );
    }
    const let_match = new SExprMatcher(s)
        .list(3)
        .at(0, (s) => s.ident("let"))
        .at(1, (s) => s.all((s) => s.ident()).capture("idents"))
        .at(2, (s) => s.capture("expr"))
        .match();
    if (let_match) {
        const m = let_match;

        return new Stmt(
            s.line,
            {
                tag: "Let",
                idents: m.idents.exprs!.map((s) => s.ident!),
                expr: parseExpr(m.expr),
            },
        );
    }
    throw new Error(`malformed statement on line ${s.line}`);
}

function parseExpr(s: SExpr): Expr {
    if (s.ty === "Ident") {
        return new Expr(s.line, { tag: "Ident", ident: s.ident! });
    } else if (s.ty === "List") {
        if (s.exprs!.length === 0) {
            throw new Error(`empty expression not allowed, on line ${s.line}`);
        }
        if (s.exprs![0].ty !== "Ident") {
            throw new Error(`callee must be an identifier, on line ${s.line}`);
        }
        return new Expr(
            s.line,
            {
                tag: "Call",
                callee: parseExpr(s.exprs![0]),
                args: s.exprs!.slice(1).map((s) => parseExpr(s)),
            },
        );
    } else {
        throw new Error(`expected expression, on line ${s.line}`);
    }
}

export function resolve(defs: Def[]): Map<Def, Map<Expr, Sym>> {
    const defSyms = new Map<string, Def>();
    for (const def of defs) {
        defSyms.set(def.ident, def);
    }
    const resols = new Map<Def, Map<Expr, Sym>>();
    for (const def of defs) {
        const defResolver = new DefResolver(def, defSyms);
        defResolver.resolve();
        resols.set(def, defResolver.resols);
    }
    return resols;
}

export type Sym =
    | { tag: "Builtin" }
    | { tag: "Input"; idx: number }
    | { tag: "Output"; idx: number }
    | { tag: "Node"; stmt: Stmt; idx: number }
    | { tag: "Def"; def: Def };

class DefResolver {
    private syms = new Map<string, Sym>([
        ["relay_default_off", { tag: "Builtin" }],
        ["relay_default_on", { tag: "Builtin" }],
        ["gnd", { tag: "Builtin" }],
        ["vin", { tag: "Builtin" }],
    ]);
    public resols = new Map<Expr, Sym>();

    constructor(private def: Def, private defSyms: Map<string, Def>) {}

    resolve() {
        for (const [idx, ident] of this.def.inputs.entries()) {
            this.syms.set(ident, { tag: "Input", idx });
        }
        for (const [idx, ident] of this.def.outputs.entries()) {
            this.syms.set(ident, { tag: "Output", idx });
        }
        const { syms, defSyms, resols } = this;
        this.def.visit({
            visitStmt(stmt) {
                stmt.visitSubtree(this);
                if (stmt.kind.tag === "Let") {
                    for (const [idx, ident] of stmt.kind.idents.entries()) {
                        syms.set(ident, { tag: "Node", stmt, idx });
                    }
                }
                return "stop";
            },
            visitExpr(expr) {
                if (expr.kind.tag === "Ident") {
                    const sym = syms.get(expr.kind.ident);
                    if (sym) {
                        resols.set(expr, sym);
                        return;
                    }
                    const defSym = defSyms.get(expr.kind.ident);
                    if (defSym) {
                        resols.set(expr, { tag: "Def", def: defSym });
                        return;
                    }
                    throw new Error(
                        `unresolved identifier '${expr.kind.ident}', on line ${expr.line}`,
                    );
                }
            },
        });
    }
}