const std = @import("std");
const ink = @import("ink");
const Ast = ink.Ast;
const AstGen = ink.AstGen;
const Ir = ink.Ir;
const Story = ink.Story;
const Compilation = ink.Compilation;
const fatal = std.process.fatal;

var stdin_buffer: [4096]u8 align(std.heap.page_size_min) = undefined;
var stdout_buffer: [4096]u8 align(std.heap.page_size_min) = undefined;
var stderr_buffer: [4096]u8 align(std.heap.page_size_min) = undefined;
var debug_allocator: std.heap.DebugAllocator(.{}) = .init;

const ArgsIterator = struct {
    args: []const []const u8,
    index: usize = 0,

    pub fn next(iter: *ArgsIterator) ?[]const u8 {
        if (iter.index >= iter.args.len) return null;
        defer iter.index += 1;
        return iter.args[iter.index];
    }

    pub fn nextOrFatal(iter: *ArgsIterator) []const u8 {
        if (iter.index >= iter.args.len) {
            fatal("expected parameter after {s}", .{iter.args[iter.index - 1]});
        }
        defer iter.index += 1;
        return iter.args[iter.index];
    }
};

const FileExtension = enum {
    /// Ink source file.
    ink,
    /// Pre-compiled source file.
    inkc,
};

pub fn main(init: std.process.Init) !void {
    const gpa = init.gpa;
    const io = init.io;
    const arena = init.arena.allocator();
    const args = try init.minimal.args.toSlice(arena);
    if (args.len < 2) {
        fatal("Not enough arguments!", .{});
    }
    return mainArgs(io, gpa, arena, args);
}

fn mainArgs(
    io: std.Io,
    gpa: std.mem.Allocator,
    arena: std.mem.Allocator,
    args_list: []const [:0]const u8,
) !void {
    var source_path: ?[]const u8 = null;
    var output_path: ?[]const u8 = null;
    var compile_only: bool = false;
    var dump_ast: bool = false;
    var dump_ir: bool = false;
    var dump_story: bool = false;
    var dump_trace: bool = false;
    var use_stdin: bool = false;
    var use_color: bool = false;

    var args_iter: ArgsIterator = .{
        .args = args_list[1..],
    };

    while (args_iter.next()) |arg| {
        if (std.mem.startsWith(u8, arg, "-")) {
            if (std.mem.eql(u8, arg, "--stdin")) {
                use_stdin = true;
            } else if (std.mem.eql(u8, arg, "--compile-only")) {
                compile_only = true;
            } else if (std.mem.eql(u8, arg, "--dump-ast")) {
                dump_ast = true;
            } else if (std.mem.eql(u8, arg, "--dump-ir")) {
                dump_ir = true;
            } else if (std.mem.eql(u8, arg, "--dump-story")) {
                dump_story = true;
            } else if (std.mem.eql(u8, arg, "--dump-trace")) {
                dump_trace = true;
            } else if (std.mem.eql(u8, arg, "--use-color")) {
                use_color = true;
            } else if (std.mem.eql(u8, arg, "--output")) {
                const next_arg = args_iter.nextOrFatal();
                output_path = next_arg;
            } else {
                fatal("invalid parameter: '{s}'", .{arg});
            }
        } else if (source_path == null) {
            source_path = arg;
        } else {
            fatal("invalid positional argument: '{s}'", .{arg});
        }
    }

    const filename = source_path orelse "<STDIN>";
    const source_bytes: [:0]const u8 = s: {
        var f = if (source_path) |p| file: {
            break :file std.Io.Dir.cwd().openFile(io, p, .{}) catch |err| {
                fatal("unable to open file '{s}': {s}", .{ p, @errorName(err) });
            };
        } else std.Io.File.stdin();
        defer if (source_path != null) f.close(io);

        var file_reader: std.Io.File.Reader = f.reader(io, &stdin_buffer);
        break :s ink.readSourceFileToEndAlloc(arena, &file_reader) catch |err| {
            fatal("unable to load file '{s}': {s}", .{ filename, @errorName(err) });
        };
    };

    const stderr = std.Io.File.stderr();
    var stderr_writer = stderr.writer(io, &stderr_buffer);
    const io_w = &stderr_writer.interface;
    const stack_size = 128;

    const file_ext: FileExtension = f: {
        const bytes = std.fs.path.extension(filename);
        if (std.mem.eql(u8, bytes, ".ink")) break :f .ink;
        if (std.mem.eql(u8, bytes, ".inkc")) break :f .inkc;
        return error.InvalidFileExtension;
    };
    switch (file_ext) {
        .ink => {
            var tree = try Ast.parse(gpa, arena, source_bytes, filename, 0);
            defer tree.deinit(gpa);

            var ir = try AstGen.generate(gpa, &tree);
            defer ir.deinit(gpa);

            var errors: std.ArrayListUnmanaged(Compilation.Error) = .empty;
            defer errors.deinit(gpa);

            if (dump_ast) {
                try io_w.writeAll("=== AST ===\n");
                try tree.render(gpa, io_w, .{
                    .use_color = use_color,
                });
                try io_w.flush();
            }
            if (dump_ir) {
                try io_w.writeAll("=== IR ===\n");
                try ir.render(io_w);
                try io_w.flush();
            }

            var cu = Compilation.build(gpa, tree, ir, &errors) catch |err| switch (err) {
                else => |e| return e,
            };
            defer cu.deinit();

            if (cu.hasCompileErrors()) {
                for (cu.errors.items) |err| {
                    try cu.renderError(io_w, err);
                }
            } else if (output_path) |_| {
                return error.NotImplemented;
            } else {
                var story: Story = try .fromCompilation(gpa, &cu, .{
                    .stack_size = stack_size,
                });
                defer story.deinit();
                if (dump_story) {
                    try story.dump(io_w);
                }
                if (dump_trace) {
                    story.dump_writer = io_w;
                }
                return if (!compile_only) run(io, gpa, &story);
            }
        },
        .inkc => {
            var story: Story = try .fromCachedCompilation(gpa, source_bytes, .{
                .stack_size = stack_size,
            });
            defer story.deinit();
            return if (!compile_only) run(io, gpa, &story);
        },
    }
}

fn run(io: std.Io, _: std.mem.Allocator, story: *Story) !void {
    const stdin = std.Io.File.stdin();
    var stdin_reader = stdin.reader(io, &stdin_buffer);
    const stdout = std.Io.File.stdin();
    var stdout_writer = stdout.writer(io, &stdout_buffer);
    const reader = &stdin_reader.interface;
    const writer = &stdout_writer.interface;

    while (!story.is_exited and story.can_advance) {
        while (try story.advance()) |content| {
            try writer.print("{s}\n", .{content});
            try writer.flush();
        }
        if (story.current_choices.items.len > 0) {
            for (story.current_choices.items, 0..) |choice, index| {
                try writer.print("[{d}]: {s}\n", .{ index + 1, choice.content });
            }
            try writer.writeAll("> ");
            try writer.flush();

            if (try reader.takeDelimiter('\n')) |bytes| {
                const index = try std.fmt.parseUnsigned(usize, bytes, 10);
                try story.selectChoiceIndex(if (index > 0) index - 1 else index);
            }
        }
    }
}