const std = @import("std");
const Io = std.Io;
const assert = std.debug.assert;

const Color = @import("./color.zig");

pub const Attribute = struct {
    name: []const u8,
    value: []const u8,

    pub const List = struct {
        items: []const Attribute,

        pub fn get(self: List, name: []const u8) ?[]const u8 { 
            for (self.items) |attr| {
                if (std.mem.eql(u8, attr.name, name)) {
                    return attr.value;
                }
            }
            return null;
        }

        pub fn getDupe(self: List, gpa: std.mem.Allocator, name: []const u8) !?[]u8 { 
            if (self.get(name)) |value| {
                return try gpa.dupe(u8, value);
            }
            return null;
        }

        pub fn getNumber(self: List, T: type, name: []const u8) !?T {
            if (self.get(name)) |value| {
                if (@typeInfo(T) == .int) {
                    return try std.fmt.parseInt(T, value, 10);
                } else if (@typeInfo(T) == .float) {
                    return try std.fmt.parseFloat(T, value);
                }
            }
            return null;
        }

        pub fn getBool(self: List, name: []const u8, true_value: []const u8, false_value: []const u8) !?bool {
            if (self.get(name)) |value| {
                if (std.mem.eql(u8, value, true_value)) {
                    return true;
                } else if (std.mem.eql(u8, value, false_value)) {
                    return false;
                } else {
                    return error.InvalidBoolean;
                }
            }
            return null;
        }

        pub fn getEnum(self: List, T: type, name: []const u8, map: std.StaticStringMap(T)) !?T {
            if (self.get(name)) |value| {
                return map.get(value) orelse return error.InvalidEnumValue;
            }
            return null;
        }

        pub fn getColor(self: List, name: []const u8, hash_required: bool) !?Color {
            if (self.get(name)) |value| {
                return try Color.parse(value, hash_required);
            }
            return null;
        }

        pub fn format(self: List, writer: *Io.Writer) Io.Writer.Error!void {
            if (self.items.len > 0) {
                try writer.writeAll("{ ");
                for (self.items, 0..) |attribute, i| {
                    if (i > 0) {
                        try writer.writeAll(", ");
                    }
                    try writer.print("{f}", .{attribute});
                }
                try writer.writeAll(" }");
            } else {
                try writer.writeAll("{ }");
            }
        }
    };

    pub fn format(self: *const Attribute, writer: *Io.Writer) Io.Writer.Error!void {
        try writer.print("{s}{{ .name='{s}', .value='{s}' }}", .{ @typeName(Attribute), self.name, self.value });
    }

    pub fn formatSlice(data: []const Attribute, writer: *Io.Writer) Io.Writer.Error!void {
        if (data.len > 0) {
            try writer.writeAll("{ ");
            for (data, 0..) |attribute, i| {
                if (i > 0) {
                    try writer.writeAll(", ");
                }
                try writer.print("{f}", .{attribute});
            }
            try writer.writeAll(" }");
        } else {
            try writer.writeAll("{ }");
        }
    }

    fn altSlice(data: []const Attribute) std.fmt.Alt(Attribute.List, Attribute.List.format) {
        return .{ .data = data };
    }
};

pub const Tag = struct {
    name: []const u8,
    attributes: Attribute.List
};

