const std = @import("std");
const sx = @import("sx");

pub fn main(init: std.process.Init) !void {
    const allocator = init.arena.allocator();
    const io = init.io;
    const args = try init.minimal.args.toSlice(allocator);

    if (args.len < 2) {
        printUsage();
        return;
    }

    const command = args[1];

    // LSP subcommand doesn't need a file argument
    if (std.mem.eql(u8, command, "lsp")) {
        runLsp(allocator, io);
        return;
    }

    if (args.len < 3) {
        printUsage();
        return;
    }

    const input_path = args[2];

    if (std.mem.eql(u8, command, "build")) {
        const output_name = deriveOutputName(input_path);
        compile(allocator, io, input_path, output_name) catch return;
        std.debug.print("compiled: {s}\n", .{output_name});
    } else if (std.mem.eql(u8, command, "ir")) {
        emitIR(allocator, io, input_path) catch return;
    } else if (std.mem.eql(u8, command, "run")) {
        const tmp_bin = "/tmp/sx_run_tmp";
        compile(allocator, io, input_path, tmp_bin) catch return;
        defer {
            std.Io.Dir.deleteFile(.cwd(), io, tmp_bin) catch {};
        }
        var child = std.process.spawn(io, .{
            .argv = &.{tmp_bin},
        }) catch {
            std.debug.print("error: failed to run program\n", .{});
            return;
        };
        _ = child.wait(io) catch {
            std.debug.print("error: program execution failed\n", .{});
            return;
        };
    } else {
        printUsage();
    }
}

fn printUsage() void {
    std.debug.print(
        \\Usage: sx <command> [file.sx]
        \\
        \\Commands:
        \\  run    Build and run immediately
        \\  build  Build binary in current directory
        \\  ir     Print LLVM IR to stdout
        \\  lsp    Start language server (LSP)
        \\
    , .{});
}

fn runLsp(allocator: std.mem.Allocator, io: std.Io) void {
    const Transport = sx.lsp.transport.Transport;
    const Server = sx.lsp.server.Server;

    const stdin_file = std.Io.File.stdin();
    const stdout_file = std.Io.File.stdout();

    var read_buf: [4096]u8 = undefined;
    var stdin_reader = stdin_file.readerStreaming(io, &read_buf);

    var transport = Transport.init(allocator, io, &stdin_reader.interface, stdout_file);
    var server = Server.init(allocator, &transport, io);

    while (true) {
        const msg = transport.readMessage() catch |err| {
            if (err == error.EndOfStream) break;
            std.debug.print("lsp: read error: {}\n", .{err});
            break;
        };

        const keep_going = server.handleMessage(msg);

        if (!keep_going) break;
    }
}

fn deriveOutputName(input_path: []const u8) []const u8 {
    // Get basename (strip directory)
    var start: usize = 0;
    for (input_path, 0..) |ch, i| {
        if (ch == '/') start = i + 1;
    }
    const basename = input_path[start..];
    // Strip .sx extension
    if (std.mem.endsWith(u8, basename, ".sx")) {
        return basename[0 .. basename.len - 3];
    }
    return basename;
}


fn readSource(allocator: std.mem.Allocator, io: std.Io, input_path: []const u8) ![:0]const u8 {
    const source_bytes = std.Io.Dir.readFileAlloc(.cwd(), io, input_path, allocator, .limited(10 * 1024 * 1024)) catch |err| {
        std.debug.print("error: cannot read '{s}': {}\n", .{ input_path, err });
        return error.CompileError;
    };
    return try allocator.dupeZ(u8, source_bytes);
}

fn emitIR(allocator: std.mem.Allocator, io: std.Io, input_path: []const u8) !void {
    const source = try readSource(allocator, io, input_path);

    var comp = sx.core.Compilation.init(allocator, io, input_path, source);
    defer comp.deinit();

    comp.parse() catch { comp.renderErrors(); return error.CompileError; };
    comp.resolveImports() catch { comp.renderErrors(); return error.CompileError; };
    comp.generateCode() catch { comp.renderErrors(); return error.CompileError; };

    var cg = &comp.cg.?;
    cg.verify() catch { comp.renderErrors(); return error.CompileError; };
    cg.printIR();
}

fn compile(allocator: std.mem.Allocator, io: std.Io, input_path: []const u8, output_path: []const u8) !void {
    const source = try readSource(allocator, io, input_path);

    var comp = sx.core.Compilation.init(allocator, io, input_path, source);
    defer comp.deinit();

    comp.parse() catch { comp.renderErrors(); return error.CompileError; };
    comp.resolveImports() catch { comp.renderErrors(); return error.CompileError; };
    comp.generateCode() catch { comp.renderErrors(); return error.CompileError; };

    var cg = &comp.cg.?;
    cg.verify() catch { comp.renderErrors(); return error.CompileError; };

    // Emit object file
    const obj_path = try std.fmt.allocPrintSentinel(allocator, "{s}.o", .{output_path}, 0);
    cg.emitObject(obj_path.ptr) catch { comp.renderErrors(); return error.CompileError; };

    // Link
    sx.codegen.CodeGen.link(io, obj_path, output_path) catch {
        std.debug.print("error: linking failed\n", .{});
        return error.CompileError;
    };

    // Clean up object file
    std.Io.Dir.deleteFile(.cwd(), io, obj_path) catch {};
}