const std = @import("std");
const ast = @import("ast.zig");
const parser = @import("parser.zig");
const errors = @import("errors.zig");
const c_import = @import("c_import.zig");
const Node = ast.Node;

pub 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 ".";
}

/// Resolve an import path. Tries (in order):
///   1. relative to `base_dir` (the importing file's directory)
///   2. relative to CWD, absolutified via `root_path` if supplied
///   3. relative to each path in `stdlib_paths` (the install-discovered stdlib)
/// Returns the first path that exists. Falls back to the raw path if nothing matches
/// so the caller's readFile produces a coherent "not found" error.
pub fn resolveImportPath(allocator: std.mem.Allocator, io: std.Io, base_dir: []const u8, raw_path: []const u8, root_path: ?[]const u8, stdlib_paths: []const []const u8) ![]const u8 {
    if (!std.mem.eql(u8, base_dir, ".")) {
        const rel_path = try std.fmt.allocPrint(allocator, "{s}/{s}", .{ base_dir, raw_path });
        // Check if it exists as file relative to base_dir
        if (std.Io.Dir.readFileAlloc(.cwd(), io, rel_path, allocator, .limited(10 * 1024 * 1024))) |_| {
            return rel_path;
        } else |_| {}
        // Check if it exists as directory relative to base_dir
        if (std.Io.Dir.openDir(.cwd(), io, rel_path, .{})) |dir| {
            dir.close(io);
            return rel_path;
        } else |_| {}
    }
    // Try CWD-relative (absolutified if root_path is known).
    const cwd_candidate = if (root_path) |rp| blk: {
        if (rp.len > 0 and raw_path.len > 0 and raw_path[0] != '/') {
            break :blk try std.fmt.allocPrint(allocator, "{s}/{s}", .{ rp, raw_path });
        }
        break :blk raw_path;
    } else raw_path;
    if (std.Io.Dir.readFileAlloc(.cwd(), io, cwd_candidate, allocator, .limited(10 * 1024 * 1024))) |_| {
        return cwd_candidate;
    } else |_| {}
    if (std.Io.Dir.openDir(.cwd(), io, cwd_candidate, .{})) |dir| {
        dir.close(io);
        return cwd_candidate;
    } else |_| {}
    // Try each stdlib search path.
    for (stdlib_paths) |sp| {
        const cand = try std.fmt.allocPrint(allocator, "{s}/{s}", .{ sp, raw_path });
        if (std.Io.Dir.readFileAlloc(.cwd(), io, cand, allocator, .limited(10 * 1024 * 1024))) |_| {
            return cand;
        } else |_| {}
        if (std.Io.Dir.openDir(.cwd(), io, cand, .{})) |dir| {
            dir.close(io);
            return cand;
        } else |_| {}
    }
    return cwd_candidate;
}

/// Discover candidate stdlib search paths from the running binary's location.
/// Honors the `SX_STDLIB_PATH` env var as an explicit override. Returns a slice
/// of absolute paths owned by the allocator.
pub fn discoverStdlibPaths(allocator: std.mem.Allocator) ![]const []const u8 {
    var out = std.ArrayList([]const u8).empty;

    // Env override via libc getenv (cross-stdlib-version stable).
    if (c_getenv("SX_STDLIB_PATH")) |env_path| {
        try out.append(allocator, try allocator.dupe(u8, std.mem.span(env_path)));
    }

    const exe_path = selfExePath(allocator) catch return try out.toOwnedSlice(allocator);
    const exe_dir = dirName(exe_path);
    // Stdlib paths are directories containing a `modules/` subdir; the import
    // directive (e.g. `#import "modules/std.sx"`) supplies the rest.
    // Dev: zig-out/bin/sx -> repo-root/library
    try out.append(allocator, try std.fmt.allocPrint(allocator, "{s}/../../library", .{exe_dir}));
    // Install: <prefix>/bin/sx -> <prefix>/library
    try out.append(allocator, try std.fmt.allocPrint(allocator, "{s}/../library", .{exe_dir}));
    // Alongside the binary.
    try out.append(allocator, try std.fmt.allocPrint(allocator, "{s}/library", .{exe_dir}));
    if (c_getenv("SX_DEBUG_STDLIB") != null) {
        std.debug.print("[sx] exe_path={s}\n", .{exe_path});
        for (out.items, 0..) |p, i| std.debug.print("[sx] stdlib_paths[{d}]={s}\n", .{ i, p });
    }
    return try out.toOwnedSlice(allocator);
}

