const std = @import("std");
const ink = @import("../root.zig");
const Ast = ink.Ast;
const Render = @This();

gpa: std.mem.Allocator,
tty_config: std.Io.tty.Config,
tree: *const Ast,
prefix: Prefix,
lines: LineCache,

const LineCache = struct {
    ranges: std.ArrayListUnmanaged(SourceRange) = .empty,

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

    pub fn build(gpa: std.mem.Allocator, bytes: []const u8) !LineCache {
        var lines: LineCache = .{};
        var start: usize = 0;

        for (bytes, 0..) |c, i| {
            if (c == '\n') {
                try lines.ranges.append(gpa, .{ .start = start, .end = i });
                start = i + 1;
            }
        }
        if (start != bytes.len or bytes.len == 0) {
            try lines.ranges.append(gpa, .{ .start = start, .end = bytes.len });
        }
        return lines;
    }

    pub fn calculateLine(self: *const LineCache, offset: usize) usize {
        for (self.ranges.items, 0..) |r, i| {
            if (offset >= r.start and offset <= r.end)
                return i;
        }
        return self.ranges.items.len - 1;
    }
};

pub const SourceRange = struct {
    start: usize,
    end: usize,
};

pub const ErrorInfo = struct {
    filename: []const u8,
    message: []const u8,
    snippet: []const u8,
    line: usize,
    column: usize,
};

fn nodeTagToString(tag: Ast.Node.Tag) []const u8 {
    return switch (tag) {
        .file => "File",
        .false_literal => "FalseLiteral",
        .true_literal => "TrueLiteral",
        .number_literal => "NumberLiteral",
        .string_literal => "StringLiteral",
        .empty_string => "EmptyString",
        .identifier => "Identifier",
        .add_expr => "AddExpr",
        .subtract_expr => "SubtractExpr",
        .multiply_expr => "MultiplyExpr",
        .divide_expr => "DivideExpr",
        .mod_expr => "ModExpr",
        .negate_expr => "NegateExpr",
        .logical_equality_expr => "LogicalEqualityExpr",
        .logical_inequality_expr => "LogicalInequalityExpr",
        .logical_and_expr => "LogicalAndExpr",
        .logical_or_expr => "LogicalOrExpr",
        .logical_not_expr => "LogicalNotExpr",
        .logical_greater_expr => "LogicalGreaterThanExpr",
        .logical_greater_or_equal_expr => "LogicalGreaterThanOrEqualExpr",
        .logical_lesser_or_equal_expr => "LogicalLesserThanOrEqualExpr",
        .logical_lesser_expr => "LogicalLesserThanExpr",
        .divert_expr => "Divert",
        .selector_expr => "SelectorExpr",
        .call_expr => "CallExpr",
        .choice_expr => "ChoiceContentExpr",
        .choice_start_expr => "ChoiceStartContentExpr",
        .choice_option_expr => "ChoiceOptionContentExpr",
        .choice_inner_expr => "ChoiceInnerContentExpr",
        .assign_stmt => "AssignStmt",
        .block_stmt => "BlockStmt",
        .content_stmt => "ContentStmt",
        .divert_stmt => "DivertStmt",
        .return_stmt => "ReturnStmt",
        .expr_stmt => "ExprStmt",
        .choice_stmt => "ChoiceStmt",
        .choice_star_stmt => "ChoiceStarStmt",
        .choice_plus_stmt => "ChoicePlusStmt",
        .gather_point_stmt => "GatherPointStmt",
        .gathered_stmt => "GatheredStmt",
        .function_prototype => "FunctionProto",
        .stitch_prototype => "StitchProto",
        .knot_prototype => "KnotProto",
        .function_decl => "FunctionDecl",
        .stitch_decl => "StitchDecl",
        .knot_decl => "KnotDecl",
        .const_decl => "ConstDecl",
        .var_decl => "VarDecl",
        .list_decl => "ListDecl",
        .temp_decl => "TempDecl",
        .parameter_decl => "ParamDecl",
        .ref_parameter_decl => "RefParamDecl",
        .argument_list => "ArgumentList",
        .parameter_list => "ParamList",
        .string_expr => "StringExpr",
        .switch_stmt => "SwitchStmt",
        .switch_case => "SwitchCase",
        .multi_if_stmt => "MultiIfStmt",
        .if_stmt => "IfStmt",
        .if_branch => "IfBranch",
        .else_branch => "ElseBranch",
        .inline_logic_expr => "InlineLogicExpr",
        .content => "Content",
        .invalid => "Invalid",
    };
}

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: anytype) !void {
        try writer.writeAll(self.buf.items);
    }

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

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

fn writeType(r: *Render, writer: *std.Io.Writer, node: *const Ast.Node) !void {
    const tag = nodeTagToString(node.tag);
    try r.tty_config.setColor(writer, .magenta);
    try r.tty_config.setColor(writer, .bold);
    try writer.writeAll(tag);
    try r.tty_config.setColor(writer, .reset);
}

