const std = @import("std");
const Ast = @import("Ast.zig");
const assert = std.debug.assert;
const Ir = @This();

string_bytes: []u8,
instructions: []Inst,
globals: []Global,
extra: []u32,
errors: []Ast.Error,

pub const Inst = struct {
    tag: Tag,
    data: Data,

    pub const Index = enum(u32) {
        file_inst,
        ref_start_index = 32,
        _,

        pub fn toRef(i: Index) Inst.Ref {
            return @enumFromInt(@intFromEnum(Index.ref_start_index) + @intFromEnum(i));
        }
    };

    pub const Ref = enum(u32) {
        bool_true,
        bool_false,
        none = std.math.maxInt(u32),
        _,

        pub fn toIndex(inst: Ref) ?Index {
            assert(inst != .none);
            const ref_int = @intFromEnum(inst);
            if (ref_int >= @intFromEnum(Index.ref_start_index)) {
                return @enumFromInt(ref_int - @intFromEnum(Index.ref_start_index));
            } else {
                return null;
            }
        }

        pub fn toIndexAllowNone(inst: Ref) ?Index {
            if (inst == .none) return null;
            return toIndex(inst);
        }
    };

    pub const Tag = enum {
        file,
        declaration,
        decl_knot,
        decl_var,
        decl_ref,
        alloc,
        load,
        store,
        add,
        sub,
        mul,
        div,
        mod,
        neg,
        not,
        cmp_eq,
        cmp_neq,
        cmp_gt,
        cmp_gte,
        cmp_lt,
        cmp_lte,
        integer,
        string,
        block,
        condbr,
        @"break",
        switch_br,
        content_push,
        content_flush,
    };

    pub const Data = union {
        payload: struct {
            payload_index: u32,
        },
        un: struct {
            lhs: Ref,
        },
        bin: struct {
            lhs: Ref,
            rhs: Ref,
        },
        integer: struct {
            value: u64,
        },
        string: struct {
            /// Offset into `string_bytes`.
            start: NullTerminatedString,

            pub fn get(self: @This(), ir: Ir) []const u8 {
                return nullTerminatedString(ir, self.start);
            }
        },
    };

    pub const Declaration = struct {
        name: NullTerminatedString,
        value: Index,
    };

    pub const Knot = struct {
        body_len: u32,
    };

    pub const Var = struct {
        body_len: u32,
    };

    pub const Block = struct {
        body_len: u32,
    };

    pub const Break = struct {
        block_inst: Index,
    };

    pub const CondBr = struct {
        condition: Ref,
        then_body_len: u32,
        else_body_len: u32,
    };

    pub const SwitchBr = struct {
        operand: Ref,
        cases_len: u32,
        else_body_len: u32,

        pub const Case = struct {
            operand: Ref,
            body_len: u32,
        };
    };
};

pub const Global = struct {
    tag: Tag,
    name: Ir.NullTerminatedString,
    is_constant: bool,

    pub const Tag = enum {
        knot,
        variable,
    };
};

pub const NullTerminatedString = enum(u32) {
    empty,
    _,
};

pub const IndexSlice = struct {
    index: u32,
    len: u32,
};

pub fn deinit(ir: *Ir, gpa: std.mem.Allocator) void {
    gpa.free(ir.string_bytes);
    gpa.free(ir.instructions);
    gpa.free(ir.globals);
    gpa.free(ir.extra);
    gpa.free(ir.errors);
    ir.* = undefined;
}