const builtin = @import("builtin");

extern "c" fn _NSGetExecutablePath(buf: [*]u8, len: *u32) c_int;
extern "c" fn getenv(name: [*:0]const u8) ?[*:0]const u8;

fn c_getenv(name: [:0]const u8) ?[*:0]const u8 {
    return getenv(name.ptr);
}

fn selfExePath(allocator: std.mem.Allocator) ![]const u8 {
    var buf: [4096]u8 = undefined;
    switch (builtin.os.tag) {
        .macos, .ios => {
            var len: u32 = buf.len;
            if (_NSGetExecutablePath(&buf, &len) != 0) return error.PathBufferTooSmall;
            const span = std.mem.sliceTo(&buf, 0);
            return try allocator.dupe(u8, span);
        },
        .linux => {
            const n = try std.posix.readlink("/proc/self/exe", &buf);
            return try allocator.dupe(u8, n);
        },
        else => return error.UnsupportedHostOS,
    }
}

/// A resolved module: the fully-resolved declarations of a single .sx file,
/// with its own scope tracking which names are defined.
pub const ResolvedModule = struct {
    path: []const u8,
    decls: []const *Node,
    scope: std.StringHashMap(void),

    /// Try to add a declaration. Returns true if added, false if name already in scope.
    pub fn addDecl(self: *ResolvedModule, allocator: std.mem.Allocator, list: *std.ArrayList(*Node), decl: *Node) !bool {
        if (decl.data.declName()) |name| {
            if (self.scope.contains(name)) return false;
            try self.scope.put(name, {});
        }
        try list.append(allocator, decl);
        return true;
    }

    /// Merge another module's decls as flat imports (skipping duplicates).
    pub fn mergeFlat(self: *ResolvedModule, allocator: std.mem.Allocator, list: *std.ArrayList(*Node), other: ResolvedModule) !void {
        for (other.decls) |decl| {
            _ = try self.addDecl(allocator, list, decl);
        }
    }

    /// Add another module as a namespaced import.
    pub fn addNamespace(self: *ResolvedModule, allocator: std.mem.Allocator, list: *std.ArrayList(*Node), name: []const u8, other: ResolvedModule, span: ast.Span) !void {
        const ns_node = try allocator.create(Node);
        ns_node.* = .{
            .span = span,
            .data = .{ .namespace_decl = .{
                .name = name,
                .decls = other.decls,
            } },
        };
        try self.scope.put(name, {});
        try list.append(allocator, ns_node);
    }

    pub fn finalize(self: *ResolvedModule, allocator: std.mem.Allocator, list: *std.ArrayList(*Node)) !void {
        self.decls = try list.toOwnedSlice(allocator);
    }
};

/// Module cache: maps resolved file paths to their ResolvedModules.
pub const ModuleCache = std.StringHashMap(ResolvedModule);

