const std = @import("std");
const ast = @import("ast.zig");
const llvm = @import("llvm_api.zig");
const Node = ast.Node;
const c = llvm.c;
const builtin = @import("builtin");

pub const CSourceLocation = struct {
    file: []const u8,
    line: u32,
};

/// Derive the NDK sysroot path from the NDK root (which by convention
/// lives in `target_config.sysroot` on Android — see target.zig's
/// Android link branch + main.zig's auto-discovery). Returns a NUL-
/// terminated path suitable for clang's `--sysroot <path>` argv.
fn androidSysrootFromNdkRoot(allocator: std.mem.Allocator, ndk_root: []const u8) ![:0]u8 {
    const host_tag: []const u8 = if (builtin.os.tag == .macos) "darwin-x86_64" else "linux-x86_64";
    return try std.fmt.allocPrintSentinel(allocator, "{s}/toolchains/llvm/prebuilt/{s}/sysroot", .{ ndk_root, host_tag }, 0);
}

pub const CImportResult = struct {
    fn_decls: []const *Node,
    /// Source locations for each fn_decl (parallel array, same indices).
    locations: []const CSourceLocation,
};

/// Info collected from c_import_decl AST nodes for native compilation.
pub const CImportInfo = struct {
    sources: []const []const u8,
    includes: []const []const u8,
    defines: []const []const u8,
    flags: []const []const u8,
};

/// Handle returned from loadCObjectsForJIT — caller must call unload() after JIT.
pub const CImportHandle = struct {
    dylib_handle: ?*anyopaque = null,
    temp_paths: []const []const u8 = &.{},
    allocator: std.mem.Allocator,

    pub fn unload(self: *CImportHandle, io: std.Io) void {
        // dlclose
        if (self.dylib_handle) |h| {
            _ = std.c.dlclose(h);
        }
        // Clean up temp files
        for (self.temp_paths) |path| {
            std.Io.Dir.deleteFile(.cwd(), io, path) catch {};
        }
    }
};

/// Parse C headers to extract function declarations as synthetic AST nodes.
/// Called during import resolution (no LLVM context needed).
pub fn processCImport(
    allocator: std.mem.Allocator,
    includes: []const []const u8,
    defines: []const []const u8,
    flags: []const []const u8,
) !CImportResult {
    // Build clang args: -I dirs, -D defines, raw flags
    var args_list = std.ArrayList([*c]const u8).empty;

    for (includes) |inc| {
        const dir = dirName(inc);
        const arg = try allocPrintZ(allocator, "-I{s}", .{dir});
        try args_list.append(allocator, arg.ptr);
    }
    for (defines) |def| {
        const arg = try allocPrintZ(allocator, "-D{s}", .{def});
        try args_list.append(allocator, arg.ptr);
    }
    for (flags) |flag| {
        const arg = try allocator.dupeZ(u8, flag);
        try args_list.append(allocator, arg.ptr);
    }

    var all_decls = std.ArrayList(*Node).empty;
    var all_locs = std.ArrayList(CSourceLocation).empty;

    for (includes) |header| {
        const header_z = try allocator.dupeZ(u8, header);
        var err_msg: [*c]u8 = null;

        const args_ptr: [*c][*c]const u8 = if (args_list.items.len > 0)
            @ptrCast(args_list.items.ptr)
        else
            null;

        const info = c.sx_clang_parse_header(
            header_z.ptr,
            args_ptr,
            @intCast(args_list.items.len),
            &err_msg,
        );

        if (info == null) {
            if (err_msg) |e| {
                std.debug.print("clang parse error for '{s}': {s}\n", .{ header, std.mem.span(e) });
            }
            return error.CompileError;
        }
        defer c.sx_clang_free_header_info(info);

        const funcs = info.*.functions;
        const num: usize = @intCast(info.*.num_functions);

        for (0..num) |i| {
            const fi = funcs[i];
            const name = try allocator.dupe(u8, std.mem.span(fi.name));

            // Build params
            var params = std.ArrayList(ast.Param).empty;
            const np: usize = @intCast(fi.num_params);
            for (0..np) |j| {
                const pi = fi.params[j];
                const pname_raw = std.mem.span(pi.name);
                const pname = if (pname_raw.len > 0)
                    try allocator.dupe(u8, pname_raw)
                else
                    try std.fmt.allocPrint(allocator, "p{d}", .{j});
                const ptype_str = std.mem.span(pi.type_spelling);
                const ptype_node = try mapCTypeToSxNode(allocator, ptype_str);

                try params.append(allocator, .{
                    .name = pname,
                    .name_span = .{ .start = 0, .end = 0 },
                    .type_expr = ptype_node,
                });
            }

            // Return type
            const ret_str = std.mem.span(fi.return_type);
            const ret_node = if (std.mem.eql(u8, ret_str, "void"))
                null
            else
                try mapCTypeToSxNode(allocator, ret_str);

            // Create foreign_expr body (no library_ref — symbols resolved at runtime)
            const foreign_body = try allocator.create(Node);
            foreign_body.* = .{
                .span = .{ .start = 0, .end = 0 },
                .data = .{ .foreign_expr = .{ .library_ref = null, .c_name = null } },
            };

            const fn_node = try allocator.create(Node);
            fn_node.* = .{
                .span = .{ .start = 0, .end = 0 },
                .data = .{ .fn_decl = .{
                    .name = name,
                    .params = try params.toOwnedSlice(allocator),
                    .return_type = ret_node,
                    .body = foreign_body,
                } },
            };

            try all_decls.append(allocator, fn_node);

            // Collect source location
            const src_file = if (fi.source_file) |sf|
                try allocator.dupe(u8, std.mem.span(sf))
            else
                header;
            try all_locs.append(allocator, .{
                .file = src_file,
                .line = @intCast(fi.source_line),
            });
        }
    }

    return .{
        .fn_decls = try all_decls.toOwnedSlice(allocator),
        .locations = try all_locs.toOwnedSlice(allocator),
    };
}