fn writeLexeme(r: *Render, writer: *std.Io.Writer, node: *const Ast.Node) !void {
    try r.tty_config.setColor(writer, .yellow);
    try writer.print("`{s}`", .{r.tree.nodeSlice(node)});
    try r.tty_config.setColor(writer, .reset);
}

fn writeLineSpan(r: *Render, writer: *std.Io.Writer, node: *const Ast.Node) !void {
    const line_start = r.lines.calculateLine(node.loc.start);
    const line_end = r.lines.calculateLine(node.loc.end);

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte('<');
    try r.tty_config.setColor(writer, .reset);

    try r.tty_config.setColor(writer, .yellow);
    try writer.print("line:{d}", .{line_start + 1});
    try r.tty_config.setColor(writer, .reset);

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte(',');
    try r.tty_config.setColor(writer, .reset);

    try writer.writeByte(' ');

    try r.tty_config.setColor(writer, .yellow);
    try writer.print("line:{d}", .{line_end + 1});
    try r.tty_config.setColor(writer, .reset);

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte('>');
    try r.tty_config.setColor(writer, .reset);
}

fn writeColumnSpan(r: *Render, writer: *std.Io.Writer, node: *const Ast.Node) !void {
    const line_start = r.lines.calculateLine(node.loc.start);
    const line_range = r.lines.ranges.items[line_start];
    const column_start = (node.loc.start - line_range.start);
    const column_end = (node.loc.end - line_range.start);

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte('<');
    try r.tty_config.setColor(writer, .reset);

    try r.tty_config.setColor(writer, .yellow);
    try writer.print("col:{d}", .{column_start + 1});
    try r.tty_config.setColor(writer, .reset);

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte(',');
    try r.tty_config.setColor(writer, .reset);

    try writer.writeByte(' ');

    try r.tty_config.setColor(writer, .yellow);
    try writer.print("col:{d}", .{column_end + 1});
    try r.tty_config.setColor(writer, .reset);

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte('>');
    try r.tty_config.setColor(writer, .reset);
}

fn writeLineColumnSpan(r: *Render, writer: *std.Io.Writer, node: *const Ast.Node) !void {
    const line_start = r.lines.calculateLine(node.loc.start);
    const line_range = r.lines.ranges.items[line_start];
    const column_start = (node.loc.start - line_range.start);
    const column_end = (node.loc.end - line_range.start);

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte('<');
    try r.tty_config.setColor(writer, .reset);

    try r.tty_config.setColor(writer, .yellow);
    try writer.print("line:{d}", .{line_start + 1});

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte(',');
    try r.tty_config.setColor(writer, .reset);

    try writer.writeByte(' ');

    try r.tty_config.setColor(writer, .yellow);
    try writer.print("col:{d}:{d}", .{ column_start + 1, column_end + 1 });
    try r.tty_config.setColor(writer, .reset);

    try r.tty_config.setColor(writer, .white);
    try writer.writeByte('>');
    try r.tty_config.setColor(writer, .reset);
}

const NodeOptions = struct {
    is_last: bool,
    is_root: bool,
};

fn renderAstNode(
    r: *Render,
    writer: *std.Io.Writer,
    node: *const Ast.Node,
    options: NodeOptions,
) !void {
    if (!options.is_root) {
        try r.tty_config.setColor(writer, .blue);
        try r.prefix.writeIndent(writer);
        if (options.is_last) {
            try writer.writeAll("`--");
        } else {
            try writer.writeAll("|--");
        }
        try r.tty_config.setColor(writer, .reset);
    }

    switch (node.tag) {
        .file => {
            try r.writeType(writer, node);
            try writer.writeByte(' ');
            try writer.writeByte('"');
            try writer.writeAll(r.tree.filename);
            try writer.writeByte('"');
        },
        .block_stmt,
        .choice_stmt,
        .function_decl,
        .gathered_stmt,
        .knot_decl,
        .stitch_decl,
        .if_stmt,
        .multi_if_stmt,
        .switch_stmt,
        => {
            try r.writeType(writer, node);
            try writer.writeByte(' ');
            try r.writeLineSpan(writer, node);
        },
        .assign_stmt,
        .choice_plus_stmt,
        .choice_star_stmt,
        .const_decl,
        .content_stmt,
        .divert_stmt,
        .expr_stmt,
        .gather_point_stmt,
        .list_decl,
        .return_stmt,
        .temp_decl,
        .var_decl,
        => {
            try r.writeType(writer, node);
            try writer.writeByte(' ');
            try r.writeLineColumnSpan(writer, node);
        },
        .choice_start_expr,
        .choice_option_expr,
        .choice_inner_expr,
        .identifier,
        .number_literal,
        .parameter_decl,
        .ref_parameter_decl,
        .string_literal,
        .string_expr,
        => {
            try r.writeType(writer, node);
            try writer.writeByte(' ');
            try r.writeLexeme(writer, node);
            try writer.writeByte(' ');
            try r.writeColumnSpan(writer, node);
        },
        else => {
            try r.writeType(writer, node);
            try writer.writeByte(' ');
            try r.writeColumnSpan(writer, node);
        },
    }
    try writer.writeByte('\n');
    return writer.flush();
}