pub fn resolveImports(
    allocator: std.mem.Allocator,
    io: std.Io,
    root: *Node,
    base_dir: []const u8,
    file_path: []const u8,
    chain: *std.StringHashMap(void),
    cache: *ModuleCache,
    source_map: ?*std.StringHashMap([:0]const u8),
    diagnostics: ?*errors.DiagnosticList,
    stdlib_paths: []const []const u8,
    import_graph: ?*std.StringHashMap(std.StringHashMap(void)),
) !ResolvedModule {
    // Record this file's edge set so `param_impl_map` lookups can filter
    // candidates by what's been imported from where. Populated as each
    // import resolves below; transitive closure computed on demand.
    if (import_graph) |g| {
        if (!g.contains(file_path)) {
            try g.put(file_path, std.StringHashMap(void).init(allocator));
        }
    }
    var mod = ResolvedModule{
        .path = file_path,
        .decls = &.{},
        .scope = std.StringHashMap(void).init(allocator),
    };

    if (root.data != .root) {
        mod.decls = &.{};
        return mod;
    }

    var decl_list = std.ArrayList(*Node).empty;

    for (root.data.root.decls) |decl| {
        if (decl.data == .c_import_decl) {
            const ci = decl.data.c_import_decl;

            // Parse headers to get synthetic function declarations
            const result = c_import.processCImport(
                allocator,
                ci.includes,
                ci.defines,
                ci.flags,
            ) catch |err| {
                if (diagnostics) |diags| {
                    diags.addFmt(.err, decl.span, "#import c failed: {}", .{err});
                }
                return error.ImportError;
            };

            if (ci.name) |ns_name| {
                // Namespaced: wrap fn_decls + c_import_decl in a namespace
                var ns_decls = std.ArrayList(*Node).empty;
                for (result.fn_decls) |fd| {
                    try ns_decls.append(allocator, fd);
                }
                // Keep c_import_decl inside namespace so codegen can find sources
                try ns_decls.append(allocator, decl);

                const ns_node = try allocator.create(Node);
                ns_node.* = .{
                    .span = decl.span,
                    .data = .{ .namespace_decl = .{
                        .name = ns_name,
                        .decls = try ns_decls.toOwnedSlice(allocator),
                    } },
                };
                ns_node.source_file = file_path;
                try mod.scope.put(ns_name, {});
                try decl_list.append(allocator, ns_node);
            } else {
                // Flat: add fn_decls directly + keep c_import_decl
                for (result.fn_decls) |fd| {
                    fd.source_file = file_path;
                    _ = try mod.addDecl(allocator, &decl_list, fd);
                }
                decl.source_file = file_path;
                _ = try mod.addDecl(allocator, &decl_list, decl);
            }
            continue;
        }
        if (decl.data != .import_decl) {
            decl.source_file = file_path;
            _ = try mod.addDecl(allocator, &decl_list, decl);
            continue;
        }
        const imp = decl.data.import_decl;

        const resolved_path = try resolveImportPath(allocator, io, base_dir, imp.path, null, stdlib_paths);

        // Record direct-import edge file_path → resolved_path. Self-imports
        // and chain duplicates are still recorded so the graph reflects what
        // the user wrote (filter happens at lookup).
        if (import_graph) |g| {
            if (g.getPtr(file_path)) |set| {
                set.put(resolved_path, {}) catch {};
            }
        }

        // Circular import check — only along the current chain
        if (chain.contains(resolved_path)) continue;

        // Resolve or retrieve the imported module
        const imported_mod = if (cache.get(resolved_path)) |cached|
            cached
        else blk: {
            // Try as file first
            if (std.Io.Dir.readFileAlloc(.cwd(), io, resolved_path, allocator, .limited(10 * 1024 * 1024))) |imp_bytes| {
                const imp_source = try allocator.dupeZ(u8, imp_bytes);

                if (source_map) |sm| {
                    sm.put(resolved_path, imp_source) catch {};
                }

                var p = parser.Parser.init(allocator, imp_source);
                const imp_root = p.parse() catch {
                    if (diagnostics) |diags| {
                        diags.addFmt(.err, decl.span, "parse error in '{s}': {s}", .{ resolved_path, p.err_msg orelse "unknown" });
                    }
                    return error.ImportError;
                };

                // Push onto chain before recursing, pop after
                try chain.put(resolved_path, {});
                const imp_dir = dirName(resolved_path);
                const result = try resolveImports(allocator, io, imp_root, imp_dir, resolved_path, chain, cache, source_map, diagnostics, stdlib_paths, import_graph);
                _ = chain.remove(resolved_path);

                // Cache
                try cache.put(resolved_path, result);
                break :blk result;
            } else |_| {
                // File read failed — try as directory import
                const result = resolveDirectoryImport(allocator, io, resolved_path, chain, cache, source_map, diagnostics, decl.span, stdlib_paths, import_graph) catch {
                    if (diagnostics) |diags| {
                        diags.addFmt(.err, decl.span, "cannot read import '{s}' (not a file or directory)", .{resolved_path});
                    }
                    return error.ImportError;
                };
                try cache.put(resolved_path, result);
                break :blk result;
            }
        };

        if (imp.name) |ns_name| {
            try mod.addNamespace(allocator, &decl_list, ns_name, imported_mod, decl.span);
        } else {
            try mod.mergeFlat(allocator, &decl_list, imported_mod);
        }
    }

    try mod.finalize(allocator, &decl_list);
    return mod;
}

