const std = @import("std");
const ast = @import("ast.zig");
const parser = @import("parser.zig");
const imports = @import("imports.zig");
const sema = @import("sema.zig");
const errors = @import("errors.zig");
const c_import = @import("c_import.zig");
const ir = @import("ir/ir.zig");
const target_mod = @import("target.zig");
const Node = ast.Node;

pub const TargetConfig = target_mod.TargetConfig;
pub const JniMainEmission = target_mod.JniMainEmission;

pub const Compilation = struct {
    allocator: std.mem.Allocator,
    io: std.Io,
    file_path: []const u8,
    source: [:0]const u8,
    diagnostics: errors.DiagnosticList,
    target_config: TargetConfig,
    stdlib_paths: []const []const u8 = &.{},

    // Pipeline results
    root: ?*Node = null,
    resolved_root: ?*Node = null,
    import_sources: std.StringHashMap([:0]const u8),
    module_scopes: std.StringHashMap(std.StringHashMap(void)),
    import_graph: std.StringHashMap(std.StringHashMap(void)),
    sema_result: ?sema.SemaResult = null,
    ir_emitter: ?ir.LLVMEmitter = null,
    /// Lowered IR module, kept alive past `generateCode` so post-link
    /// callbacks can re-enter the interpreter to invoke sx functions
    /// (e.g. `platform.bundle.bundle_main` after `target.link`).
    ir_module: ?*ir.Module = null,
    /// C sources requested by the lowering pass (not in the user's AST).
    /// E.g. the JNI env TL runtime when `#jni_env` is used. Merged with
    /// AST sources in `collectCImportSources`.
    lowering_extra_c_sources: std.ArrayList(c_import.CImportInfo) = .empty,
    /// `#jni_main #jni_class("...")` declarations whose Java sources were
    /// rendered during lowering. Surfaced to the sx Android bundler
    /// (`library/modules/platform/bundle.sx`) via `BuildConfig.jni_main_*`
    /// in `compiler_hooks.zig`; the bundler writes `.java` files + runs
    /// `javac` + `d8` + bundles `classes.dex` into the APK.
    lowering_jni_main_decls: std.ArrayList(JniMainEmission) = .empty,

    pub fn init(allocator: std.mem.Allocator, io: std.Io, file_path: []const u8, source: [:0]const u8, target_config: TargetConfig, stdlib_paths: []const []const u8) Compilation {
        return .{
            .allocator = allocator,
            .io = io,
            .file_path = file_path,
            .source = source,
            .diagnostics = errors.DiagnosticList.init(allocator, source, file_path),
            .import_sources = std.StringHashMap([:0]const u8).init(allocator),
            .module_scopes = std.StringHashMap(std.StringHashMap(void)).init(allocator),
            .import_graph = std.StringHashMap(std.StringHashMap(void)).init(allocator),
            .target_config = target_config,
            .stdlib_paths = stdlib_paths,
        };
    }

    pub fn deinit(self: *Compilation) void {
        if (self.ir_emitter) |*e| e.deinit();
        if (self.ir_module) |m| {
            m.deinit();
            self.allocator.destroy(m);
        }
        self.diagnostics.deinit();
    }

    pub fn parse(self: *Compilation) !void {
        var p = parser.Parser.init(self.allocator, self.source);
        p.diagnostics = &self.diagnostics;
        self.root = p.parse() catch return error.CompileError;
    }

    /// Derive the comptime evaluation context (OS / ARCH / POINTER_SIZE
    /// values) from the build target. Used by `imports.resolveImports`
    /// to hoist top-level `inline if OS == .X { ... }` body decls
    /// before resolution; mirrors `injectComptimeConstants` in lowering.
    fn comptimeContext(self: *const Compilation) imports.ComptimeContext {
        const tc = self.target_config;
        const os: []const u8 = if (tc.isWasm()) "wasm" else if (tc.isWindows()) "windows" else if (tc.isAndroid()) "android" else if (tc.isLinux()) "linux" else if (tc.isIOS()) "ios" else if (tc.isMacOS()) "macos" else "unknown";
        const arch: []const u8 = if (tc.isWasm32()) "wasm32" else if (tc.isWasm64()) "wasm64" else if (tc.isAarch64()) "aarch64" else if (tc.isX86_64()) "x86_64" else "unknown";
        const ptr_size: i64 = if (tc.isWasm32()) 4 else 8;
        return .{ .os = os, .arch = arch, .pointer_size = ptr_size };
    }

    pub fn resolveImports(self: *Compilation) !void {
        const root = self.root orelse return error.CompileError;
        var chain = std.StringHashMap(void).init(self.allocator);
        var cache = imports.ModuleCache.init(self.allocator);
        const base_dir = imports.dirName(self.file_path);
        const mod = imports.resolveImports(
            self.allocator,
            self.io,
            root,
            base_dir,
            self.file_path,
            &chain,
            &cache,
            &self.import_sources,
            &self.diagnostics,
            self.stdlib_paths,
            &self.import_graph,
            self.comptimeContext(),
        ) catch return error.CompileError;

        // Preserve per-module visibility scopes for C import access checking
        self.module_scopes.put(self.file_path, mod.scope) catch {};
        var cache_it = cache.iterator();
        while (cache_it.next()) |entry| {
            self.module_scopes.put(entry.key_ptr.*, entry.value_ptr.scope) catch {};
        }

        // Store main file source in import_sources so error reporting can find it
        self.import_sources.put(self.file_path, self.source) catch {};

        // Wire import_sources to diagnostics for file-aware error rendering
        self.diagnostics.import_sources = &self.import_sources;

        // Build a root node from the resolved module's decls
        const new_root = try self.allocator.create(Node);
        new_root.* = .{
            .span = root.span,
            .data = .{ .root = .{ .decls = mod.decls } },
        };
        self.resolved_root = new_root;
    }

    pub fn analyze(self: *Compilation) !void {
        const root = self.resolved_root orelse self.root orelse return error.CompileError;
        var analyzer = sema.Analyzer.init(self.allocator);
        self.sema_result = analyzer.analyze(root) catch return error.CompileError;
        // Merge sema diagnostics into our list
        if (self.sema_result) |sr| {
            for (sr.diagnostics) |d| {
                self.diagnostics.add(d.level, d.message, d.span);
            }
        }
    }

    /// Generate code via the IR pipeline: lower AST → IR → LLVM.
    pub fn generateCode(self: *Compilation) !void {
        // Heap-allocate the IR module so its address is stable during emit
        const ir_mod_ptr = try self.allocator.create(ir.Module);
        ir_mod_ptr.* = try self.lowerToIR();
        var emitter = ir.LLVMEmitter.init(self.allocator, ir_mod_ptr, "sx_module", self.target_config);
        emitter.emit();
        // Keep the IR module alive past LLVM emission so post-link
        // callbacks can re-enter the interpreter via `invokeByName`.
        self.ir_module = ir_mod_ptr;
        self.ir_emitter = emitter;
    }

    /// Re-enter the IR interpreter after `generateCode` (and after linking,
    /// if applicable) to invoke a named sx function. Used for the post-link
    /// bundling callback. Returns the function's return value, or null if the
    /// name doesn't resolve to a function in the lowered module.
    pub fn invokeByName(self: *Compilation, name: []const u8, args: []const ir.Value) !?ir.Value {
        const mod = self.ir_module orelse return null;
        var found_id: ?ir.FuncId = null;
        for (mod.functions.items, 0..) |func, i| {
            const fname = mod.types.getString(func.name);
            if (std.mem.eql(u8, fname, name)) {
                found_id = ir.FuncId.fromIndex(@intCast(i));
                break;
            }
        }
        const fid = found_id orelse return null;
        return try self.invokeByFuncId(fid, args);
    }

    /// Re-enter the IR interpreter and call a previously-resolved function
    /// id. Companion to `invokeByName` — used when the FuncId was captured
    /// at `#run` time (e.g. by `set_post_link_callback`) and we want to
    /// invoke it later without name lookup.
    pub fn invokeByFuncId(self: *Compilation, id: ir.FuncId, args: []const ir.Value) !ir.Value {
        const mod = self.ir_module orelse return error.NoIRModule;
        var interp = ir.Interpreter.init(mod, self.allocator);
        defer interp.deinit();
        if (self.ir_emitter) |*e| interp.build_config = &e.build_config;
        ir.Interpreter.last_bail_op = null;
        ir.Interpreter.last_bail_builtin = null;
        ir.Interpreter.last_bail_detail = null;
        const result = interp.call(id, args) catch |err| {
            flushInterpOutput(interp.output.items);
            return err;
        };
        flushInterpOutput(interp.output.items);
        return result;
    }

    /// #run / post-link callback `print` output lands here. Routes to
    /// fd 1 (stdout) so it joins the JIT-executed runtime's output
    /// stream — the user wrote `print(...)` in both call sites, so
    /// the stream split is invisible to them. issue-0047.
    fn flushInterpOutput(bytes: []const u8) void {
        if (bytes.len == 0) return;
        _ = std.c.write(1, bytes.ptr, bytes.len);
    }

    /// Get link flags accumulated from #run build blocks.
    pub fn getBuildLinkFlags(self: *Compilation) []const []const u8 {
        if (self.ir_emitter) |*e| return e.build_config.link_flags.items;
        return &.{};
    }

    /// Get frameworks accumulated from #run build blocks (BuildOptions.add_framework).
    pub fn getBuildFrameworks(self: *Compilation) []const []const u8 {
        if (self.ir_emitter) |*e| return e.build_config.frameworks.items;
        return &.{};
    }

    /// Get output path set from #run build blocks, if any.
    pub fn getBuildOutputPath(self: *Compilation) ?[]const u8 {
        if (self.ir_emitter) |*e| return e.build_config.output_path;
        return null;
    }

    /// Get custom WASM shell template path set from #run build blocks, if any.
    pub fn getBuildWasmShell(self: *Compilation) ?[]const u8 {
        if (self.ir_emitter) |*e| return e.build_config.wasm_shell_path;
        return null;
    }

    /// Get the post-link callback function id (set via
    /// `BuildOptions.set_post_link_callback(fn)`), if any.
    pub fn getPostLinkCallback(self: *Compilation) ?ir.FuncId {
        if (self.ir_emitter) |*e| return e.build_config.post_link_callback_fn;
        return null;
    }

    /// Get the post-link module name (set via
    /// `BuildOptions.set_post_link_module("name")`), if any.
    pub fn getPostLinkModule(self: *Compilation) ?[]const u8 {
        if (self.ir_emitter) |*e| return e.build_config.post_link_module;
        return null;
    }

    /// Collect C import source info — both from user-written `#import c { ... }`
    /// blocks in the AST AND from lowering-time auto-injections (currently:
    /// the JNI env TL runtime when `#jni_env` / `#jni_call`-with-omitted-env
    /// is used). The lower-side auto-injections live in
    /// `lowering_extra_c_sources` and are populated by `lowerToIR` based on
    /// `Lowering.needs_jni_env_tl_runtime` etc.
    pub fn collectCImportSources(self: *Compilation) ![]c_import.CImportInfo {
        const root = self.resolved_root orelse self.root orelse return &.{};
        const ast_sources = try c_import.collectCImportSources(self.allocator, root);
        if (self.lowering_extra_c_sources.items.len == 0) return ast_sources;
        var merged = std.ArrayList(c_import.CImportInfo).empty;
        try merged.appendSlice(self.allocator, ast_sources);
        try merged.appendSlice(self.allocator, self.lowering_extra_c_sources.items);
        return merged.toOwnedSlice(self.allocator);
    }

    /// Resolve a stdlib-relative path through the configured `stdlib_paths`.
    /// Returns the first candidate whose absolute path resolves to an
    /// existing file. Used by lower-side auto-injected C sources.
    fn resolveStdlibPath(self: *Compilation, rel: []const u8) !?[]const u8 {
        for (self.stdlib_paths) |root_path| {
            const candidate = try std.fmt.allocPrint(self.allocator, "{s}/{s}", .{ root_path, rel });
            if (std.Io.Dir.readFileAlloc(.cwd(), self.io, candidate, self.allocator, .limited(1024 * 1024))) |buf| {
                self.allocator.free(buf);
                return candidate;
            } else |_| {
                self.allocator.free(candidate);
            }
        }
        return null;
    }

    /// Lower the parsed AST to the sx IR module (shadow pipeline).
    pub fn lowerToIR(self: *Compilation) !ir.Module {
        const root = self.resolved_root orelse self.root orelse return ir.Module.init(self.allocator);
        var module = ir.Module.init(self.allocator);
        //TODO: find a better place for this
        if (self.target_config.isWasm32()) {
            module.types.pointer_size = 4;
        }
        var lowering = ir.Lowering.init(&module);
        lowering.main_file = self.file_path;
        lowering.resolved_root = root;
        lowering.target_config = self.target_config;
        lowering.diagnostics = &self.diagnostics;
        lowering.module_scopes = &self.module_scopes;
        lowering.import_graph = &self.import_graph;
        lowering.lowerRoot(root);
        if (self.diagnostics.hasErrors()) return error.CompileError;

        // Auto-link the JNI env TL runtime when lowering used it. The .c file
        // ships with the sx library; we resolve it through stdlib_paths so
        // consumers don't need to vendor a copy.
        if (lowering.needs_jni_env_tl_runtime) {
            if (try self.resolveStdlibPath("vendors/sx_jni_runtime/sx_jni_env_tl.c")) |abs_path| {
                var sources = std.ArrayList([]const u8).empty;
                try sources.append(self.allocator, abs_path);
                try self.lowering_extra_c_sources.append(self.allocator, .{
                    .sources = try sources.toOwnedSlice(self.allocator),
                    .includes = &.{},
                    .defines = &.{},
                    .flags = &.{},
                });
            }
        }

        try self.collectJniMainEmissions(&lowering);

        return module;
    }

    /// Walk `lowering.foreign_class_map` and render Java sources for every
    /// `#jni_main #jni_class("...")` declaration. Renders happen here so the
    /// AST + class-registry snapshot stay confined to the lowering pass; the
    /// downstream APK pipeline only needs `{foreign_path, java_source}` pairs.
    fn collectJniMainEmissions(self: *Compilation, lowering: *ir.Lowering) !void {
        // `foreign_class_map` registers each decl under bare + qualified names —
        // dedupe by foreign_path so a single decl emits one .java.
        var seen = std.StringHashMap(void).init(self.allocator);
        defer seen.deinit();

        // Class registry passed to jni_java_emit for `*Foo` cross-class refs
        // and `#extends Alias` resolution.
        var registry = std.StringHashMap([]const u8).init(self.allocator);
        defer registry.deinit();
        var it_reg = lowering.foreign_class_map.iterator();
        while (it_reg.next()) |entry| {
            try registry.put(entry.key_ptr.*, entry.value_ptr.*.foreign_path);
        }

        // Derive the `System.loadLibrary` argument from the `-o` basename
        // (e.g. `/tmp/libsxchess.so` → `sxchess`). When `-o` is unset the
        // emitter omits the static init block; the user must then arrange
        // .so loading via another class.
        const lib_name = libNameFromOutputPath(self.target_config.output_path);

        var it = lowering.foreign_class_map.iterator();
        while (it.next()) |entry| {
            const fcd = entry.value_ptr.*;
            if (!fcd.is_main) continue;
            if (fcd.is_foreign) continue;
            if (fcd.runtime != .jni_class) continue;
            if (seen.contains(fcd.foreign_path)) continue;
            try seen.put(fcd.foreign_path, {});

            const java_source = try ir.jni_java_emit.emitJavaSource(self.allocator, fcd, .{
                .classes = &registry,
                .lib_name = lib_name,
            });
            try self.lowering_jni_main_decls.append(self.allocator, .{
                .foreign_path = try self.allocator.dupe(u8, fcd.foreign_path),
                .java_source = java_source,
            });
        }
    }

    /// `/path/to/libfoo.so` → `foo`. Anything else → null (caller skips
    /// emitting the `System.loadLibrary` init block).
    fn libNameFromOutputPath(output_path: ?[]const u8) ?[]const u8 {
        const path = output_path orelse return null;
        const basename = std.fs.path.basename(path);
        if (!std.mem.startsWith(u8, basename, "lib")) return null;
        if (!std.mem.endsWith(u8, basename, ".so")) return null;
        return basename[3 .. basename.len - 3];
    }

    /// Java sources rendered from `#jni_main #jni_class("...")` decls during
    /// lowering. Empty unless `lowerToIR` has run.
    pub fn getJniMainEmissions(self: *const Compilation) []const JniMainEmission {
        return self.lowering_jni_main_decls.items;
    }

    pub fn renderErrors(self: *const Compilation) void {
        self.diagnostics.renderStderr();
    }
};