// ---------------------------------------------------------------------------
// Native C compilation (compile to .o, not LLVM module)
// ---------------------------------------------------------------------------

/// Compile C sources to native object files (in memory).
/// Returns list of LLVMMemoryBufferRef (each containing a .o file).
pub fn compileCToObjects(
    allocator: std.mem.Allocator,
    infos: []const CImportInfo,
    target_config: @import("target.zig").TargetConfig,
) ![]c.LLVMMemoryBufferRef {
    var obj_bufs = std.ArrayList(c.LLVMMemoryBufferRef).empty;

    for (infos) |info| {
        if (info.sources.len == 0) continue;

        // Build clang args: -I dirs, -D defines, raw flags
        var args_list = std.ArrayList([*c]const u8).empty;
        // Cross-compile target: forward -target / -isysroot when set.
        if (target_config.triple) |t| {
            try args_list.append(allocator, "-target");
            try args_list.append(allocator, t);
        }
        if (target_config.sysroot) |sr| {
            try args_list.append(allocator, "-isysroot");
            try args_list.append(allocator, (try allocator.dupeZ(u8, sr)).ptr);
        }
        // Android: route through the NDK sysroot so bionic headers resolve.
        // The embedded clang library doesn't know how to be an Android cross-
        // compiler on its own. `target_config.sysroot` holds the NDK root
        // by convention (main.zig auto-fills it for --target android), so
        // derive the headers/libs sysroot inside it.
        if (target_config.isAndroid()) {
            if (target_config.sysroot) |ndk_root| {
                const sysroot = try androidSysrootFromNdkRoot(allocator, ndk_root);
                try args_list.append(allocator, "--sysroot");
                try args_list.append(allocator, sysroot.ptr);
            }
        }
        for (info.includes) |inc| {
            const dir = dirName(inc);
            try args_list.append(allocator, (try allocPrintZ(allocator, "-I{s}", .{dir})).ptr);
        }
        for (info.defines) |def| {
            try args_list.append(allocator, (try allocPrintZ(allocator, "-D{s}", .{def})).ptr);
        }
        for (info.flags) |flag| {
            try args_list.append(allocator, (try allocator.dupeZ(u8, flag)).ptr);
        }

        const args_ptr: [*c][*c]const u8 = if (args_list.items.len > 0)
            @ptrCast(args_list.items.ptr)
        else
            null;
        const args_len: c_int = @intCast(args_list.items.len);

        for (info.sources) |src| {
            const src_z = try allocator.dupeZ(u8, src);
            var err_msg: [*c]u8 = null;

            const obj_buf = c.sx_clang_compile_to_object(
                src_z.ptr,
                args_ptr,
                args_len,
                &err_msg,
            );

            if (obj_buf == null) {
                if (err_msg) |e| {
                    std.debug.print("clang compile error for '{s}': {s}\n", .{ src, std.mem.span(e) });
                }
                return error.CompileError;
            }

            try obj_bufs.append(allocator, obj_buf);
        }
    }

    return try obj_bufs.toOwnedSlice(allocator);
}