const Render = struct {
    gpa: std.mem.Allocator,
    prefix: Prefix,
    writer: *std.Io.Writer,

    pub const Error = error{
        OutOfMemory,
        WriteFailed,
    };

    const Prefix = struct {
        buf: std.ArrayListUnmanaged(u8) = .empty,

        pub fn deinit(self: *Prefix, gpa: std.mem.Allocator) void {
            self.buf.deinit(gpa);
        }

        pub fn writeIndent(self: *const Prefix, writer: *std.Io.Writer) !void {
            try writer.writeAll(self.buf.items);
        }

        pub fn pushChildPrefix(self: *Prefix, gpa: std.mem.Allocator) !usize {
            const old_len = self.buf.items.len;
            const seg: []const u8 = "  ";
            try self.buf.appendSlice(gpa, seg);
            return old_len;
        }

        pub fn restore(self: *Prefix, new_len: usize) void {
            self.buf.shrinkRetainingCapacity(new_len);
        }
    };

    fn renderInstIndex(r: *Render, index: Inst.Index) !void {
        const io_w = r.writer;
        return io_w.print("%{d}", .{@intFromEnum(index)});
    }

    fn renderInstRef(r: *Render, ref: Inst.Ref) !void {
        const io_w = r.writer;
        if (ref == .none) {
            return io_w.writeAll(".none");
        } else if (ref.toIndex()) |i| {
            return r.renderInstIndex(i);
        } else {
            return io_w.print("@{s}", .{@tagName(ref)});
        }
    }

    fn renderSimple(r: *Render, inst: Inst) Error!void {
        const io_w = r.writer;
        return io_w.print("{s}(?)", .{@tagName(inst.tag)});
    }

    fn renderUnary(r: *Render, inst: Inst) Error!void {
        const io_w = r.writer;
        const data = inst.data.un;
        try io_w.print("{s}(", .{@tagName(inst.tag)});

        const lhs = data.lhs;
        try renderInstRef(r, lhs);
        return io_w.writeAll(")");
    }

    fn renderBinary(r: *Render, inst: Inst) Error!void {
        const io_w = r.writer;
        const data = inst.data.bin;
        try io_w.print("{s}(", .{@tagName(inst.tag)});
        try renderInstRef(r, data.lhs);
        try io_w.writeAll(", ");
        try renderInstRef(r, data.rhs);
        return io_w.writeAll(")");
    }

    fn renderBodyInner(r: *Render, ir: Ir, body_list: []const Inst.Index) Error!void {
        const io_w = r.writer;
        try io_w.writeAll("{\n");
        {
            const old_len = try r.prefix.pushChildPrefix(r.gpa);
            defer r.prefix.restore(old_len);

            for (body_list) |inst| try r.renderInst(ir, inst);
        }
        try r.prefix.writeIndent(io_w);
        try io_w.writeAll("}");
    }

    fn renderBlock(r: *Render, ir: Ir, inst: Inst) Error!void {
        const io_w = r.writer;
        const extra = ir.extraData(Inst.Block, inst.data.payload.payload_index);
        const body_slice = ir.bodySlice(extra.end, extra.data.body_len);

        try io_w.print("{s}(", .{@tagName(inst.tag)});
        try renderBodyInner(r, ir, body_slice);
        try io_w.writeAll(")");
    }

    fn renderBreak(r: *Render, ir: Ir, inst: Inst) Error!void {
        const io_w = r.writer;
        const extra = ir.extraData(Inst.Break, inst.data.payload.payload_index);

        try io_w.print("{s}(%{d})", .{ @tagName(inst.tag), extra.data.block_inst });
    }

    fn renderCondbr(r: *Render, ir: Ir, inst: Inst) Error!void {
        const io_w = r.writer;
        const extra = ir.extraData(Inst.CondBr, inst.data.payload.payload_index);
        const then_body = ir.bodySlice(extra.end, extra.data.then_body_len);
        const else_body = ir.bodySlice(extra.end + then_body.len, extra.data.else_body_len);

        try io_w.print("{s}(", .{@tagName(inst.tag)});
        try renderInstRef(r, extra.data.condition);
        try io_w.writeAll(", ");
        try renderBodyInner(r, ir, then_body);
        try io_w.writeAll(", ");
        try renderBodyInner(r, ir, else_body);
        try io_w.writeAll(")");
    }

    fn renderSwitchBr(r: *Render, ir: Ir, inst: Inst) Error!void {
        const io_w = r.writer;
        const switch_node = inst.data.payload;
        const switch_extra = ir.extraData(Inst.SwitchBr, switch_node.payload_index);
        const cases_slice = ir.bodySlice(switch_extra.end, switch_extra.data.cases_len);

        try io_w.print("{s}(", .{@tagName(inst.tag)});
        try renderInstRef(r, switch_extra.data.operand);
        try io_w.print(", \n", .{});
        for (cases_slice) |case_index| {
            const case_extra = ir.extraData(Inst.SwitchBr.Case, @intFromEnum(case_index));
            const body_slice = ir.bodySlice(case_extra.end, case_extra.data.body_len);
            const old_len = try r.prefix.pushChildPrefix(r.gpa);
            defer r.prefix.restore(old_len);

            try r.prefix.writeIndent(io_w);
            try renderInstRef(r, case_extra.data.operand);
            try io_w.print(" = ", .{});
            try renderBodyInner(r, ir, body_slice);
            try io_w.writeAll(",\n");
        }

        const else_body = ir.bodySlice(
            switch_extra.end + switch_extra.data.cases_len,
            switch_extra.data.else_body_len,
        );
        {
            const old_len = try r.prefix.pushChildPrefix(r.gpa);
            defer r.prefix.restore(old_len);

            try r.prefix.writeIndent(io_w);
            try io_w.print("else = ", .{});
            try renderBodyInner(r, ir, else_body);
        }
        try io_w.writeAll(")");
    }

    fn renderKnotDecl(r: *Render, ir: Ir, inst: Inst) Error!void {
        const io_w = r.writer;
        const extra = ir.extraData(Inst.Knot, inst.data.payload.payload_index);
        const body_slice = ir.bodySlice(extra.end, extra.data.body_len);

        try io_w.print("{s}(body=", .{@tagName(inst.tag)});
        try renderBodyInner(r, ir, body_slice);
        try io_w.writeAll(")");
    }

    fn renderVarDecl(r: *Render, ir: Ir, inst: Inst) Error!void {
        const io_w = r.writer;
        const extra = ir.extraData(Inst.Knot, inst.data.payload.payload_index);
        const body_slice = ir.bodySlice(extra.end, extra.data.body_len);

        try io_w.print("{s}(body=", .{@tagName(inst.tag)});
        try renderBodyInner(r, ir, body_slice);
        try io_w.writeAll(")");
    }

    fn renderDeclaration(r: *Render, ir: Ir, inst: Inst) Error!void {
        const io_w = r.writer;
        const extra = ir.extraData(Inst.Declaration, inst.data.payload.payload_index);
        const ident_str = ir.nullTerminatedString(extra.data.name);

        try io_w.print("{s}(name=\"{s}\", value={{\n", .{ @tagName(inst.tag), ident_str });
        {
            const old_len = try r.prefix.pushChildPrefix(r.gpa);
            defer r.prefix.restore(old_len);
            try renderInst(r, ir, extra.data.value);
        }
        try r.prefix.writeIndent(io_w);
        return io_w.writeAll(")");
    }

    fn renderFile(r: *Render, ir: Ir, inst: Inst) Error!void {
        const io_w = r.writer;
        const extra = ir.extraData(Inst.Block, inst.data.payload.payload_index);
        const body_slice = ir.bodySlice(extra.end, extra.data.body_len);

        try io_w.print("{s}(", .{@tagName(inst.tag)});
        try renderBodyInner(r, ir, body_slice);
        return io_w.writeAll(")");
    }

    fn renderInst(r: *Render, ir: Ir, index: Inst.Index) Error!void {
        const io_w = r.writer;
        const inst_index: u32 = @intFromEnum(index);
        const inst = ir.instructions[inst_index];

        try r.prefix.writeIndent(io_w);
        try io_w.print("%{d} = ", .{inst_index});
        switch (inst.tag) {
            .file => try r.renderFile(ir, inst),
            .declaration => try r.renderDeclaration(ir, inst),
            .decl_knot => try r.renderKnotDecl(ir, inst),
            .decl_var => try r.renderVarDecl(ir, inst),
            .decl_ref => {
                const str_bytes = inst.data.string.get(ir);
                try io_w.print("{s}(\"{s}\")", .{ @tagName(inst.tag), str_bytes });
            },
            .condbr => try r.renderCondbr(ir, inst),
            .@"break" => try r.renderBreak(ir, inst),
            .switch_br => try r.renderSwitchBr(ir, inst),
            .alloc => try r.renderSimple(inst),
            .load => try r.renderUnary(inst),
            .store => try r.renderBinary(inst),
            .block => try r.renderBlock(ir, inst),
            .add => try r.renderBinary(inst),
            .sub => try r.renderBinary(inst),
            .mul => try r.renderBinary(inst),
            .div => try r.renderBinary(inst),
            .mod => try r.renderBinary(inst),
            .neg => try r.renderUnary(inst),
            .not => try r.renderUnary(inst),
            .cmp_eq => try r.renderBinary(inst),
            .cmp_neq => try r.renderBinary(inst),
            .cmp_gt => try r.renderBinary(inst),
            .cmp_gte => try r.renderBinary(inst),
            .cmp_lt => try r.renderBinary(inst),
            .cmp_lte => try r.renderBinary(inst),
            .integer => {
                const value = inst.data.integer.value;
                try io_w.print("{s}({d})", .{ @tagName(inst.tag), value });
            },
            .string => {
                const str_bytes = inst.data.string.get(ir);
                try io_w.print("{s}(\"{s}\")", .{ @tagName(inst.tag), str_bytes });
            },
            .content_push => try r.renderUnary(inst),
            .content_flush => try r.renderUnary(inst),
        }
        try io_w.writeAll("\n");
    }
};

