const std = @import("std");
const Io = std.Io;
const Type = std.builtin.Type;

pub fn pretty(value: anytype) Pretty(@TypeOf(value)) {
    return Pretty(@TypeOf(value)).init(value);
}

pub fn prettyO(value: anytype, comptime opts: Options) PrettyWithOptions(@TypeOf(value), opts) {
    return PrettyWithOptions(@TypeOf(value), opts).init(value);
}

const root = @import("root");
const default_options: Options = if (@hasDecl(root, "pretty_options"))
    root.pretty_options
else
    Options{};

pub const Options = struct {
    /// Disable pretty printing.
    ///
    /// If enabled, will format values using `{any}`
    disable: bool = false,

    /// Show type names
    show_type_names: bool = true,
    /// Don't show type names for the root value
    skip_root_type_name: bool = false,

    /// Disable color output
    no_color: bool = false,
    /// Customize the theme, change colors or separators
    theme: Theme = .{},

    /// Inline print arrays
    array_inline: bool = false,
    /// Always show array indices, even in inline mode
    array_always_show_index: bool = false,

    /// Treat pointers to `[]u8` / `[]const u8` as strings
    ptr_array_u8_as_string: bool = true,
    /// Treat pointers to a u8 slice as strings
    ptr_slice_u8_as_string: bool = true,
    /// Inline print slices
    ptr_slice_inline: bool = false,
    /// Always show slice indices, even in inline mode
    ptr_slice_always_show_index: bool = false,

    /// Inline print structs
    struct_inline: bool = false,
    /// Highlight tag type of unions
    union_highlight_tag: bool = true,
};

pub const Theme = struct {
    pub const Color = enum {
        dim,
        field,
        value,
        @"error",
        type,
        string,
        many_ptr,
        null,
        true,
        false,
    };

    indent_width: comptime_int = 2,

    field_name_type_sep: []const u8 = ": ",
    type_value_sep: []const u8 = " = ",
    index_value_sep: []const u8 = " = ",

    index_open: []const u8 = "[",
    index_close: []const u8 = "]",

    array_open: []const u8 = "{",
    array_close: []const u8 = "}",

    vec_open: []const u8 = "(",
    vec_close: []const u8 = ")",

    color_dim: Io.Terminal.Color = .dim,
    color_field: Io.Terminal.Color = .green,
    color_value: Io.Terminal.Color = .blue,
    color_error: Io.Terminal.Color = .red,
    color_type: Io.Terminal.Color = .bright_blue,
    color_string: Io.Terminal.Color = .magenta,
    color_many_ptr: Io.Terminal.Color = .magenta,
    color_null: Io.Terminal.Color = .cyan,
    color_true: Io.Terminal.Color = .bright_green,
    color_false: Io.Terminal.Color = .bright_red,

    pub inline fn getColor(comptime this: Theme, comptime color: Color) Io.Terminal.Color {
        return @field(this, "color_" ++ @tagName(color));
    }
};

pub fn Pretty(comptime T: type) type {
    return PrettyWithOptions(T, default_options);
}

pub fn PrettyWithOptions(comptime T: type, comptime options: Options) type {
    if (default_options.disable) return struct {
        value: T,

        pub fn init(val: T) @This() {
            return .{ .value = val };
        }

        pub fn format(this: *const @This(), w: *std.Io.Writer) error{WriteFailed}!void {
            try w.print("{any}", .{this.value});
        }
    };

    const global = struct {
        var tty: ?Io.Terminal = null;
    };
    const ctx: Context = Context{ .options = options };

    return struct {
        value: T,

        pub fn init(val: T) @This() {
            return .{ .value = val };
        }

        pub fn format(this: *const @This(), w: *std.Io.Writer) error{WriteFailed}!void {
            if (global.tty == null) {
                var buffer: [1]u8 = undefined;
                const stderr = std.debug.lockStderr(&buffer).terminal();
                defer std.debug.unlockStderr();

                global.tty = stderr;
                global.tty.?.writer = w;
            }

            var run = Runtime{
                .out = w,
                .tty = global.tty.?,
                .no_color = options.no_color,
            };

            return run.pretty(ctx, this.value, .{});
        }
    };
}

