const std = @import("std");
const Loc = @import("loc.zig").Loc;
const Expr = @import("expr.zig").Expr;
const Pat = @import("pat.zig").Pat;
const Stmt = @import("stmt.zig").Stmt;
const BlockStmt = @import("stmt.zig").BlockStmt;
const VarStmt = @import("stmt.zig").VarStmt;
const VarDeclarator = @import("stmt.zig").VarDeclarator;
const VarKind = @import("stmt.zig").VarKind;
const FnType = @import("expr.zig").FnType;
const ClassElem = @import("expr.zig").ClassElem;
const Parser = @import("../parser.zig").Parser;
const Formatter = @import("../formatter.zig").Formatter;
const ParseError = @import("../parser.zig").ParseError;

pub const DeclType = union(enum) {
    @"fn": FnDecl,
    class: ClassDecl,
    @"var": VarDecl,
    using: UsingDecl,
    import: ImportDecl,
    @"export": ExportDecl,
    export_default: ExportDefault,
};

pub const Decl = struct {
    loc: Loc,
    data: DeclType,

    pub fn fmt(self: @This(), f: *Formatter) void {
        switch (self.data) {
            .@"fn" => |d| d.fmt(f),
            .class => |c| c.fmt(f),
            .@"var" => |v| {
                f.write(@tagName(v.kind));
                f.writeSpace();
                for (v.decls, 0..) |d, i| {
                    if (i > 0) f.write(", ");
                    d.fmt(f);
                }
            },
            .using => |u| {
                f.write("using ");
                for (u.decls, 0..) |d, i| {
                    if (i > 0) f.write(", ");
                    d.fmt(f);
                }
            },
            .import => |i| i.fmt(f),
            .@"export" => |e| e.fmt(f),
            .export_default => |e| e.fmt(f),
        }
    }

    pub fn parse(p: *Parser) ParseError!Decl {
        switch (p.peek()) {
            .@"var" => {
                const s = try VarStmt.parse(p);
                return Decl{ .loc = s.loc, .data = .{ .@"var" = .{ .loc = s.loc, .kind = s.kind, .decls = s.decls } } };
            },
            else => return ParseError.ExpectedStatement,
        }
    }
};

pub const FnDecl = struct {
    loc: Loc,
    fn_type: FnType,
    id: ?[]const u8,
    params: []const Pat,
    body: BlockStmt,

    pub fn parse(p: *Parser, fn_type_arg: FnType, mode: enum { stmt, expr }) ParseError!?Stmt {
        const start = p.tok().loc.start;
        var fn_type = fn_type_arg;

        if (fn_type == .async) {
            _ = p.advanceTok();
        } else if (fn_type == .async_generator) {
            _ = p.advanceTok();
            _ = p.advanceTok();
        }

        _ = p.advanceTok();

        if (fn_type == .normal and p.peek() == .@"*") {
            p.advance();
            fn_type = .generator;
        }

        var id: ?[]const u8 = null;
        if (p.peek().isIdentifier()) {
            const id_tok = p.advanceTok();
            id = p.tokenSlice(id_tok);
        }

        const params = try p.parseParams();
        const body = try p.parseFunctionBody();

        const end = body.loc.end;

        if (mode == .stmt) {
            return Stmt{
                .loc = .{ .start = start, .end = end },
                .data = .{ .function = .{ .loc = .{ .start = start, .end = end }, .fn_type = fn_type, .id = id, .params = params, .body = body } },
            };
        }
        return null;
    }

    pub fn fmt(self: @This(), f: *Formatter) void {
        f.write("function");
        if (self.fn_type == .generator or self.fn_type == .async_generator) {
            f.write("*");
        }
        if (self.fn_type == .async or self.fn_type == .async_generator) {
            f.write(" async");
        }
        if (self.id) |id| {
            f.writeSpace();
            f.write(id);
        }
        f.writeChar('(');
        for (self.params, 0..) |p, i| {
            if (i > 0) f.write(", ");
            p.fmt(f);
        }
        f.write(") ");
        self.body.fmt(f);
    }
};

