import * as parser from "./parser.out.js";
import * as yaml from "jsr:@std/yaml";
import * as ast from "./ast.ts";

async function main() {
    const filename = Deno.args[0];
    const text = await Deno.readTextFile(filename);

    const ast = ((): ast.Struct[] | undefined => {
        try {
            return parser.parse(text, {
                grammarSource: filename,
            });
        } catch (e) {
            if (e instanceof parser.SyntaxError) {
                console.error(
                    e.format([{ source: filename, text }]),
                );
                return undefined;
            }
            throw e;
        }
    })();
    if (!ast) {
        return;
    }
    //console.log(yaml.stringify(ast));

    const rep = new Reporter(text);

    const re = new Resolver(ast, rep).resolve();
    const units = new DependencyTracer(ast, re, rep).trace();
    const unitMap = mapUnitIds(units);

    new CircularDependencyFinder(unitMap, rep)
        .findCircularDependencies();

    const entries = entryPoints(units, rep);

    if (rep.errorOccured) {
        console.log("Errors occured. Stopping...");
    }
}

type MonoUnit = {
    id: number;
    unit: Unit;
    generics?: Map<number, ast.Ty>;
};

class Monomorphizer {
    public constructor(
        private units: Unit[],
    ) {}

    public monomorph(): MonoUnit[] {
        const monoUnits: MonoUnit[] = [];
        for (const unit of this.units) {
        }
        return monoUnits;
    }
}

class CircularDependencyFinder {
    private hasCirc = new Set<number>();

    public constructor(
        private units: Map<number, Unit>,
        private rep: Reporter,
    ) {}

    public findCircularDependencies() {
        for (const id of this.units.keys()) {
            this.searchDependencies(id);
        }
    }

    private searchDependencies(
        id: number,
        defined = new Set<number>(),
    ) {
        if (this.hasCirc.has(id)) {
            return;
        }
        defined.add(id);
        const unit = this.units.get(id)!;
        for (const depId of unit.dependencies) {
            const dep = this.units.get(depId)!;
            if (defined.has(depId)) {
                this.rep.error(
                    `circular dependency between '${unit.struct.ident}' and '${dep.struct.ident}'`,
                );
                this.hasCirc.add(unit.id);
                return;
            }
            this.searchDependencies(depId, new Set(defined));
        }
    }
}

function entryPoints(units: Unit[], rep: Reporter): Set<number> {
    const entries = units.filter((unit) => unit.dependencies.size === 0);
    if (entries.length === 0) {
        rep.error(
            "no entry points in dependency graph (everything depends on something else)",
        );
    }
    return new Set(entries.map((entry) => entry.struct.id));
}

function mapUnitIds(units: Unit[]): Map<number, Unit> {
    return new Map(units.map((v, i) => [i, v]));
}

type Unit = {
    id: number;
    struct: ast.Struct;
    dependencies: Set<number>;
};

class DependencyTracer {
    public constructor(
        private ast: ast.Struct[],
        private re: Map<number, Resol>,
        private rep: Reporter,
    ) {}

    public trace(): Unit[] {
        const units: Unit[] = [];
        for (const struct of this.ast) {
            const dependencies = new StructTracer(struct, this.re).trace();
            units.push({ id: struct.id, struct, dependencies });
        }
        return units;
    }
}

class StructTracer {
    private dependencies = new Set<number>();

    public constructor(
        private struct: ast.Struct,
        private re: Map<number, Resol>,
    ) {}

    public trace(): Set<number> {
        for (const field of this.struct.fields) {
            this.traceTy(field.ty);
        }
        return this.dependencies;
    }

    private traceTy(ty: ast.Ty) {
        const k = ty.kind;
        switch (k.tag) {
            case "array":
                this.traceTy(k.ty);
                return;
            case "generic": {
                const re = this.re.get(ty.id)!;
                if (re.tag !== "struct") {
                    return;
                }
                this.dependencies.add(re.struct.id);
                for (const ty of k.args) {
                    this.traceTy(ty);
                }
                return;
            }
            case "struct_literal":
                for (const field of k.fields) {
                    this.traceTy(field.ty);
                }
                return;
            case "ident": {
                const re = this.re.get(ty.id)!;
                if (re.tag !== "struct") {
                    return;
                }
                this.dependencies.add(re.struct.id);
                return;
            }
            case "int":
                return;
        }
        const _: never = k;
    }
}

type Resol =
    | { tag: "error" }
    | { tag: "undefined" }
    | { tag: "struct"; struct: ast.Struct }
    | { tag: "generic"; struct: ast.Struct; idx: number }
    | { tag: "primitive" };

type Rib = {
    defs: Map<string, Resol>;
};

class Resolver {
    private resols = new Map<number, Resol>();
    private ribs: Rib[] = [{ defs: new Map() }];

    public constructor(
        private ast: ast.Struct[],
        private rep: Reporter,
    ) {}

