const std = @import("std");

pub fn main() !void {
    try execute(std.io.getStdOut().writer(), @embedFile("./test.tmp.html"), .{
        .community = .{ .name = "abcd" },
        .foo = [_][]const u8{ "5", "4", "3", "2", "1" },
        .baz = [_][]const []const u8{
            &.{ "5", "4", "3", "2", "1" },
            &.{ "5", "4", "3", "2", "1" },
        },
        .qux = true,
        .quxx = false,
    });
}

pub fn execute(writer: anytype, comptime template: []const u8, args: anytype) !void {
    @setEvalBranchQuota(@intCast(u32, template.len * 8));

    const tokens = comptime parseTemplateTokens(ControlTokenIter{ .text = template });
    const tmpl = comptime parseTemplate(tokens, 0, .root);
    try executeTemplate(writer, tmpl.item, args, .{});
}

fn executeTemplate(writer: anytype, comptime items: []const TemplateItem, args: anytype, captures: anytype) !void {
    inline for (items) |it| {
        switch (it) {
            .text => |text| try writer.writeAll(text),
            .statement => |stmt| try executeStatement(writer, stmt, args, captures),
        }
    }
}

fn executeStatement(writer: anytype, comptime stmt: Statement, args: anytype, captures: anytype) !void {
    switch (stmt) {
        .expression => |expr| {
            const val = evaluateExpression(expr, args, captures);
            try print(writer, val);
        },
        .@"for" => |loop| {
            const iterable = evaluateExpression(loop.header.iterable, args, captures);
            const subtemplate = loop.subtemplate;
            //std.log.debug("{any}", .{subtemplate});
            for (iterable) |v| {
                try executeTemplate(
                    writer,
                    subtemplate,
                    args,
                    addCapture(captures, loop.header.capture, v),
                );
            }
        },
        .@"if" => |if_stmt| {
            const condition = evaluateExpression(if_stmt.header.condition, args, captures);
            const subtemplate = if_stmt.subtemplate;
            if (condition) try executeTemplate(writer, subtemplate, args, captures);
        },
        //else => @compileError("TODO"),
    }
}

fn print(writer: anytype, arg: anytype) !void {
    const T = @TypeOf(arg);
    if (comptime std.meta.trait.isZigString(T)) return writer.writeAll(arg);
    if (comptime std.meta.trait.isNumber(T)) return std.fmt.format(writer, "{}", .{arg});
    @compileLog(@TypeOf(arg));

    @compileError("TODO");
}

fn Deref(comptime T: type, comptime names: []const []const u8) type {
    if (names.len == 0) return T;

    // Compiler segfaults when I use std.meta to get this info so we search it manually
    const field = for (@typeInfo(T).Struct.fields) |f| {
        if (std.mem.eql(u8, f.name, names[0])) break f;
    } else @compileError("Unknown field " ++ names[0] ++ " in type " ++ @typeName(T));

    return Deref(field.field_type, names[1..]);
}

fn deref(arg: anytype, comptime names: []const []const u8) Deref(@TypeOf(arg), names) {
    if (names.len == 0) return arg;
    return deref(@field(arg, names[0]), names[1..]);
}

fn EvaluateExpression(comptime expression: Expression, comptime Args: type, comptime Captures: type) type {
    return switch (expression) {
        .arg_deref => |names| Deref(Args, names),
        .capture_deref => |names| Deref(Captures, names),
    };
}

fn evaluateExpression(
    comptime expression: Expression,
    args: anytype,
    captures: anytype,
) EvaluateExpression(expression, @TypeOf(args), @TypeOf(captures)) {
    return switch (expression) {
        .arg_deref => |names| deref(args, names),
        .capture_deref => |names| deref(captures, names),
    };
}

fn AddCapture(comptime Root: type, comptime name: []const u8, comptime Val: type) type {
    var fields = std.meta.fields(Root) ++ [_]std.builtin.Type.StructField{.{
        .name = name,
        .field_type = Val,
        .default_value = null,
        .is_comptime = false,
        .alignment = @alignOf(Val),
    }};

    return @Type(.{ .Struct = .{
        .layout = .Auto,
        .fields = fields,
        .decls = &.{},
        .is_tuple = false,
    } });
}

fn addCapture(root: anytype, comptime name: []const u8, val: anytype) AddCapture(@TypeOf(root), name, @TypeOf(val)) {
    var result = std.mem.zeroInit(AddCapture(@TypeOf(root), name, @TypeOf(val)), root);
    @field(result, name) = val;
    return result;
}