fn renderAstWalk(
    r: *Render,
    writer: *std.Io.Writer,
    node: *const Ast.Node,
    options: NodeOptions,
) !void {
    var children: std.ArrayListUnmanaged(?*const Ast.Node) = .empty;
    defer children.deinit(r.gpa);

    try renderAstNode(r, writer, node, options);
    switch (node.tag) {
        .false_literal,
        .true_literal,
        .number_literal,
        .string_literal,
        .empty_string,
        .identifier,
        .parameter_decl,
        .ref_parameter_decl,
        .choice_start_expr,
        .choice_option_expr,
        .choice_inner_expr,
        => {},
        .file,
        .argument_list,
        .parameter_list,
        .block_stmt,
        .choice_stmt,
        .content,
        => {
            const list = node.data.list.items;
            if (list) |items| for (items) |n| {
                try children.append(r.gpa, n);
            };
        },
        .choice_expr => {
            const lhs = node.data.choice_expr.start_expr;
            const mhs = node.data.choice_expr.option_expr;
            const rhs = node.data.choice_expr.inner_expr;

            if (lhs) |n| try children.append(r.gpa, n);
            if (mhs) |n| try children.append(r.gpa, n);
            if (rhs) |n| try children.append(r.gpa, n);
        },
        .add_expr,
        .subtract_expr,
        .multiply_expr,
        .divide_expr,
        .mod_expr,
        .negate_expr,
        .logical_equality_expr,
        .logical_inequality_expr,
        .logical_and_expr,
        .logical_or_expr,
        .logical_not_expr,
        .logical_greater_expr,
        .logical_greater_or_equal_expr,
        .logical_lesser_or_equal_expr,
        .logical_lesser_expr,
        .assign_stmt,
        .divert_stmt,
        .return_stmt,
        .expr_stmt,
        .const_decl,
        .var_decl,
        .list_decl,
        .temp_decl,
        .content_stmt,
        .choice_star_stmt,
        .choice_plus_stmt,
        .gather_point_stmt,
        .gathered_stmt,
        .function_prototype,
        .stitch_prototype,
        .knot_prototype,
        .function_decl,
        .stitch_decl,
        .string_expr,
        .divert_expr,
        .selector_expr,
        .call_expr,
        .switch_case,
        .if_branch,
        .else_branch,
        .invalid,
        => {
            const lhs = node.data.bin.lhs;
            const rhs = node.data.bin.rhs;

            if (lhs) |n| try children.append(r.gpa, n);
            if (rhs) |n| try children.append(r.gpa, n);
        },
        .knot_decl => {
            const lhs = node.data.knot_decl.prototype;
            try children.append(r.gpa, lhs);

            const list = node.data.knot_decl.children;
            if (list) |items| for (items) |n| {
                try children.append(r.gpa, n);
            };
        },
        .switch_stmt, .if_stmt, .multi_if_stmt => {
            const expr = node.data.switch_stmt.condition_expr;
            if (expr) |n| try children.append(r.gpa, n);

            const list = node.data.switch_stmt.cases;
            for (list) |case_stmt| {
                try children.append(r.gpa, case_stmt);
            }
        },
        .inline_logic_expr => {
            const lhs = node.data.bin.lhs;
            if (lhs) |n| try children.append(r.gpa, n);
        },
    }
    for (children.items, 0..) |maybe_child, i| {
        const old_len = if (!options.is_root)
            try r.prefix.pushChildPrefix(r.gpa, options.is_last)
        else
            0;
        const child_is_last = (i == children.items.len - 1);

        if (maybe_child) |child| {
            try renderAstWalk(r, writer, child, .{
                .is_last = child_is_last,
                .is_root = false,
            });
        }
        if (!options.is_root) r.prefix.restore(old_len);
    }
}

pub const Options = struct {
    use_color: bool,
};

pub fn renderTree(
    gpa: std.mem.Allocator,
    writer: *std.Io.Writer,
    ast: *const Ast,
    options: Options,
) !void {
    var r: Render = .{
        .gpa = gpa,
        .tree = ast,
        .tty_config = if (options.use_color) .escape_codes else .no_color,
        .prefix = .{},
        .lines = try LineCache.build(gpa, ast.source),
    };
    defer r.prefix.deinit(gpa);
    defer r.lines.deinit(gpa);

    if (ast.root) |root| {
        try r.renderAstWalk(writer, root, .{
            .is_root = true,
            .is_last = true,
        });
    }
    try writer.flush();
}