pub const Lexer = struct {
    pub const Buffers = struct {
        scratch: std.heap.ArenaAllocator,
        text: std.ArrayList(u8),

        pub fn init(allocator: std.mem.Allocator) Buffers {
            return Buffers{
                .scratch = std.heap.ArenaAllocator.init(allocator),
                .text = .empty
            };
        }

        pub fn clear(self: *Buffers) void {
            self.text.clearRetainingCapacity();
            _ = self.scratch.reset(.retain_capacity);
        }

        pub fn deinit(self: *Buffers) void {
            const allocator = self.scratch.child_allocator;
            self.scratch.deinit();
            self.text.deinit(allocator);
        }
    };

    pub const Token = union(enum) {
        start_tag: Tag,
        end_tag: []const u8,
        text: []const u8,

        pub fn isStartTag(self: Token, name: []const u8) bool {
            if (self == .start_tag) {
                return std.mem.eql(u8, self.start_tag.name, name);
            }
            return false;
        }

        pub fn isEndTag(self: Token, name: []const u8) bool {
            if (self == .end_tag) {
                return std.mem.eql(u8, self.end_tag, name);
            }
            return false;
        }
    };

    pub const StepResult = struct {
        token: ?Token,
        self_closing: bool = false,
    };

    pub const TestingContext = struct {
        io_reader: Io.Reader,
        buffers: Buffers,
        lexer: Lexer,

        pub fn init(self: *TestingContext, allocator: std.mem.Allocator, body: []const u8) void {
            self.* = TestingContext{
                .lexer = undefined,
                .io_reader = Io.Reader.fixed(body),
                .buffers = Buffers.init(allocator)
            };
            self.lexer = Lexer.init(&self.io_reader, &self.buffers);
        }

        pub fn deinit(self: *TestingContext) void {
            self.buffers.deinit();
        }
    };

    io_reader: *Io.Reader,
    buffers: *Buffers,

    peeked_value: ?Token,
    cursor: usize,
    queued_end_tag: ?[]const u8,

    pub fn init(reader: *Io.Reader, buffers: *Buffers) Lexer {
        buffers.clear();

        return Lexer{
            .io_reader = reader,
            .buffers = buffers,
            .cursor = 0,
            .queued_end_tag = null,
            .peeked_value = null
        };
    }

    fn step(self: *Lexer) !StepResult {
        _ = self.buffers.scratch.reset(.retain_capacity);

        if (try self.peekByte() == '<') {
            self.tossByte();

            if (try self.peekByte() == '/') {
                // End tag
                self.tossByte();

                const name = try self.parseName();
                try self.skipWhiteSpace();

                if (!std.mem.eql(u8, try self.takeBytes(1), ">")) {
                    return error.InvalidEndTag;
                }

                const token = Token{ .end_tag = name };
                return .{ .token = token };

            } else if (try self.peekByte() == '?') {
                // Prolog tag
                self.tossByte();
                if (!std.mem.eql(u8, try self.takeBytes(4), "xml ")) {
                    return error.InvalidPrologTag;
                }

                const attributes = try self.parseAttributes();
                try self.skipWhiteSpace();

                if (!std.mem.eql(u8, try self.takeBytes(2), "?>")) {
                    return error.MissingPrologEnd;
                }

                const version = attributes.get("version") orelse return error.InvalidProlog;
                if (!std.mem.eql(u8, version, "1.0")) {
                    return error.InvalidPrologVersion;
                }
                const encoding = attributes.get("encoding") orelse return error.InvalidProlog;
                if (!std.mem.eql(u8, encoding, "UTF-8")) {
                    return error.InvalidPrologEncoding;
                }

                return .{ .token = null };

            } else {
                // Start tag
                const name = try self.parseName();
                const attributes = try self.parseAttributes();
                try self.skipWhiteSpace();

                const token = Token{
                    .start_tag = .{
                        .name = name,
                        .attributes = attributes
                    }
                };

                var self_closing = false;
                if (std.mem.eql(u8, try self.peekBytes(1), ">")) {
                    self.tossBytes(1);

                } else if (std.mem.eql(u8, try self.peekBytes(2), "/>")) {
                    self.tossBytes(2);

                    self_closing = true;

                } else {
                    return error.UnfinishedStartTag;
                }

                return .{
                    .token = token,
                    .self_closing = self_closing
                };
            }

        } else {
            try self.skipWhiteSpace();

            const text_start = self.cursor;
            while (try self.peekByte() != '<') {
                self.tossByte();
            }
            var text: []const u8 = self.buffers.text.items[text_start..self.cursor];
            text = std.mem.trimEnd(u8, text, &std.ascii.whitespace);

            var token: ?Token = null;
            if (text.len > 0) {
                token = Token{ .text = text };
            }
            return .{ .token = token };
        }
    }

    pub fn next(self: *Lexer) !?Token {
        if (self.peeked_value) |value| {
            self.peeked_value = null;
            return value;
        }

        if (self.queued_end_tag) |name| {
            self.queued_end_tag = null;
            return Token{
                .end_tag = name
            };
        }

        while (true) {
            if (self.buffers.text.items.len == 0) {
                self.readIntoTextBuffer() catch |e| switch (e) {
                    error.EndOfStream => break,
                    else => return e
                };
            }

            const saved_cursor = self.cursor;
            const result = self.step() catch |e| switch(e) {
                error.EndOfTextBuffer => {
                    self.cursor = saved_cursor;

                    const unused_capacity = self.buffers.text.capacity - self.buffers.text.items.len;
                    if (unused_capacity == 0 and self.cursor > 0) {
                        self.rebaseBuffer();
                    } else {
                        self.readIntoTextBuffer() catch |read_err| switch (read_err) {
                            error.EndOfStream => break,
                            else => return read_err
                        };
                    }

                    continue;
                },
                else => return e
            };

            if (result.token) |token| {
                if (token == .start_tag and result.self_closing) {
                    self.queued_end_tag = token.start_tag.name;
                }
                return token;
            }
        }

        return null;
    }

    pub fn nextExpectEndTag(self: *Lexer, name: []const u8) !void {
        const value = try self.next() orelse return error.MissingEndTag;
        if (!value.isEndTag(name)) return error.MissingEndTag;
    }

    pub fn nextExpectStartTag(self: *Lexer, name: []const u8) !Attribute.List {
        const value = try self.next() orelse return error.MissingStartTag;
        if (!value.isStartTag(name)) return error.MissingStartTag;
        return value.start_tag.attributes;
    }

    pub fn nextExpectText(self: *Lexer) ![]const u8 {
        const value = try self.next() orelse return error.MissingTextTag;
        if (value != .text) return error.MissingTextTag;
        return value.text;
    }

    pub fn skipUntilMatchingEndTag(self: *Lexer, name: ?[]const u8) !void {
        var depth: usize = 0;
        while (true) {
            const value = try self.next() orelse return error.MissingEndTag;

            if (depth == 0 and value == .end_tag) {
                if (name != null and !std.mem.eql(u8, value.end_tag, name.?)) {
                    return error.MismatchedEndTag;
                }
                break;
            }

            if (value == .start_tag) {
                depth += 1;
            } else if (value == .end_tag) {
                depth -= 1;
            }
        }
    }

    pub fn peek(self: *Lexer) !?Token {
        if (try self.next()) |value| {
            self.peeked_value = value;
            return value;
        }

        return null;
    }

    fn readIntoTextBuffer(self: *Lexer) !void {
        const gpa = self.buffers.scratch.child_allocator;
        const text = &self.buffers.text;
        try text.ensureUnusedCapacity(gpa, 1);

        var writer = Io.Writer.fixed(text.allocatedSlice());
        writer.end = text.items.len;

        _ = self.io_reader.stream(&writer, .limited(text.capacity - text.items.len)) catch |e| switch (e) {
            error.WriteFailed => unreachable,
            else => |ee| return ee
        };
        text.items.len = writer.end;
    }

    fn rebaseBuffer(self: *Lexer) void {
        if (self.cursor == 0) {
            return;
        }

        const text = &self.buffers.text;
        @memmove(
            text.items[0..(text.items.len - self.cursor)],
            text.items[self.cursor..]
        );

        text.items.len -= self.cursor;
        self.cursor = 0;
    }

    fn isNameStartChar(c: u8) bool {
        return c == ':' or c == '_' or std.ascii.isAlphabetic(c);
    }

    fn isNameChar(c: u8) bool {
        return isNameStartChar(c) or c == '-' or c == '.' or ('0' <= c and c <= '9');
    }

    fn hasBytes(self: *Lexer, n: usize) bool {
        const text = self.buffers.text.items;
        return self.cursor + n <= text.len;
    }

    fn peekBytes(self: *Lexer, n: usize) ![]const u8 {
        if (self.hasBytes(n)) {
            const text = self.buffers.text.items;
            return text[self.cursor..][0..n];
        }
        return error.EndOfTextBuffer;
    }

    fn tossBytes(self: *Lexer, n: usize) void {
        assert(self.hasBytes(n));
        self.cursor += n;
    }

    fn takeBytes(self: *Lexer, n: usize) ![]const u8 {
        const result = try self.peekBytes(n);
        self.tossBytes(n);
        return result;
    }

    fn peekByte(self: *Lexer) !u8 {
        return (try self.peekBytes(1))[0];
    }

    fn tossByte(self: *Lexer) void {
        self.tossBytes(1);
    }

    fn takeByte(self: *Lexer) !u8 {
        return (try self.takeBytes(1))[0];
    }

    fn parseName(self: *Lexer) ![]const u8 {
        const name_start = self.cursor;

        if (isNameStartChar(try self.peekByte())) {
            self.tossByte();

            while (isNameChar(try self.peekByte())) {
                self.tossByte();
            }
        }

        return self.buffers.text.items[name_start..self.cursor];
    }

    fn skipWhiteSpace(self: *Lexer) !void {
        while (std.ascii.isWhitespace(try self.peekByte())) {
            self.tossByte();
        }
    }

    fn parseAttributeValue(self: *Lexer) ![]const u8 {
        const quote = try self.takeByte();
        if (quote != '"' and quote  != '\'') {
            return error.InvalidAttributeValue;
        }

        const value_start: usize = self.cursor;
        var value_len: usize = 0;

        while (true) {
            const c = try self.takeByte();
            if (c == '<' or c == '&') {
                return error.InvalidAttributeValue;
            }
            if (c == quote) {
                break;
            }
            value_len += 1;
        }

        return self.buffers.text.items[value_start..][0..value_len];
    }

    fn parseAttributes(self: *Lexer) !Attribute.List {
        const arena = self.buffers.scratch.allocator();
        var attributes: std.ArrayList(Attribute) = .empty;

        while (true) {
            try self.skipWhiteSpace();
            const name = try self.parseName();
            if (name.len == 0) {
                break;
            }

            try self.skipWhiteSpace();
            if (try self.takeByte() != '=') {
                return error.MissingAttributeEquals;
            }
            try self.skipWhiteSpace();
            const value = try self.parseAttributeValue();

            const list = Attribute.List{ .items = attributes.items };
            if (list.get(name) != null) {
                return error.DuplicateAttribute;
            }

            try attributes.append(arena, Attribute{
                .name = name,
                .value = value
            });
        }

        return Attribute.List{
            .items = attributes.items
        };
    }

    test "self closing tag" {
        const allocator = std.testing.allocator;
        var ctx: TestingContext = undefined;
        ctx.init(allocator,
            \\ <hello />
        );
        defer ctx.deinit();

        try std.testing.expect((try ctx.lexer.next()).?.isStartTag("hello"));
        try std.testing.expect((try ctx.lexer.next()).?.isEndTag("hello"));
        try std.testing.expect((try ctx.lexer.next()) == null);
    }

    test "tag" {
        const allocator = std.testing.allocator;
        var ctx: TestingContext = undefined;
        ctx.init(allocator,
            \\ <hello></hello>
        );
        defer ctx.deinit();

        try std.testing.expect((try ctx.lexer.next()).?.isStartTag("hello"));
        try std.testing.expect((try ctx.lexer.next()).?.isEndTag("hello"));
        try std.testing.expect((try ctx.lexer.next()) == null);
    }

    test "tag with prolog" {
        const allocator = std.testing.allocator;
        var ctx: TestingContext = undefined;
        ctx.init(allocator,
            \\ <?xml version="1.0" encoding="UTF-8"?>
            \\ <hello></hello>
        );
        defer ctx.deinit();

        try std.testing.expect((try ctx.lexer.next()).?.isStartTag("hello"));
        try std.testing.expect((try ctx.lexer.next()).?.isEndTag("hello"));
        try std.testing.expect((try ctx.lexer.next()) == null);
    }

    test "text content" {
        const allocator = std.testing.allocator;
        var ctx: TestingContext = undefined;
        ctx.init(allocator,
            \\ <hello> Hello World </hello>
        );
        defer ctx.deinit();

        try std.testing.expect((try ctx.lexer.next()).?.isStartTag("hello"));
        try std.testing.expectEqualStrings("Hello World", (try ctx.lexer.next()).?.text);
        try std.testing.expect((try ctx.lexer.next()).?.isEndTag("hello"));
        try std.testing.expect((try ctx.lexer.next()) == null);
    }

    test "attributes" {
        const allocator = std.testing.allocator;
        var ctx: TestingContext = undefined;
        ctx.init(allocator,
            \\ <hello a='1' b='2'/>
        );
        defer ctx.deinit();

        const token = try ctx.lexer.next();
        const attrs = token.?.start_tag.attributes;
        try std.testing.expectEqualStrings("1", attrs.get("a").?);
        try std.testing.expectEqualStrings("2", attrs.get("b").?);
    }
};

