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 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,

    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();
        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;
    }

    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,
        ) 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();
        // IR module is no longer needed after LLVM IR has been generated
        ir_mod_ptr.deinit();
        self.allocator.destroy(ir_mod_ptr);
        self.ir_emitter = emitter;
    }

    /// 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;
    }

    /// Collect C import source info from the resolved AST.
    pub fn collectCImportSources(self: *Compilation) ![]c_import.CImportInfo {
        const root = self.resolved_root orelse self.root orelse return &.{};
        return c_import.collectCImportSources(self.allocator, root);
    }

    /// 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;
        return module;
    }

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