import "@ffi/c"
import "@std/mem"
import "@std/arraylist"
import "@std/enummap"
import "@std/io"
import "@std/process"

import "@source/diagnostics"
import "@source/compilation"
import "@source/ir"
import "@source/interpreter"
import "@source/llvm"

test import "@std/mem"
test import "@std/enummap"
test import "@std/arraylist"
test import "@std/hashmap"
test import "@std/strmap"
test import "@std/meta"
test import "@std/io"
test import "@std/process"
test import "@std/fs"
test import "@std/strconv"
test import "@source/diagnostics"
test import "@source/strpool"
test import "@source/lexer"
test import "@source/ast"
test import "@source/parser"
test import "@source/sema"
test import "@source/sema/check"
test import "@source/ir"
test import "@source/lower"
test import "@source/interpreter"
test import "@source/compilation"
test import "@source/loader"

@hide:file
DriverError :: alias (
    compilation.Error
    | io.OpenError
    | io.CloseError
    | io.ReadError
    | io.WriteError
    | process.RunError
    | NativeBuildError
)

@hide:file
NativeBuildError :: enum { temporary_file_failed, compiler_failed }

main proc(init process.Init) i32 {
    return command(mem.c_allocator, init.io, init.args) catch |err| {
        if (init.args.values.len > 2) {
            io.print("{s}: infrastructure error: {s}\n", io.stderr(init.io), {
                init.args.values[2], enummap.get(&infrastructure_messages, err)?,
            }) catch |_| {}
        } else {
            io.print("honey: infrastructure error: {s}\n", io.stderr(init.io), {
                enummap.get(&infrastructure_messages, err)?,
            }) catch |_| {}
        }
        return 2
    }
}

@hide:file
infrastructure_messages :: $enummap.init(DriverError, []u8, {
    out_of_memory = "out of memory",
    diagnostic_limit_reached = "diagnostic limit reached",
    storage_limit_reached = "aggregate storage limit reached",
    source_too_large = "source file is too large",
    out_of_space = "identifier storage is full",
    open_failed = "cannot open source file",
    close_failed = "cannot close source file",
    not_open_for_reading = "source file is not readable",
    read_failed = "cannot read source file",
    not_open_for_writing = "output is not writable",
    write_failed = "cannot write output",
    no_progress = "output writer made no progress",
    spawn_failed = "cannot start native compiler",
    wait_failed = "cannot wait for native compiler",
    invalid_arguments = "invalid native compiler arguments",
    temporary_file_failed = "cannot create temporary LLVM file",
    compiler_failed = "native compiler failed",
    empty_package = "package contains no production source files",
    invalid_path = "invalid filesystem path",
    path_failed = "cannot resolve filesystem path",
    directory_failed = "cannot read package directory",
})

@hide:file
command proc(allocator mem.Allocator, runtime io.Io, args process.Args) i32 ! DriverError {
    if (args.values.len == 3) {
        if (mem.eql(args.values[1], "run")) return try run_file(allocator, runtime, args.values[2])
        if (mem.eql(args.values[1], "emit-llvm")) return try emit_file(allocator, runtime, args.values[2])
    }
    if (args.values.len == 5 or args.values.len == 6) {
        if (mem.eql(args.values[1], "build") and mem.eql(args.values[3], "-o") and args.values[4].len != 0) {
            optimization [;0]u8 := "-O0"
            if (args.values.len == 6) {
                optimization = args.values[5]
                if (!mem.eql(optimization, "-O0") and !mem.eql(optimization, "-O2")) {
                    try usage(io.stderr(runtime))
                    return 2
                }
            }
            return try build_file(allocator, runtime, args.values[2], args.values[4], optimization)
        }
    }
    try usage(io.stderr(runtime))
    return 2
}

@hide:file
usage proc(output io.Writer) void ! io.WriteError {
    try io.write_all(output, "usage: honey run <file-or-directory>\n")
    try io.write_all(output, "       honey emit-llvm <file-or-directory>\n")
    try io.write_all(output, "       honey build <file-or-directory> -o <executable> [-O0|-O2]\n")
}

@hide:file
emit_file proc(allocator mem.Allocator, runtime io.Io, path [;0]u8) i32 ! DriverError {
    session := compilation.init(allocator)
    defer compilation.deinit(&session)
    diagnostic_store := diagnostics.init(allocator)
    defer diagnostics.deinit(&diagnostic_store)
    root :: try compilation.load(&session, runtime, path, &diagnostic_store)
    module := try compilation.compile(&session, root, &diagnostic_store)
    defer ir.deinit(&module)
    sources :: compilation.source_files(&session)
    primary :: session.packages.items[usize(root)].sources.items[0]
    try llvm.emit(allocator, io.stdout(runtime), &module, &diagnostic_store, sources, primary)
    try report_diagnostics(allocator, io.stderr(runtime), sources, &diagnostic_store, null)
    return if (diagnostics.count(&diagnostic_store) == 0) 0 else 1
}