/// For JIT mode: write .o files to temp, link into a shared library, dlopen it.
/// Returns a handle that must be unloaded after JIT execution.
pub fn loadCObjectsForJIT(
    allocator: std.mem.Allocator,
    io: std.Io,
    obj_bufs: []c.LLVMMemoryBufferRef,
) !CImportHandle {
    if (obj_bufs.len == 0) return .{ .allocator = allocator };

    var temp_paths = std.ArrayList([]const u8).empty;

    // Write each .o buffer to a temp file
    var obj_paths = std.ArrayList([]const u8).empty;
    for (obj_bufs, 0..) |buf, i| {
        const path = try std.fmt.allocPrint(allocator, "/tmp/sx_c_{d}.o", .{i});
        const start = c.LLVMGetBufferStart(buf);
        const size = c.LLVMGetBufferSize(buf);
        const data = @as([*]const u8, @ptrCast(start))[0..size];
        std.Io.Dir.writeFile(.cwd(), io, .{ .sub_path = path, .data = data }) catch {
            std.debug.print("failed to write temp object: {s}\n", .{path});
            return error.CompileError;
        };
        try obj_paths.append(allocator, path);
        try temp_paths.append(allocator, path);
        c.LLVMDisposeMemoryBuffer(buf);
    }

    // Link into a shared library
    const dylib_path = "/tmp/sx_c_import.dylib";
    try temp_paths.append(allocator, try allocator.dupe(u8, dylib_path));

    var argv = std.ArrayList([]const u8).empty;
    try argv.append(allocator, "cc");
    if (comptime builtin.os.tag == .macos) {
        try argv.append(allocator, "-dynamiclib");
    } else {
        try argv.append(allocator, "-shared");
    }
    try argv.append(allocator, "-o");
    try argv.append(allocator, dylib_path);
    for (obj_paths.items) |op| {
        try argv.append(allocator, op);
    }

    const argv_slice = try argv.toOwnedSlice(allocator);
    var child = std.process.spawn(io, .{
        .argv = argv_slice,
    }) catch {
        std.debug.print("failed to spawn linker for C import shared library\n", .{});
        return error.CompileError;
    };
    const result = child.wait(io) catch {
        std.debug.print("linker wait failed for C import shared library\n", .{});
        return error.CompileError;
    };
    if (result != .exited or result.exited != 0) {
        std.debug.print("linker failed for C import shared library (exit={})\n", .{result.exited});
        return error.CompileError;
    }

    // dlopen the shared library
    const dylib_z = try allocator.dupeZ(u8, dylib_path);
    const handle = std.c.dlopen(dylib_z.ptr, .{ .NOW = true });
    if (handle == null) {
        const err = std.c.dlerror();
        if (err) |e| {
            std.debug.print("dlopen failed: {s}\n", .{std.mem.span(e)});
        }
        return error.CompileError;
    }

    return .{
        .dylib_handle = handle,
        .temp_paths = try temp_paths.toOwnedSlice(allocator),
        .allocator = allocator,
    };
}