const TemplateType = enum {
    root,
    for_block,
    if_block,
};

fn parseTemplate(
    comptime tokens: []const TemplateToken,
    comptime start: usize,
    comptime template_type: TemplateType,
) ParseResult(usize, []const TemplateItem) {
    comptime {
        var i: usize = start;
        var current_text: []const u8 = "";
        var items: []const TemplateItem = &.{};

        while (i < tokens.len) : (i += 1) {
            switch (tokens[i]) {
                .text => |text| current_text = current_text ++ text,
                .whitespace => |wsp| {
                    if (i != tokens.len - 1 and tokens[i + 1] == .control_block)
                        if (tokens[i + 1].control_block.strip_before)
                            continue;
                    current_text = current_text ++ wsp;
                },
                .control_block => |cb| {
                    if (current_text.len != 0) {
                        items = items ++ [_]TemplateItem{.{ .text = current_text }};
                        current_text = "";
                    }

                    switch (cb.block) {
                        .expression => |expr| items = items ++ [_]TemplateItem{.{ .statement = .{ .expression = expr } }},
                        .if_header => |header| {
                            if (i != tokens.len - 1 and tokens[i + 1] == .whitespace and cb.strip_after) i += 1;
                            const subtemplate = parseTemplate(tokens, i + 1, .if_block);
                            items = items ++ [_]TemplateItem{.{
                                .statement = .{
                                    .@"if" = .{
                                        .subtemplate = subtemplate.item,
                                        .header = header,
                                    },
                                },
                            }};
                            i = subtemplate.new_iter;
                        },
                        .for_header => |header| {
                            if (i != tokens.len - 1 and tokens[i + 1] == .whitespace and cb.strip_after) i += 1;
                            const subtemplate = parseTemplate(tokens, i + 1, .for_block);
                            items = items ++ [_]TemplateItem{.{
                                .statement = .{
                                    .@"for" = .{
                                        .subtemplate = subtemplate.item,
                                        .header = header,
                                    },
                                },
                            }};
                            i = subtemplate.new_iter;
                        },
                        .end_for => if (template_type == .for_block)
                            break
                        else
                            @compileError("Unexpected /for tag"),
                        .end_if => if (template_type == .if_block)
                            break
                        else
                            @compileError("Unexpected /if tag"),
                    }

                    if (i != tokens.len - 1 and tokens[i] == .control_block) {
                        if (tokens[i].control_block.strip_after and tokens[i + 1] == .whitespace) {
                            i += 1;
                        }
                    }
                },
            }
        } else if (template_type != .root) @compileError("End tag not found");

        if (current_text.len != 0) items = items ++ [_]TemplateItem{.{ .text = current_text }};

        return .{
            .new_iter = i,
            .item = items,
        };
    }
}

const TemplateToken = union(enum) {
    text: []const u8,
    whitespace: []const u8,
    control_block: ControlBlock,
};

fn parseTemplateTokens(comptime tokens: ControlTokenIter) []const TemplateToken {
    comptime {
        var iter = tokens;
        var items: []const TemplateToken = &.{};

        while (iter.next()) |token| switch (token) {
            .whitespace => |wsp| items = items ++ [_]TemplateToken{.{ .whitespace = wsp }},
            .text => |text| items = items ++ [_]TemplateToken{.{ .text = text }},
            .open_bracket => {
                const next = iter.next() orelse @compileError("Unexpected end of template");
                if (next == .open_bracket) {
                    items = items ++ [_]TemplateToken{.{ .text = "{" }};
                } else {
                    iter.putBack(next);
                    const result = parseControlBlock(iter);
                    iter = result.new_iter;
                    items = items ++ [_]TemplateToken{.{ .control_block = result.item }};
                }
            },
            .close_bracket => {
                const next = iter.next() orelse @compileError("Unexpected end of template");
                if (next == .close_bracket)
                    items = items ++ [_]TemplateToken{.{ .text = "}" }}
                else
                    @compileError("Unpaired close bracket, did you mean \"}}\"?");
            },
            .period => items = items ++ [_]TemplateToken{.{ .text = "." }},
            .pound => items = items ++ [_]TemplateToken{.{ .text = "#" }},
            .pipe => items = items ++ [_]TemplateToken{.{ .text = "|" }},
            .dollar => items = items ++ [_]TemplateToken{.{ .text = "$" }},
            .slash => items = items ++ [_]TemplateToken{.{ .text = "/" }},
            .equals => items = items ++ [_]TemplateToken{.{ .text = "=" }},
        };

        return items;
    }
}