    public resolve(): Map<number, Resol> {
        for (const struct of this.ast) {
            const res = this.define(struct.ident, { tag: "struct", struct });
            if (res === "already defined") {
                this.reportAlreadyDefined(struct.ident, struct.loc);
            }
        }
        for (const struct of this.ast) {
            this.ribs.push({ defs: new Map() });
            if (struct.generics) {
                for (const [idx, ident] of struct.generics.params.entries()) {
                    const res = this.define(ident, {
                        tag: "generic",
                        struct,
                        idx,
                    });
                    if (res === "already defined") {
                        this.reportAlreadyDefined(ident, struct.loc);
                    }
                }
            }
            const fields = new Map<string, ast.Field>();
            for (const field of struct.fields) {
                this.resolveTy(field.ty);
                if (fields.has(field.ident)) {
                    this.rep.error(
                        `field '${field.ident}' already defined`,
                        field.loc,
                    );
                    const otherField = fields.get(field.ident)!;
                    this.rep.info(
                        `original field '${otherField.ident}' defined here`,
                        otherField.loc,
                    );
                    continue;
                }
                fields.set(field.ident, field);
            }
            this.ribs.pop();
        }
        return this.resols;
    }

    private resolveTy(ty: ast.Ty) {
        const k = ty.kind;
        switch (k.tag) {
            case "array":
                this.resolveTy(k.ty);
                return;
            case "generic": {
                const re = this.resolveIdent(k.ident);
                if (re.tag !== "struct") {
                    this.rep.error(`identifier '${k.ident}' is not a struct`);
                    this.resols.set(ty.id, { tag: "error" });
                    return;
                }
                if (!re.struct.generics) {
                    this.rep.error(
                        `struct '${re.struct.ident}' does not accept generics`,
                    );
                    this.resols.set(ty.id, { tag: "error" });
                    return;
                }
                if (re.struct.generics.params.length !== k.args.length) {
                    this.rep.error(
                        `incorrect amount of generic paramters, expected ${re.struct.generics.params.length}, got ${k.args.length}`,
                    );
                    this.resols.set(ty.id, { tag: "error" });
                    return;
                }
                this.resols.set(ty.id, re);
                for (const arg of k.args) {
                    this.resolveTy(arg);
                }
                return;
            }
            case "struct_literal": {
                const fields = new Map<string, ast.Field>();
                for (const field of k.fields) {
                    this.resolveTy(field.ty);
                    if (fields.has(field.ident)) {
                        this.rep.error(
                            `field '${field.ident}' already defined`,
                            field.loc,
                        );
                        const otherField = fields.get(field.ident)!;
                        this.rep.info(
                            `original field '${otherField.ident}' defined here`,
                            otherField.loc,
                        );
                        continue;
                    }
                    fields.set(field.ident, field);
                }
                return;
            }
            case "ident": {
                const re = this.resolveIdent(k.ident);
                if (re.tag === "undefined") {
                    this.rep.error(
                        `identifier '${k.ident}' is not defined`,
                        ty.loc,
                    );
                }
                this.resols.set(ty.id, re);
                return;
            }
            case "int":
                return;
        }
        const _: never = k;
    }

    private reportAlreadyDefined(ident: string, loc: ast.Loc) {
        const re = this.resolveIdent(ident);
        switch (re.tag) {
            case "error":
            case "undefined":
                throw new Error();
            case "struct":
                this.rep.error(`identifier '${ident}' already defined`, loc);
                this.rep.info(
                    `struct '${re.struct.ident}' defined here`,
                    re.struct.loc,
                );
                break;
            case "generic":
                this.rep.error(`identifier '${ident}' already defined`, loc);
                this.rep.info(
                    `generic parameter '${
                        re.struct.generics!.params[re.idx]
                    }' defined here`,
                    re.struct.loc,
                );
                break;
            case "primitive":
                this.rep.error(`cannot redefine primitive '${ident}'`, loc);
                break;
        }
    }

    private static primitives = ["bool", "int", "float", "str"];

    private resolveIdent(ident: string): Resol {
        if (Resolver.primitives.includes(ident)) {
            return { tag: "primitive" };
        }
        for (const rib of this.ribs.toReversed()) {
            if (rib.defs.has(ident)) {
                return rib.defs.get(ident)!;
            }
        }
        return { tag: "undefined" };
    }

    private define(ident: string, re: Resol): "ok" | "already defined" {
        if (this.rib().defs.has(ident)) {
            return "already defined";
        }
        this.rib().defs.set(ident, re);
        return "ok";
    }

    private rib(): Rib {
        return this.ribs.at(-1)!;
    }
}

class Reporter {
    public errorOccured = false;

    public constructor(
        private text: string,
    ) {}

    public error(msg: string, loc?: ast.Loc) {
        this.errorOccured = true;
        console.error(
            `%cerror%c: ${msg}%c`,
            "color: red;   font-weight: bold",
            "font-weight: bold",
            "",
        );
        if (loc) {
            this.printLoc(loc);
        }
    }

    public info(msg: string, loc?: ast.Loc) {
        console.error(
            `%cerror%c: ${msg}%c`,
            "color: cyan; font-weight: bold",
            "font-weight: bold",
            "",
        );
        if (loc) {
            this.printLoc(loc);
        }
    }

    private printLoc(loc: ast.Loc) {
        const line = this.text.split("\n")[loc.start.line - 1];
        const posPad = " ".repeat(loc.start.column - 1);
        const lineNr = loc.start.line.toString().padStart(3, " ");
        const lPad = " ".repeat(lineNr.length + 1);
        const pos = `./${loc.source}:${loc.start.line}:${loc.start.column}`;
        console.error(
            `%c --> ${pos}\n${lPad}|\n${lineNr} |%c${line}%c\n${lPad}|${posPad}^%c`,
            "color: gray",
            "color: lightgray",
            "color: gray",
            "",
        );
    }
}

if (import.meta.main) {
    main();
}