export class Mod {
    public fns: Fn[] = [];

    constructor() {}

    createFn(): Fn {
        return new Fn(this);
    }
}

export class Fn {
    public bbs: BasicBlock[] = [];

    constructor(
        public mod: Mod,
    ) {}

    createBasicBlock(): BasicBlock {
        return new BasicBlock(this);
    }

    pushBasicBlock(bb: BasicBlock) {
        this.bbs.push(bb);
    }
}

export class BasicBlock {
    public insts: Inst[] = [];

    constructor(
        public fn: Fn,
    ) {}
}

export class Inst {
    constructor(
        public bb: BasicBlock,
        public ty: Ty,
        public kind: InstKind,
    ) {}
}

export type InstKind =
    | { tag: "Int"; value: number }
    | { tag: "Bool"; value: boolean }
    | { tag: "Fn"; fn: Fn }
    | { tag: "Param"; idx: number }
    | { tag: "GetElemPtr"; base: Inst; offset: number }
    | { tag: "Call"; callee: Inst; args: Inst[] }
    | { tag: "Alloca" }
    | { tag: "Load"; source: Inst }
    | { tag: "Store"; target: Inst; source: Inst }
    | { tag: "ReturnVoid" }
    | { tag: "Return"; source: Inst }
    | { tag: "Jump"; target: BasicBlock }
    | { tag: "Branch"; cond: Inst; truthy: BasicBlock; falsy: BasicBlock }
    | { tag: "Select"; cond: Inst; truthy: Inst; falsy: Inst }
    | { tag: "Trunc"; source: Inst }
    | { tag: "ZeroExt"; source: Inst }
    | { tag: "SignExt"; source: Inst }
    | { tag: UnaryOp; source: Inst; signed: boolean }
    | { tag: BinaryOp; left: Inst; right: Inst; signed: boolean };

export type UnaryOp = "Not" | "Neg";

export type BinaryOp =
    | "Eq"
    | "Ne"
    | "Lt"
    | "Gt"
    | "Lte"
    | "Gte"
    | "Or"
    | "Xor"
    | "And"
    | "Shl"
    | "Shr"
    | "Add"
    | "Sub"
    | "Mul"
    | "Div"
    | "Rem";

export class Ty {
    constructor(
        public kind: TyKind,
    ) {}

    hash(): string {
        return JSON.stringify(this.kind);
    }
}

export type TyKind =
    | { tag: "I8" | "I16" | "I32" | "I64" }
    | { tag: "Ptr" }
    | { tag: "Array"; ty: Ty; length: number };

export class Context {
    private tys = new Map<string, Ty>();
    private instsHashIds = new Map<Inst, number>();
    private fnHashIds = new Map<Fn, number>();
    private bbHashIds = new Map<BasicBlock, number>();

    private insts = new Map<string, Inst>();

    createTy<
        Tag extends TyKind["tag"],
    >(
        tag: Tag,
        kind: Omit<TyKind & { tag: Tag }, "tag">,
    ): Ty {
        const ty = new Ty({ ...kind, tag } as TyKind);
        const key = ty.hash();
        const found = this.tys.get(key);
        if (found) {
            return found;
        }
        this.tys.set(key, ty);
        return ty;
    }

    createInst(
        bb: BasicBlock,
        ty: Ty,
        kind: InstKind,
    ): Inst {
        const newId = this.instsHashIds.size;
        const inst = new Inst(bb, ty, kind);
        const key = this.hashInst(inst);
        const found = this.insts.get(key);
        if (found) {
            return found;
        }
        this.insts.set(key, inst);
        this.instsHashIds.set(inst, newId);
        return inst;
    }

    hashInst(inst: Inst): string {
        const id = this.instsHashIds.get(inst);
        if (id) {
            return `%${id}`;
        }

        const i = (inst: Inst) => this.hashInst(inst);
        const fn = (fn: Fn): string => {
            let id = this.fnHashIds.get(fn);
            if (!id) {
                id = this.fnHashIds.size;
                this.fnHashIds.set(fn, id);
            }
            return `%${id}`;
        };
        const bb = (bb: BasicBlock): string => {
            let id = this.bbHashIds.get(bb);
            if (!id) {
                id = this.bbHashIds.size;
                this.bbHashIds.set(bb, id);
            }
            return `%${id}`;
        };

        const k = inst.kind;
        switch (k.tag) {
            case "Int":
            case "Bool":
                return `${k.tag}-${k.value}`;
            case "Fn":
                return `${k.tag}-${fn(k.fn)}`;
            case "Param":
                return `${k.tag}-${k.idx}`;
            case "GetElemPtr":
                return `${k.tag}-${i(k.base)}-${k.offset}`;
            case "Call":
                return `${k.tag}-${i(k.callee)}-${k.args.map(i).join("-")}`;
            case "Alloca":
                return `${k.tag}`;
            case "Load":
                return `${k.tag}-${i(k.source)}`;
            case "Store":
                return `${k.tag}-${i(k.target)}-${i(k.source)}`;
            case "ReturnVoid":
                return `${k.tag}`;
            case "Return":
                return `${k.tag}-${i(k.source)}`;
            case "Jump":
                return `${k.tag}-${bb(k.target)}`;
            case "Branch":
                return `${k.tag}-${i(k.cond)}-${bb(k.truthy)}-${bb(k.falsy)}`;
            case "Select":
                return `${k.tag}-${i(k.cond)}-${i(k.truthy)}-${i(k.falsy)}`;
            case "Trunc":
                return `${k.tag}-${i(k.source)}`;
            case "ZeroExt":
                return `${k.tag}-${i(k.source)}`;
            case "SignExt":
                return `${k.tag}-${i(k.source)}`;
            case "Not":
            case "Neg":
                return `${k.tag}-${i(k.source)}-${k.signed}`;
            case "Eq":
            case "Ne":
            case "Lt":
            case "Gt":
            case "Lte":
            case "Gte":
            case "Or":
            case "Xor":
            case "And":
            case "Shl":
            case "Shr":
            case "Add":
            case "Sub":
            case "Mul":
            case "Div":
            case "Rem":
                return `${k.tag}-${i(k.left)}-${i(k.right)}-${k.signed}`;
        }
    }
}

export class Builder {
    constructor(
        private bb: BasicBlock,
        private cx: Context,
    ) {}

    createInt(ty: Ty, value: number) {
        this.push(ty, { tag: "Int", value });
    }

    private push(ty: Ty, kind: InstKind) {
        this.bb.insts.push(this.cx.createInst(this.bb, ty, kind));
    }
}