/// Compile C sources using emcc for Emscripten/WASM targets.
/// Shells out to `emcc -c` for each source file, returns temp object file paths.
pub fn compileCWithEmcc(
    allocator: std.mem.Allocator,
    io: std.Io,
    infos: []const CImportInfo,
    target_config: @import("target.zig").TargetConfig,
    tmp_dir: []const u8,
) ![]const []const u8 {
    var paths = std.ArrayList([]const u8).empty;
    var obj_idx: usize = 0;

    for (infos) |info| {
        if (info.sources.len == 0) continue;

        for (info.sources) |src| {
            const out_path = try std.fmt.allocPrint(allocator, "{s}/sx_emcc_{d}.o", .{ tmp_dir, obj_idx });
            obj_idx += 1;

            var argv = std.ArrayList([]const u8).empty;
            try argv.appendSlice(allocator, &.{ "emcc", "-c", "-O2", src, "-o", out_path });
            // wasm64: compile C sources with memory64 support
            if (target_config.isWasm64()) {
                try argv.append(allocator, "-sMEMORY64");
            }
            // Add include paths
            for (info.includes) |inc| {
                try argv.append(allocator, try std.fmt.allocPrint(allocator, "-I{s}", .{dirName(inc)}));
            }
            for (info.defines) |def| {
                try argv.append(allocator, try std.fmt.allocPrint(allocator, "-D{s}", .{def}));
            }
            for (info.flags) |flag| {
                try argv.append(allocator, flag);
            }

            const argv_slice = try argv.toOwnedSlice(allocator);
            var child = std.process.spawn(io, .{ .argv = argv_slice }) catch {
                std.debug.print("error: failed to spawn emcc for '{s}'\n", .{src});
                return error.CompileError;
            };
            const result = child.wait(io) catch return error.CompileError;
            if (result != .exited or result.exited != 0) {
                std.debug.print("error: emcc failed for '{s}'\n", .{src});
                return error.CompileError;
            }
            try paths.append(allocator, out_path);
        }
    }

    return try paths.toOwnedSlice(allocator);
}

/// For build mode: write .o buffers to temp files, return paths for the linker.
pub fn writeCObjectFiles(
    allocator: std.mem.Allocator,
    io: std.Io,
    obj_bufs: []c.LLVMMemoryBufferRef,
    tmp_dir: []const u8,
) ![]const []const u8 {
    var paths = std.ArrayList([]const u8).empty;

    for (obj_bufs, 0..) |buf, i| {
        const path = try std.fmt.allocPrint(allocator, "{s}/sx_c_{d}.o", .{ tmp_dir, i });
        const start = c.LLVMGetBufferStart(buf);
        const size = c.LLVMGetBufferSize(buf);
        const data = @as([*]const u8, @ptrCast(start))[0..size];
        std.Io.Dir.writeFile(.cwd(), io, .{ .sub_path = path, .data = data }) catch {
            std.debug.print("failed to write temp object: {s}\n", .{path});
            return error.CompileError;
        };
        try paths.append(allocator, path);
        c.LLVMDisposeMemoryBuffer(buf);
    }

    return try paths.toOwnedSlice(allocator);
}

/// Walk the resolved AST and collect CImportInfo from all c_import_decl nodes.
/// Deduplicates by source pointer identity (shared nodes from import propagation).
pub fn collectCImportSources(allocator: std.mem.Allocator, root: *const Node) ![]CImportInfo {
    if (root.data != .root) return &.{};

    var infos = std.ArrayList(CImportInfo).empty;
    var seen = std.AutoHashMap([*]const []const u8, void).init(allocator);
    defer seen.deinit();

    for (root.data.root.decls) |decl| {
        switch (decl.data) {
            .c_import_decl => |ci| {
                if (ci.sources.len > 0) {
                    const key = ci.sources.ptr;
                    if (!seen.contains(key)) {
                        try seen.put(key, {});
                        try infos.append(allocator, .{
                            .sources = ci.sources,
                            .includes = ci.includes,
                            .defines = ci.defines,
                            .flags = ci.flags,
                        });
                    }
                }
            },
            .namespace_decl => |ns| {
                for (ns.decls) |nd| {
                    if (nd.data == .c_import_decl) {
                        const nci = nd.data.c_import_decl;
                        if (nci.sources.len > 0) {
                            const key = nci.sources.ptr;
                            if (!seen.contains(key)) {
                                try seen.put(key, {});
                                try infos.append(allocator, .{
                                    .sources = nci.sources,
                                    .includes = nci.includes,
                                    .defines = nci.defines,
                                    .flags = nci.flags,
                                });
                            }
                        }
                    }
                }
            },
            else => {},
        }
    }

    return try infos.toOwnedSlice(allocator);
}