fn parseExpression(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, Expression) {
    comptime {
        var iter = tokens;

        var expr: Expression = while (iter.next()) |token| switch (token) {
            .whitespace => {},
            .period => {
                const names = parseDeref(iter);
                iter = names.new_iter;
                break .{ .arg_deref = names.item };
            },
            .dollar => {
                const names = parseDeref(iter);
                iter = names.new_iter;
                break .{ .capture_deref = names.item };
            },
            else => @compileError("TODO"),
        };

        return .{
            .new_iter = iter,
            .item = expr,
        };
    }
}

fn parseControlBlock(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, ControlBlock) {
    comptime {
        var iter = tokens;

        const strip_before = if (iter.next()) |first| blk: {
            if (first == .equals) {
                break :blk true;
            }

            iter.putBack(first);
            break :blk false;
        } else @compileError("Unexpected end of template");

        var stmt: ControlBlock.Data = while (iter.next()) |token| switch (token) {
            .equals => @compileError("Unexpected '='"),
            .whitespace => {},
            .pound => {
                const next = iter.next() orelse @compileError("Unexpected end of template");
                if (next != .text) @compileError("Expected keyword following '#' character");
                const text = next.text;
                const keyword = std.meta.stringToEnum(Keyword, text) orelse @compileError("Unknown keyword: " ++ text);

                switch (keyword) {
                    .@"for" => {
                        const result = parseForHeader(iter);
                        iter = result.new_iter;
                        break .{ .for_header = result.item };
                    },
                    .@"if" => {
                        const result = parseIfHeader(iter);
                        iter = result.new_iter;
                        break .{ .if_header = result.item };
                    },

                    //else => @compileError("TODO"),
                }
            },
            .slash => {
                const next = iter.next() orelse @compileError("Unexpected end of template");
                if (next != .text) @compileError("Expected keyword following '/' character");
                const text = next.text;
                const keyword = std.meta.stringToEnum(EndKeyword, text) orelse @compileError("Unknown keyword: " ++ text);

                switch (keyword) {
                    .@"for" => break .{ .end_for = {} },
                    .@"if" => break .{ .end_if = {} },
                }
            },
            .period, .dollar => {
                iter.putBack(token);
                const expr = parseExpression(iter);
                iter = expr.new_iter;
                break .{ .expression = expr.item };
            },
            else => @compileError("TODO"),
        };

        // search for end of statement
        var strip_after: bool = false;
        while (iter.next()) |token| switch (token) {
            .whitespace => {},
            .equals => {
                if (iter.peek()) |t| {
                    if (t == .close_bracket) {
                        strip_after = true;
                        continue;
                    }
                }
                @compileError("Unexpected '='");
            },
            .close_bracket => return .{
                .new_iter = iter,
                .item = .{
                    .block = stmt,
                    .strip_before = strip_before,
                    .strip_after = strip_after,
                },
            },
            else => {
                @compileLog(iter.row);
                @compileError("TODO" ++ @tagName(token));
            },
        };

        @compileError("Unexpected end of template");
    }
}

fn skipWhitespace(comptime tokens: ControlTokenIter) ControlTokenIter {
    comptime {
        var iter = tokens;
        while (iter.peek()) |token| switch (token) {
            .whitespace => _ = iter.next(),
            else => break,
        };

        return iter;
    }
}

fn endControlBlock(comptime tokens: ControlTokenIter) ControlTokenIter {
    comptime {
        var iter = skipWhitespace(tokens);

        const token = iter.next() orelse @compileError("Unexpected end of template");
        if (token != .close_bracket) @compileError("Unexpected token");
        return iter;
    }
}

fn parseForHeader(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, ForHeader) {
    comptime {
        const iterable = parseExpression(tokens);
        var iter = iterable.new_iter;

        iter = skipWhitespace(iter);
        {
            const token = iter.next() orelse @compileError("Unexpected end of template");
            if (token != .pipe) @compileError("Unexpected token");
        }
        {
            const token = iter.next() orelse @compileError("Unexpected end of template");
            if (token != .dollar) @compileError("Unexpected token");
        }
        const capture = blk: {
            const token = iter.next() orelse @compileError("Unexpected end of template");
            if (token != .text) @compileError("Unexpected token");
            break :blk token.text;
        };
        {
            const token = iter.next() orelse @compileError("Unexpected end of template");
            if (token != .pipe) @compileError("Unexpected token");
        }

        return .{
            .new_iter = iter,
            .item = .{
                .iterable = iterable.item,
                .capture = capture,
            },
        };
    }
}