pub const Runtime = struct {
    out: *Io.Writer,
    tty: Io.Terminal,

    no_color: bool = default_options.no_color,

    pub inline fn print(
        this: *const Runtime,
        comptime fmt: []const u8,
        args: anytype,
    ) error{WriteFailed}!void {
        return this.out.print(fmt, args);
    }

    pub inline fn write(this: *const Runtime, text: []const u8) error{WriteFailed}!void {
        _ = try this.out.write(text);
    }

    pub fn pretty(
        this: *const Runtime,
        comptime ctx: Context,
        value: anytype,
        comptime opts: InnerFmtOptions,
    ) error{WriteFailed}!void {
        return innerFmt(@TypeOf(value), ctx, this, value, opts);
    }

    pub inline fn setColor(
        this: *const Runtime,
        comptime ctx: Context,
        comptime color: Theme.Color,
    ) void {
        if (this.no_color or default_options.no_color) return;
        this.tty.setColor(comptime ctx.options.theme.getColor(color)) catch {};
    }

    pub inline fn setColorRaw(this: *const Runtime, color: Io.Terminal.Color) void {
        if (this.no_color or default_options.no_color) return;
        this.tty.setColor(color) catch {};
    }

    pub inline fn resetColor(this: *const Runtime) void {
        if (this.no_color or default_options.no_color) return;
        this.tty.setColor(.reset) catch {};
    }
};

pub const Context = struct {
    depth: comptime_int = 0,

    options: Options,
    exited_comptime: bool = false,

    pub fn inComptime(comptime this: Context) bool {
        if (!this.exited_comptime) return false;
        return @inComptime();
    }
};

pub const InnerFmtOptions = struct {
    skip_type_name: bool = !default_options.show_type_names,
};

fn innerFmt(
    comptime T: type,
    comptime ctx: Context,
    run: *const Runtime,
    value: T,
    comptime opts: InnerFmtOptions,
) error{WriteFailed}!void {
    const info = @typeInfo(T);

    if (!opts.skip_type_name)
        try printType(T, ctx, run, value);

    if (std.meta.hasMethod(T, "pretty")) {
        return value.pretty(ctx, run);
    }

    return switch (info) {
        .bool => formatBool(ctx, run, value),
        .null => formatNull(ctx, run),
        .type => formatType(ctx, run, value),

        // comptime types
        .comptime_int,
        .comptime_float,
        // number types
        .int,
        .float,
        // enum types
        .@"enum",
        .enum_literal,
        => formatValue(ctx, run, value),

        .optional => |opt| formatOptional(opt.child, ctx, run, value),
        .@"struct" => |st| formatStruct(T, ctx, st, run, value),
        .@"union" => formatUnion(T, ctx, run, value),

        .error_set => formatErrorSet(ctx, run, value),
        .error_union => formatErrorUnion(ctx, run, value),

        .array => |arr| formatArray(ctx, arr, run, value),
        .vector => |vec| formatVector(ctx, vec, run, value),

        .@"fn" => formatFn(T, ctx, run),

        .pointer => |ptr| switch (ptr.size) {
            .one => formatPtrOne(ctx, ptr, run, value),
            .slice => formatPtrSlice(ctx, ptr, run, value),
            .many => formatPtrMany(ctx, ptr, run, value),
            else => {
                run.setColor(ctx, .@"error");
                try run.print("Unimplemented! ({} = {any})", .{ info, value });
                run.resetColor();
            },
        },

        else => {
            run.setColor(ctx, .@"error");
            try run.print("Unimplemented! ({} = {any})", .{ info, value });
            run.resetColor();
        },
    };
}

inline fn printType(
    comptime T: type,
    comptime ctx: Context,
    run: *const Runtime,
    value: T,
) error{WriteFailed}!void {
    const active_type = comptime std.meta.activeTag(@typeInfo(T));
    const excluded = comptime switch (active_type) {
        .void, .noreturn, .undefined, .null, .@"fn" => true,
        else => false,
    };

    if ((ctx.depth != 0 or !ctx.options.skip_root_type_name) and !excluded) {
        run.setColor(ctx, .dim);

        if (ctx.options.show_type_names) {
            try run.write(@typeName(T));

            if (active_type == .@"union") {
                if (ctx.options.union_highlight_tag) {
                    run.resetColor();
                    run.setColor(ctx, .field);
                }

                try run.print("{}", .{std.meta.activeTag(value)});

                if (ctx.options.union_highlight_tag) {
                    run.resetColor();
                    run.setColor(ctx, .dim);
                }
            }
        }
        try run.write(ctx.options.theme.type_value_sep);

        run.resetColor();
    }
}