// ---------------------------------------------------------------------------
// C type → sx type mapping
// ---------------------------------------------------------------------------

fn mapCTypeToSxNode(
    allocator: std.mem.Allocator,
    c_type: []const u8,
) !*Node {
    const trimmed = std.mem.trim(u8, c_type, " ");

    // Pointer types (trailing *)
    if (std.mem.endsWith(u8, trimmed, "*")) {
        const base = std.mem.trim(u8, trimmed[0 .. trimmed.len - 1], " ");

        // const char * → [*]u8 (raw pointer, matches C ABI)
        if (std.mem.eql(u8, base, "const char") or std.mem.eql(u8, base, "char const")) {
            return makeManyPointerTypeNode(allocator, "u8");
        }
        // char * → [*]u8
        if (std.mem.eql(u8, base, "char")) {
            return makeManyPointerTypeNode(allocator, "u8");
        }
        // unsigned char * / const unsigned char * → [*]u8
        if (std.mem.eql(u8, base, "unsigned char") or
            std.mem.eql(u8, base, "const unsigned char") or
            std.mem.eql(u8, base, "unsigned char const"))
        {
            return makeManyPointerTypeNode(allocator, "u8");
        }
        // void * / const void * → *void
        if (std.mem.eql(u8, base, "void") or std.mem.eql(u8, base, "const void")) {
            return makePointerTypeNode(allocator, "void");
        }
        // int * → *s32
        if (std.mem.eql(u8, base, "int") or std.mem.eql(u8, base, "const int")) {
            return makePointerTypeNode(allocator, "s32");
        }
        // unsigned int * / unsigned * → *u32
        if (std.mem.eql(u8, base, "unsigned int") or std.mem.eql(u8, base, "unsigned") or std.mem.eql(u8, base, "const unsigned int")) {
            return makePointerTypeNode(allocator, "u32");
        }
        // float * → *f32
        if (std.mem.eql(u8, base, "float") or std.mem.eql(u8, base, "const float")) {
            return makePointerTypeNode(allocator, "f32");
        }
        // double * → *f64
        if (std.mem.eql(u8, base, "double") or std.mem.eql(u8, base, "const double")) {
            return makePointerTypeNode(allocator, "f64");
        }
        // short * → *s16
        if (std.mem.eql(u8, base, "short") or std.mem.eql(u8, base, "const short")) {
            return makePointerTypeNode(allocator, "s16");
        }
        // Pointer to pointer → *void
        if (std.mem.endsWith(u8, base, "*")) {
            return makePointerTypeNode(allocator, "void");
        }
        // Remove const qualifier and retry
        if (std.mem.startsWith(u8, base, "const ")) {
            const without_const = try std.fmt.allocPrint(allocator, "{s} *", .{base[6..]});
            return mapCTypeToSxNode(allocator, without_const);
        }
        // Default: struct/opaque pointer → *void
        return makePointerTypeNode(allocator, "void");
    }

    // Direct types
    if (std.mem.eql(u8, trimmed, "int") or std.mem.eql(u8, trimmed, "signed int")) return makeTypeExprNode(allocator, "s32");
    if (std.mem.eql(u8, trimmed, "unsigned int") or std.mem.eql(u8, trimmed, "unsigned")) return makeTypeExprNode(allocator, "u32");
    if (std.mem.eql(u8, trimmed, "long") or std.mem.eql(u8, trimmed, "long int") or std.mem.eql(u8, trimmed, "signed long")) return makeTypeExprNode(allocator, "s64");
    if (std.mem.eql(u8, trimmed, "unsigned long") or std.mem.eql(u8, trimmed, "unsigned long int")) return makeTypeExprNode(allocator, "u64");
    if (std.mem.eql(u8, trimmed, "long long") or std.mem.eql(u8, trimmed, "long long int")) return makeTypeExprNode(allocator, "s64");
    if (std.mem.eql(u8, trimmed, "unsigned long long") or std.mem.eql(u8, trimmed, "unsigned long long int")) return makeTypeExprNode(allocator, "u64");
    if (std.mem.eql(u8, trimmed, "short") or std.mem.eql(u8, trimmed, "short int") or std.mem.eql(u8, trimmed, "signed short")) return makeTypeExprNode(allocator, "s16");
    if (std.mem.eql(u8, trimmed, "unsigned short") or std.mem.eql(u8, trimmed, "unsigned short int")) return makeTypeExprNode(allocator, "u16");
    if (std.mem.eql(u8, trimmed, "char") or std.mem.eql(u8, trimmed, "signed char")) return makeTypeExprNode(allocator, "u8");
    if (std.mem.eql(u8, trimmed, "unsigned char")) return makeTypeExprNode(allocator, "u8");
    if (std.mem.eql(u8, trimmed, "float")) return makeTypeExprNode(allocator, "f32");
    if (std.mem.eql(u8, trimmed, "double")) return makeTypeExprNode(allocator, "f64");
    if (std.mem.eql(u8, trimmed, "size_t")) return makeTypeExprNode(allocator, "u64");
    if (std.mem.eql(u8, trimmed, "_Bool") or std.mem.eql(u8, trimmed, "bool")) return makeTypeExprNode(allocator, "u8");

    // Default: unknown type → s64 (treat as opaque integer-sized value)
    return makeTypeExprNode(allocator, "s64");
}