fn parseIfHeader(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, IfHeader) {
    comptime {
        const condition = parseExpression(tokens);
        var iter = condition.new_iter;

        return .{
            .new_iter = iter,
            .item = .{
                .condition = condition.item,
            },
        };
    }
}

fn parseDeref(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, []const []const u8) {
    comptime {
        var iter = tokens;
        var fields: []const []const u8 = &.{};
        var wants = .text;
        while (iter.peek()) |token| {
            switch (token) {
                .whitespace => {},
                .text => |text| {
                    if (wants != .text) @compileError("Unexpected token \"" ++ text ++ "\"");
                    fields = fields ++ [1][]const u8{text};
                    wants = .period;
                },
                .period => {
                    if (wants != .period) @compileError("Unexpected token \".\"");
                    wants = .text;
                },
                else => if (wants == .period) return .{
                    .new_iter = iter,
                    .item = fields,
                } else @compileError("Unexpected token"),
            }
            _ = iter.next();
        }
    }
}

fn ParseResult(comptime It: type, comptime T: type) type {
    return struct {
        new_iter: It,
        item: T,
    };
}

const TemplateItem = union(enum) {
    text: []const u8,
    statement: Statement,
};

const Expression = union(enum) {
    arg_deref: []const []const u8,
    capture_deref: []const []const u8,
};

const For = struct {
    subtemplate: []const TemplateItem,
    header: ForHeader,
};

const ForHeader = struct {
    iterable: Expression,
    capture: []const u8,
};

const If = struct {
    subtemplate: []const TemplateItem,
    header: IfHeader,
};

const IfHeader = struct {
    condition: Expression,
};

const Statement = union(enum) {
    expression: Expression,
    @"for": For,
    @"if": If,
};

const ControlBlock = struct {
    const Data = union(enum) {
        expression: Expression,
        for_header: ForHeader,
        end_for: void,
        if_header: IfHeader,
        end_if: void,
    };
    block: Data,
    strip_before: bool,
    strip_after: bool,
};

const Keyword = enum {
    @"for",
    @"if",
};

const EndKeyword = enum {
    @"for",
    @"if",
};

const ControlToken = union(enum) {
    text: []const u8,
    open_bracket: void,
    close_bracket: void,
    period: void,
    whitespace: []const u8,
    pound: void,
    pipe: void,
    dollar: void,
    slash: void,
    equals: void,
};

const ControlTokenIter = struct {
    start: usize = 0,
    text: []const u8,
    peeked_token: ?ControlToken = null,

    row: usize = 0,

    fn next(self: *ControlTokenIter) ?ControlToken {
        if (self.peeked_token) |token| {
            self.peeked_token = null;
            return token;
        }

        const remaining = self.text[self.start..];
        if (remaining.len == 0) return null;

        const ch = remaining[0];
        self.start += 1;
        switch (ch) {
            '{' => return .{ .open_bracket = {} },
            '}' => return .{ .close_bracket = {} },
            '.' => return .{ .period = {} },
            '#' => return .{ .pound = {} },
            '|' => return .{ .pipe = {} },
            '$' => return .{ .dollar = {} },
            '/' => return .{ .slash = {} },
            '=' => return .{ .equals = {} },
            ' ', '\t', '\n', '\r' => {
                var idx: usize = 0;
                while (idx < remaining.len and std.mem.indexOfScalar(u8, " \t\n\r", remaining[idx]) != null) : (idx += 1) {}
                const newline_count = std.mem.count(u8, remaining[0..idx], "\n");
                self.row += newline_count;

                self.start += idx - 1;
                return .{ .whitespace = remaining[0..idx] };
            },
            else => {
                var idx: usize = 0;
                while (idx < remaining.len and std.mem.indexOfScalar(u8, "{}.#|$/= \t\n\r", remaining[idx]) == null) : (idx += 1) {}

                self.start += idx - 1;
                return .{ .text = remaining[0..idx] };
            },
        }
    }

    fn peek(self: *ControlTokenIter) ?ControlToken {
        const token = self.next();
        self.peeked_token = token;
        return token;
    }

    fn putBack(self: *ControlTokenIter, token: ControlToken) void {
        std.debug.assert(self.peeked_token == null);
        self.peeked_token = token;
    }
};