inline fn formatBool(
    comptime ctx: Context,
    run: *const Runtime,
    value: bool,
) !void {
    if (value)
        run.setColor(ctx, .true)
    else
        run.setColor(ctx, .false);

    try run.print("{}", .{value});
    run.resetColor();
}

inline fn formatNull(comptime ctx: Context, run: *const Runtime) !void {
    run.setColor(ctx, .null);
    try run.write("null");
    run.resetColor();
}

inline fn formatType(
    comptime ctx: Context,
    run: *const Runtime,
    value: type,
) !void {
    run.setColor(ctx, .type);
    run.setColorRaw(.bold);
    try run.write(@typeName(value));
    run.resetColor();
}

inline fn formatValue(
    comptime ctx: Context,
    run: *const Runtime,
    value: anytype,
) !void {
    run.setColor(ctx, .value);
    try run.print("{}", .{value});
    run.resetColor();
}

inline fn formatOptional(
    comptime T: type,
    comptime ctx: Context,
    run: *const Runtime,
    value: ?T,
) !void {
    return if (value) |val|
        innerFmt(T, ctx, run, val, .{ .skip_type_name = true })
    else
        formatNull(ctx, run);
}

inline fn formatStruct(
    comptime T: type,
    comptime ctx: Context,
    comptime st: Type.Struct,
    run: *const Runtime,
    value: T,
) !void {
    const next_ctx = Context{
        .depth = ctx.depth + 1,
        .exited_comptime = ctx.exited_comptime,
        .options = ctx.options,
    };

    if (ctx.options.struct_inline) {
        run.setColor(ctx, .dim);
        try run.write(".{ ");
        run.resetColor();
    }

    comptime var index = 0;
    inline for (st.fields) |field| {
        indent(next_ctx, next_ctx.options.struct_inline, run);
        if (index != 0 and ctx.options.struct_inline) {
            run.setColor(ctx, .dim);
            try run.write(", ");
            run.resetColor();
        }

        run.setColor(ctx, .field);
        try run.write("." ++ field.name ++
            ctx.options.theme.field_name_type_sep);
        run.resetColor();

        try innerFmt(field.type, next_ctx, run, @field(value, field.name), .{});

        index += 1;
    }

    if (ctx.options.struct_inline) {
        run.setColor(ctx, .dim);
        try run.write(" }");
        run.resetColor();
    }
}

inline fn formatUnion(
    comptime T: type,
    comptime ctx: Context,
    run: *const Runtime,
    value: T,
) !void {
    switch (value) {
        inline else => |val| try innerFmt(@TypeOf(val), ctx, run, val, .{ .skip_type_name = true }),
    }
}

inline fn formatArray(
    comptime ctx: Context,
    comptime arr: Type.Array,
    run: *const Runtime,
    value: anytype,
) !void {
    if (arr.len == 0) {
        try run.write(ctx.options.theme.array_open);
        try run.write(ctx.options.theme.array_close);
        return;
    }

    const next_ctx = Context{
        .depth = ctx.depth + 1,
        .exited_comptime = ctx.exited_comptime,
        .options = ctx.options,
    };

    if (ctx.options.array_inline) try arrayOpen(ctx, run);

    comptime var index = 0;
    inline for (value) |val| {
        indent(
            next_ctx,
            next_ctx.options.array_inline,
            run,
        );

        run.setColor(ctx, .dim);
        if (index != 0 and ctx.options.array_inline)
            try run.write(", ");

        if (!ctx.options.array_inline or ctx.options.array_always_show_index) {
            try run.write(ctx.options.theme.index_open);
            try run.print("{}", .{index});
            try run.write(ctx.options.theme.index_close ++
                ctx.options.theme.index_value_sep);
        }

        run.resetColor();

        try innerFmt(arr.child, next_ctx, run, val, .{ .skip_type_name = true });

        index += 1;
    }

    if (ctx.options.array_inline) try arrayClose(ctx, run);
}

inline fn formatVector(
    comptime ctx: Context,
    comptime vec: Type.Vector,
    run: *const Runtime,
    value: @Vector(vec.len, vec.child),
) !void {
    run.setColor(ctx, .dim);
    try run.write(ctx.options.theme.vec_open ++ " ");
    run.resetColor();

    inline for (0..vec.len) |idx| {
        if (idx != 0) {
            run.setColor(ctx, .dim);
            try run.write(", ");
            run.resetColor();
        }

        try innerFmt(vec.child, ctx, run, value[idx], .{ .skip_type_name = true });
    }

    run.setColor(ctx, .dim);
    try run.write(" " ++ ctx.options.theme.vec_close);
    run.resetColor();
}