// ---------------------------------------------------------------------------
// AST node construction helpers
// ---------------------------------------------------------------------------

fn makeTypeExprNode(allocator: std.mem.Allocator, name: []const u8) !*Node {
    const node = try allocator.create(Node);
    node.* = .{
        .span = .{ .start = 0, .end = 0 },
        .data = .{ .type_expr = .{ .name = name } },
    };
    return node;
}

fn makePointerTypeNode(allocator: std.mem.Allocator, pointee: []const u8) !*Node {
    const inner = try makeTypeExprNode(allocator, pointee);
    const node = try allocator.create(Node);
    node.* = .{
        .span = .{ .start = 0, .end = 0 },
        .data = .{ .pointer_type_expr = .{ .pointee_type = inner } },
    };
    return node;
}

fn makeManyPointerTypeNode(allocator: std.mem.Allocator, element: []const u8) !*Node {
    const inner = try makeTypeExprNode(allocator, element);
    const node = try allocator.create(Node);
    node.* = .{
        .span = .{ .start = 0, .end = 0 },
        .data = .{ .many_pointer_type_expr = .{ .element_type = inner } },
    };
    return node;
}

fn makeSliceTypeNode(allocator: std.mem.Allocator, element: []const u8) !*Node {
    const inner = try makeTypeExprNode(allocator, element);
    const node = try allocator.create(Node);
    node.* = .{
        .span = .{ .start = 0, .end = 0 },
        .data = .{ .slice_type_expr = .{ .element_type = inner } },
    };
    return node;
}

// ---------------------------------------------------------------------------
// Utility
// ---------------------------------------------------------------------------

fn allocPrintZ(allocator: std.mem.Allocator, comptime fmt: []const u8, args: anytype) ![:0]u8 {
    return allocator.dupeZ(u8, try std.fmt.allocPrint(allocator, fmt, args));
}

fn dirName(path: []const u8) []const u8 {
    var last_sep: usize = 0;
    var found = false;
    for (path, 0..) |ch, i| {
        if (ch == '/') {
            last_sep = i;
            found = true;
        }
    }
    return if (found) path[0..last_sep] else ".";
}