@hide:file
build_file proc(
    allocator mem.Allocator,
    runtime io.Io,
    path [;0]u8,
    output [;0]u8,
    optimization [;0]u8,
) i32 ! DriverError {
    session := compilation.init(allocator)
    defer compilation.deinit(&session)
    diagnostic_store := diagnostics.init(allocator)
    defer diagnostics.deinit(&diagnostic_store)
    root :: try compilation.load(&session, runtime, path, &diagnostic_store)
    module := try compilation.compile(&session, root, &diagnostic_store)
    defer ir.deinit(&module)

    template [18]mut u8 := undefined
    memcopy!(template[..], "/tmp/honey-XXXXXX\0")
    descriptor :: c.mkstemp(ptrcast!(template[..].ptr))
    if (descriptor < 0) return .temporary_file_failed
    defer _ = c.unlink(ptrcast!(template[..].ptr))
    file :: io.File{ io = runtime, handle = io.Handle{ file_desc = descriptor } }
    closing := false
    defer {
        if (!closing) io.close(file) catch |_| {}
    }
    sources :: compilation.source_files(&session)
    primary :: session.packages.items[usize(root)].sources.items[0]
    try llvm.emit(allocator, io.writer(file), &module, &diagnostic_store, sources, primary)
    closing = true
    try io.close(file)
    temporary [;0]u8 :: template[..].ptr[..17;0]
    arguments [9][;0]u8 :: [
        "clang", "-target", "arm64-apple-macosx11.0.0",
        "-x", "ir", optimization, temporary, "-o", output,
    ]
    status :: try process.run(allocator, arguments[..])
    if (status != 0) return .compiler_failed
    try report_diagnostics(allocator, io.stderr(runtime), sources, &diagnostic_store, null)
    return if (diagnostics.count(&diagnostic_store) == 0) 0 else 1
}

@hide:file
run_file proc(allocator mem.Allocator, runtime io.Io, path [;0]u8) i32 ! DriverError {
    session := compilation.init(allocator)
    defer compilation.deinit(&session)
    diagnostic_store := diagnostics.init(allocator)
    defer diagnostics.deinit(&diagnostic_store)
    root :: try compilation.load(&session, runtime, path, &diagnostic_store)
    module := try compilation.compile(&session, root, &diagnostic_store)
    defer ir.deinit(&module)
    outcome := try interpreter.run(allocator, &module, &diagnostic_store)
    defer interpreter.deinit(&outcome)
    if (outcome.trap == null) try print_result(io.stdout(runtime), &module, outcome.value)
    try report_diagnostics(allocator, io.stderr(runtime), compilation.source_files(&session), &diagnostic_store, outcome.trap)
    if (outcome.trap != null or diagnostics.count(&diagnostic_store) != 0) return 1
    return 0
}

@hide:file
print_result proc(output io.Writer, module @ir.Module, value ir.Value) void ! io.WriteError {
    if (value.type == .void) return
    try print_value(output, module, value)
    try io.print("\n", output, {})
}

@hide:file
print_value proc(output io.Writer, module @ir.Module, value ir.Value) void ! io.WriteError {
    match value.type {
        .int, .i8, .i16, .i32, .isize: try io.print("{}", output, {value.data.integer})
        .usize, .u8, .u16, .u32, .u64: try io.print("{}", output, {value.data.unsigned})
        .float: try io.print("{d}", output, {value.data.real})
        .bool: try io.print("{}", output, {value.data.boolean})
        .string: try io.print("{s}", output, {ir.string_bytes(module, value.data.string)})
        .array, .record: {
            shape :: module.arrays.items[usize(value.array_type)]
            count :: if (value.type == .array) shape.len else shape.fields.len
            opening :: if (value.type == .array) "[" else "{"
            closing :: if (value.type == .array) "]" else "}"
            try io.write_all(output, opening)
            for (0..count) |index| {
                if (index != 0) try io.write_all(output, ", ")
                try print_value(output, module, ir.aggregate_item(module, value, index))
            }
            try io.write_all(output, closing)
        }
        else: unreachable
    }
}

@hide:file
report_diagnostics proc(
    allocator mem.Allocator,
    output io.Writer,
    sources []diagnostics.Source,
    diagnostic_store @diagnostics.Store,
    trap ?diagnostics.DiagnosticId,
) void ! (mem.AllocError | io.WriteError) {
    if (diagnostics.count(diagnostic_store) == 0) return
    line_starts := arraylist.init(arraylist.ArrayList(usize), allocator)
    defer {
        for (line_starts.items) |@starts| arraylist.deinit(starts)
        arraylist.deinit(&line_starts)
    }
    for (sources) |source| {
        starts := arraylist.init(usize, allocator)
        errdefer arraylist.deinit(&starts)
        try arraylist.append(&starts, usize(0))
        for (source.text) |byte, index| {
            if (byte == '\n') try arraylist.append(&starts, index + 1)
        }
        try arraylist.append(&line_starts, starts)
    }
    for (0..diagnostics.count(diagnostic_store)) |index| {
        if (diagnostics.get(diagnostic_store, diagnostics.DiagnosticId(index))) |diagnostic| {
            source :: sources[usize(diagnostic.span.source)]
            starts :: line_starts.items[usize(diagnostic.span.source)].items
            offset :: usize(diagnostic.span.start)
            line := usize(0)
            end := starts.len
            while (line + 1 < end) {
                middle :: line + divtrunc!(end - line, 2)
                if (starts[middle] <= offset) line = middle
                else end = middle
            }
            label []u8 := "error"
            if (trap) |id| {
                if (id == diagnostics.DiagnosticId(index)) label = "runtime trap"
            }
            try io.print("{s}:{}:{}: {s}: {s}\n", output, {
                source.path,
                line + 1,
                offset - starts[line] + 1,
                label,
                enummap.get(diagnostics.Code, []u8, &diagnostics.messages, diagnostic.code)?,
            })
        }
    }
}