inline fn formatPtrOne(
    comptime ctx: Context,
    comptime ptr: Type.Pointer,
    run: *const Runtime,
    value: anytype,
) !void {
    const is_string = switch (@typeInfo(ptr.child)) {
        .array => |arr| arr.child == u8 and ctx.options.ptr_array_u8_as_string,
        else => false,
    };
    if (is_string) return formatString(ctx, run, value);

    try innerFmt(ptr.child, ctx, run, value.*, .{ .skip_type_name = true });
}

inline fn formatPtrSlice(
    comptime ctx: Context,
    comptime ptr: Type.Pointer,
    run: *const Runtime,
    value: anytype,
) !void {
    if (ptr.child == u8 and ctx.options.ptr_slice_u8_as_string)
        return formatString(ctx, run, value);

    const next_ctx = Context{
        .depth = if (!ctx.options.ptr_slice_inline) ctx.depth + 1 else ctx.depth,
        .exited_comptime = ctx.exited_comptime,
        .options = ctx.options,
    };

    if (ctx.options.ptr_slice_inline) try arrayOpen(ctx, run);

    var count: usize = 0;
    for (value, 0..) |val, idx| {
        count += 1;

        indent(
            next_ctx,
            next_ctx.options.ptr_slice_inline,
            run,
        );

        run.setColor(ctx, .dim);
        if (idx != 0 and ctx.options.ptr_slice_inline)
            try run.write(", ");

        if (!ctx.options.ptr_slice_inline or ctx.options.ptr_slice_always_show_index) {
            try run.write(ctx.options.theme.index_open);
            try run.print("{}", .{idx});
            try run.write(ctx.options.theme.index_close ++
                ctx.options.theme.index_value_sep);
        }

        run.resetColor();

        try innerFmt(ptr.child, next_ctx, run, val, .{ .skip_type_name = true });
    }

    if (count == 0) {
        run.setColor(ctx, .dim);
        try run.write(ctx.options.theme.array_open);
        try run.write(ctx.options.theme.array_close);
        run.resetColor();
    }

    if (ctx.options.ptr_slice_inline) try arrayClose(ctx, run);
}

inline fn formatPtrMany(
    comptime ctx: Context,
    comptime ptr: Type.Pointer,
    run: *const Runtime,
    value: anytype,
) !void {
    _ = ptr;
    run.setColor(ctx, .dim);
    run.setColor(ctx, .many_ptr);
    try run.print("{*}", .{value});
    run.resetColor();
}

inline fn formatString(
    comptime ctx: Context,
    run: *const Runtime,
    value: anytype,
) !void {
    run.setColor(ctx, .string);
    try run.write("\"");
    try run.write(value);
    try run.write("\"");
    run.resetColor();
}

inline fn formatErrorSet(
    comptime ctx: Context,
    run: *const Runtime,
    value: anyerror,
) !void {
    run.setColor(ctx, .@"error");
    try run.write("error.");
    try run.write(@errorName(value));
    run.resetColor();
}

inline fn formatErrorUnion(
    comptime ctx: Context,
    run: *const Runtime,
    value: anytype,
) !void {
    const val = value catch |err| return formatErrorSet(ctx, run, err);
    return innerFmt(
        @TypeOf(val),
        ctx,
        run,
        val,
        .{ .skip_type_name = true },
    );
}

inline fn formatFn(
    comptime T: type,
    comptime ctx: Context,
    run: *const Runtime,
) !void {
    run.setColor(ctx, .dim);
    run.setColor(ctx, .type);
    run.setColorRaw(.bold);

    try run.write(@typeName(T));

    run.resetColor();
}

inline fn indent(
    comptime ctx: Context,
    comptime inline_option: bool,
    run: *const Runtime,
) void {
    if (inline_option) return;

    const text: [ctx.depth * ctx.options.theme.indent_width]u8 = @splat(' ');
    run.write("\n" ++ text) catch {};
}

fn arrayOpen(comptime ctx: Context, run: *const Runtime) !void {
    run.setColor(ctx, .dim);
    try run.write(ctx.options.theme.array_open ++ " ");
    run.resetColor();
}

fn arrayClose(comptime ctx: Context, run: *const Runtime) !void {
    run.setColor(ctx, .dim);
    try run.write(" " ++ ctx.options.theme.array_close);
    run.resetColor();
}