pub const ClassDecl = struct {
    loc: Loc,
    id: ?[]const u8,
    super_class: ?*const Expr,
    body: []const ClassElem,

    pub fn parse(p: *Parser, mode: enum { stmt, expr }) ParseError!?Stmt {
        const start = p.advanceTok().loc.start;

        var id: ?[]const u8 = null;
        if (p.peek() == .identifier) {
            const id_tok = p.advanceTok();
            id = p.tokenSlice(id_tok);
        }

        var super_class: ?*const Expr = null;
        if (p.matchKeyword(.extends)) {
            super_class = try Expr.parse(p, 0);
        }

        try p.expect(.lbrace);
        var elems: [128]ClassElem = undefined;
        var elem_count: usize = 0;

        while (p.peek() != .rbrace and p.peek() != .eof) {
            const elem = try ClassElem.parse(p);
            elems[elem_count] = elem;
            elem_count += 1;
            if (elem_count >= elems.len) return ParseError.TooManyTokens;
        }

        const end_tok = try p.expectTok(.rbrace);

        if (mode == .stmt) {
            return Stmt{
                .loc = .{ .start = start, .end = end_tok.loc.end },
                .data = .{ .class = .{ .loc = .{ .start = start, .end = end_tok.loc.end }, .id = id, .super_class = super_class, .body = elems[0..elem_count] } },
            };
        }
        return null;
    }

    pub fn fmt(self: @This(), f: *Formatter) void {
        f.write("class");
        if (self.id) |id| {
            f.writeSpace();
            f.write(id);
        }
        if (self.super_class) |sc| {
            f.write(" extends ");
            sc.*.fmt(f);
        }
        f.write(" {");
        f.newline();
        f.indent_level += 1;
        for (self.body) |elem| {
            f.writeIndent();
            elem.fmt(f);
            f.newline();
        }
        f.indent_level -= 1;
        f.writeIndent();
        f.writeChar('}');
    }
};

pub const VarDecl = struct {
    loc: Loc,
    kind: VarKind,
    decls: []const VarDeclarator,
};

pub const UsingDecl = struct {
    loc: Loc,
    is_await: bool,
    decls: []const VarDeclarator,
};

pub const ImportDecl = struct {
    loc: Loc,
    specifiers: []const ImportSpecifier,
    source: []const u8,
    attributes: []const ImportAttribute,

    pub fn parse(p: *Parser) ParseError!ImportDecl {
        const start = p.advanceTok().loc.start;
        var specifiers: [32]ImportSpecifier = undefined;
        var spec_count: usize = 0;

        if (p.peek() != .string and p.peek() != .lparen) {
            if (p.peek() == .identifier or p.peek().isKeyword()) {
                const local_tok = p.advanceTok();
                const local = p.tokenSlice(local_tok);
                specifiers[spec_count] = .{ .default = local };
                spec_count += 1;
                if (p.match(.comma)) {}
            }

            if (p.match(.@"*")) {
                try p.expect(.identifier);
                const ns_tok = try p.expectTok(.identifier);
                const ns = p.tokenSlice(ns_tok);
                specifiers[spec_count] = .{ .namespace = ns };
                spec_count += 1;
            } else if (p.peek() == .lbrace) {
                p.advance();
                if (p.peek() != .rbrace) {
                    while (true) {
                        const imported_tok = try p.expectTok(.identifier);
                        const imported = p.tokenSlice(imported_tok);
                        var local = imported;
                        if (p.match(.identifier)) {
                            local = p.tokenSlice(p.tokens[p.pos - 1]);
                        }
                        specifiers[spec_count] = .{ .named = .{ .imported = imported, .local = local } };
                        spec_count += 1;
                        if (!p.match(.comma)) break;
                    }
                }
                try p.expect(.rbrace);
            }
        }

        var source: []const u8 = "";
        var attrs: [8]ImportAttribute = undefined;
        var attr_count: usize = 0;

        if (p.matchKeyword(.from)) {
            const str_tok = try p.expectTok(.string);
            source = p.tokenSlice(str_tok);
        } else if (p.peek() == .string) {
            const str_tok = p.advanceTok();
            source = p.tokenSlice(str_tok);
        }

        if (p.matchKeyword(.with)) {
            try p.expect(.lbrace);
            while (p.peek() != .rbrace and p.peek() != .eof) {
                const key_tok = p.advanceTok();
                const key = p.tokenSlice(key_tok);
                try p.expect(.@":");
                const val_tok = try p.expectTok(.string);
                const val = p.tokenSlice(val_tok);
                attrs[attr_count] = .{ .key = key, .value = val };
                attr_count += 1;
                _ = p.match(.comma);
            }
            try p.expect(.rbrace);
        }

        try p.expect(.semicolon);
        return .{ .loc = .{ .start = start, .end = p.tok().loc.start }, .specifiers = specifiers[0..spec_count], .source = source, .attributes = attrs[0..attr_count] };
    }

    pub fn fmt(self: @This(), f: *Formatter) void {
        f.write("import ");
        if (self.specifiers.len > 0) {
            for (self.specifiers, 0..) |spec, j| {
                if (j > 0) f.write(", ");
                spec.fmt(f);
            }
            f.write(" from ");
        }
        f.writeChar('"');
        f.write(self.source);
        f.writeChar('"');
        for (self.attributes) |a| {
            f.write(" /* ");
            a.fmt(f);
            f.write(" */");
        }
        f.writeChar(';');
    }
};