/// Resolve a directory import by aggregating all .sx files in the directory.
fn resolveDirectoryImport(
    allocator: std.mem.Allocator,
    io: std.Io,
    dir_path: []const u8,
    chain: *std.StringHashMap(void),
    cache: *ModuleCache,
    source_map: ?*std.StringHashMap([:0]const u8),
    diagnostics: ?*errors.DiagnosticList,
    span: ast.Span,
    stdlib_paths: []const []const u8,
    import_graph: ?*std.StringHashMap(std.StringHashMap(void)),
) anyerror!ResolvedModule {
    // Open the directory with iteration capability
    const dir = std.Io.Dir.openDir(.cwd(), io, dir_path, .{ .iterate = true }) catch {
        return error.ImportError;
    };
    defer dir.close(io);

    // Collect all .sx file names
    var file_names = std.ArrayList([]const u8).empty;
    var it = dir.iterate();
    while (it.next(io) catch null) |entry| {
        if (entry.kind != .file) continue;
        if (!std.mem.endsWith(u8, entry.name, ".sx")) continue;
        const name_copy = try allocator.dupe(u8, entry.name);
        try file_names.append(allocator, name_copy);
    }

    // Sort alphabetically for deterministic ordering
    std.mem.sort([]const u8, file_names.items, {}, struct {
        fn lessThan(_: void, a: []const u8, b: []const u8) bool {
            return std.mem.order(u8, a, b) == .lt;
        }
    }.lessThan);

    // Add directory to chain for circular import detection
    try chain.put(dir_path, {});
    defer _ = chain.remove(dir_path);

    // Merge all files into a combined module
    var combined = ResolvedModule{
        .path = dir_path,
        .decls = &.{},
        .scope = std.StringHashMap(void).init(allocator),
    };
    var decl_list = std.ArrayList(*Node).empty;

    for (file_names.items) |file_name| {
        const file_path = try std.fmt.allocPrint(allocator, "{s}/{s}", .{ dir_path, file_name });

        if (chain.contains(file_path)) continue;

        const file_mod = if (cache.get(file_path)) |cached|
            cached
        else file_blk: {
            const imp_bytes = std.Io.Dir.readFileAlloc(.cwd(), io, file_path, allocator, .limited(10 * 1024 * 1024)) catch {
                if (diagnostics) |diags| {
                    diags.addFmt(.err, span, "cannot read '{s}' in directory import", .{file_path});
                }
                return error.ImportError;
            };
            const imp_source = try allocator.dupeZ(u8, imp_bytes);

            if (source_map) |sm| {
                sm.put(file_path, imp_source) catch {};
            }

            var p = parser.Parser.init(allocator, imp_source);
            const imp_root = p.parse() catch {
                if (diagnostics) |diags| {
                    diags.addFmt(.err, span, "parse error in '{s}': {s}", .{ file_path, p.err_msg orelse "unknown" });
                }
                return error.ImportError;
            };

            try chain.put(file_path, {});
            const result = try resolveImports(allocator, io, imp_root, dir_path, file_path, chain, cache, source_map, diagnostics, stdlib_paths, import_graph);
            _ = chain.remove(file_path);

            try cache.put(file_path, result);
            break :file_blk result;
        };

        try combined.mergeFlat(allocator, &decl_list, file_mod);
    }

    try combined.finalize(allocator, &decl_list);
    return combined;
}