pub fn render(ir: Ir, gpa: std.mem.Allocator, writer: *std.Io.Writer) !void {
    std.debug.assert(ir.instructions.len > 0);
    var r = Render{ .gpa = gpa, .writer = writer, .prefix = .{} };
    defer r.prefix.deinit(gpa);

    try r.renderInst(ir, @enumFromInt(0));
    return writer.flush();
}

pub fn dumpInfo(ir: Ir, writer: *std.Io.Writer) !void {
    const bytes = ir.string_bytes;

    var start: usize = 0;
    while (start < bytes.len) {
        const end = std.mem.indexOfScalarPos(u8, bytes, start, 0) orelse break;
        const str = bytes[start..end];

        try writer.print("[{d:04}] ", .{start});
        for (str) |b| try writer.print("{x:02} ", .{b});
        try writer.print("00: {s}\n", .{str});
        start = end + 1;
    }
    for (ir.globals) |global| {
        try writer.print("{any}\n", .{global});
    }
    return writer.flush();
}

fn ExtraData(comptime T: type) type {
    return struct {
        data: T,
        end: usize,
    };
}

pub fn extraData(ir: Ir, comptime T: type, index: usize) ExtraData(T) {
    const fields = @typeInfo(T).@"struct".fields;
    var i: usize = index;
    var result: T = undefined;
    inline for (fields) |field| {
        @field(result, field.name) = switch (field.type) {
            u32 => ir.extra[i],
            Inst.Index => @enumFromInt(ir.extra[i]),
            Inst.Ref => @enumFromInt(ir.extra[i]),
            NullTerminatedString => @enumFromInt(ir.extra[i]),
            else => @compileError("bad field type"),
        };
        i += 1;
    }
    return .{ .data = result, .end = i };
}

pub fn nullTerminatedString(ir: Ir, index: NullTerminatedString) [:0]const u8 {
    const slice = ir.string_bytes[@intFromEnum(index)..];
    return slice[0..std.mem.indexOfScalar(u8, slice, 0).? :0];
}

pub fn bodySlice(ir: Ir, start: usize, len: usize) []Inst.Index {
    return @ptrCast(ir.extra[start..][0..len]);
}