pub const ImportSpecifier = union(enum) {
    default: []const u8,
    namespace: []const u8,
    named: struct { imported: []const u8, local: []const u8 },
};

pub const ImportAttribute = struct {
    key: []const u8,
    value: []const u8,

    pub fn fmt(self: @This(), f: *Formatter) void {
        f.write(self.key);
        f.write(": \"");
        f.write(self.value);
        f.writeChar('"');
    }
};

pub const ExportDecl = struct {
    loc: Loc,
    declaration: ?*const Decl,
    specifiers: ?[]const ExportSpecifier,
    source: ?[]const u8,

    pub fn parse(p: *Parser) ParseError!ExportDecl {
        const start = p.advanceTok().loc.start;

        if (p.matchKeyword(.default)) {
            if (p.peek() == .function or p.peek() == .async or p.peek() == .class) {
                const decl = try Stmt.parse(p);
                p.semicolon();
                return .{ .loc = .{ .start = start, .end = decl.loc.end }, .declaration = null, .specifiers = null, .source = null };
            }
            const expr = try Expr.parse(p, 0);
            p.semicolon();
            return .{ .loc = .{ .start = start, .end = expr.loc.end }, .declaration = null, .specifiers = null, .source = null };
        }

        if (p.peek() == .@"*") {
            p.advance();
            try p.expect(.identifier);
            const source_tok = try p.expectTok(.string);
            const src = p.tokenSlice(source_tok);
            try p.expect(.semicolon);
            return .{ .loc = .{ .start = start, .end = p.tok().loc.start }, .declaration = null, .specifiers = null, .source = src };
        }

        if (p.peek() == .lbrace) {
            p.advance();
            var specs: [32]ExportSpecifier = undefined;
            var spec_count: usize = 0;
            while (p.peek() != .rbrace and p.peek() != .eof) {
                const exported_tok = p.advanceTok();
                const exported = p.tokenSlice(exported_tok);
                var local = exported;
                if (p.match(.identifier)) {
                    local = p.tokenSlice(p.tokens[p.pos - 1]);
                }
                specs[spec_count] = .{ .loc = .{ .start = exported_tok.loc.start, .end = p.tok().loc.start }, .exported = exported, .local = local };
                spec_count += 1;
                if (!p.match(.comma)) break;
            }
            try p.expect(.rbrace);
            var source: ?[]const u8 = null;
            if (p.matchKeyword(.from)) {
                const src_tok = try p.expectTok(.string);
                source = p.tokenSlice(src_tok);
            }
            try p.expect(.semicolon);
            return .{ .loc = .{ .start = start, .end = p.tok().loc.start }, .declaration = null, .specifiers = specs[0..spec_count], .source = source };
        }

        const decl = try Stmt.parse(p);
        p.semicolon();
        return .{ .loc = .{ .start = start, .end = decl.loc.end }, .declaration = null, .specifiers = null, .source = null };
    }

    pub fn fmt(self: @This(), f: *Formatter) void {
        f.write("export ");
        if (self.declaration) |d| {
            d.*.fmt(f);
        } else if (self.specifiers) |s| {
            f.writeChar('{');
            for (s, 0..) |sp, i| {
                if (i > 0) f.write(", ");
                sp.fmt(f);
            }
            f.writeChar('}');
            if (self.source) |src| {
                f.write(" from \"");
                f.write(src);
                f.writeChar('"');
            }
            f.writeChar(';');
        } else if (self.source) |src| {
            f.write("* from \"");
            f.write(src);
            f.writeChar("\";");
        }
    }
};

pub const ExportSpecifier = struct {
    loc: Loc,
    exported: []const u8,
    local: []const u8,

    pub fn fmt(self: @This(), f: *Formatter) void {
        f.write(self.local);
        if (!std.mem.eql(u8, self.local, self.exported)) {
            f.write(" as ");
            f.write(self.exported);
        }
    }
};

pub const ExportDefault = struct {
    loc: Loc,
    declaration: *const Decl,

    pub fn fmt(self: @This(), f: *Formatter) void {
        f.write("export default ");
        self.declaration.*.fmt(f);
    }
};