// TODO: The API for this is easy to misuse.
// Design a better API for using Reader
// As a compromise `assert` was used to guard against some of the ways this can be misused
pub const TagParser = struct {
    lexer: *Lexer,

    begin_called: bool = false,
    finish_called: bool = false,

    pub const Node = union(enum) {
        tag: Tag,
        text: []const u8,

        pub fn isTag(self: Node, name: []const u8) bool {
            if (self == .tag) {
                return std.mem.eql(u8, self.tag.name, name);
            }
            return false;
        }
    };

    pub fn init(lexer: *Lexer) TagParser {
        return TagParser{
            .lexer = lexer,
        };
    }

    pub fn begin(self: *TagParser, name: []const u8) !Attribute.List {
        assert(!self.begin_called);
        self.begin_called = true;

        return try self.lexer.nextExpectStartTag(name);
    }

    pub fn finish(self: *TagParser, name: []const u8) !void {
        assert(self.begin_called);
        assert(!self.finish_called);
        self.finish_called = true;

        try self.lexer.skipUntilMatchingEndTag(name);
    }

    pub fn next(self: *TagParser) !?Node {
        assert(self.begin_called);
        assert(!self.finish_called);

        const value = try self.lexer.peek() orelse return error.MissingEndTag;
        if (value == .end_tag) {
            return null;
        }

        return switch (value) {
            .text => |text| Node{ .text = text },
            .start_tag => |start_tag| Node{ .tag = start_tag },
            .end_tag => unreachable,
        };
    }

    pub fn skip(self: *TagParser) !void {
        assert(self.begin_called);
        assert(!self.finish_called);

        const value = try self.lexer.next() orelse return error.MissingNode; 
        if (value == .end_tag) {
            return error.UnexpectedEndTag;
        } else if (value == .start_tag) {
            // TODO: Make this configurable
            var name_buffer: [64]u8 = undefined;
            var name: std.ArrayList(u8) = .initBuffer(&name_buffer);
            try name.appendSliceBounded(value.start_tag.name);

            try self.lexer.skipUntilMatchingEndTag(name.items);
        }
    }
};