const std = @import("std");

pub fn build(b: *std.Build) void {
    const target = b.standardTargetOptions(.{});
    const optimize = b.standardOptimizeOption(.{});

    const vosk_dep_name = selectVoskDependency(target.result);
    const vosk_dep = b.dependency(vosk_dep_name, .{});
    const alsa_dep = b.dependency("alsa", .{
        .target = target,
        .optimize = optimize,
    });

    // We need to use curl for this as the domain doesn't work with zig TLS
    const model_step = ModelDownloadStep.create(b);

    // Install the model to the output directory
    const install_model = b.addInstallDirectory(.{
        .source_dir = model_step.getOutputPath(),
        .install_dir = .bin,
        .install_subdir = "vosk-model-small-en-us-0.15",
    });
    install_model.step.dependOn(&model_step.step);
    b.getInstallStep().dependOn(&install_model.step);

    // Create the STT library
    const stt_lib = b.addLibrary(.{
        .name = "stt",
        .linkage = .static,
        .root_module = b.createModule(.{
            .root_source_file = b.path("src/root.zig"),
            .target = target,
            .optimize = optimize,
            .link_libc = true,
        }),
    });

    // Link with Vosk library
    stt_lib.addIncludePath(vosk_dep.path(""));
    stt_lib.addLibraryPath(vosk_dep.path(""));
    stt_lib.linkSystemLibrary("vosk");

    const alsa_lib = alsa_dep.artifact("asound");
    stt_lib.linkLibrary(alsa_lib);
    stt_lib.addIncludePath(alsa_dep.path("zig-out/include"));

    b.installArtifact(stt_lib);

    // Create the demo executable
    const exe = b.addExecutable(.{
        .name = "stt-demo",
        .root_module = b.createModule(.{
            .root_source_file = b.path("src/main.zig"),
            .target = target,
            .optimize = optimize,
            .link_libc = true,
        }),
    });

    exe.linkLibrary(stt_lib);
    exe.linkLibrary(alsa_lib);
    exe.addIncludePath(alsa_dep.path("zig-out/include"));

    // Link with Vosk for the executable
    exe.addIncludePath(vosk_dep.path(""));
    exe.addLibraryPath(vosk_dep.path(""));
    exe.linkSystemLibrary("vosk");

    b.installArtifact(exe);

    const run_step = b.step("run", "Run the app");
    const run_cmd = b.addRunArtifact(exe);
    run_step.dependOn(&run_cmd.step);
    run_cmd.step.dependOn(b.getInstallStep());

    if (b.args) |args| {
        run_cmd.addArgs(args);
    }
    // Creates a step for unit testing the library
    const lib_unit_tests = b.addTest(.{
        .root_module = b.createModule(.{
            .root_source_file = b.path("src/root.zig"),
            .target = target,
            .optimize = optimize,
            .link_libc = true,
        }),
    });

    // Link the same dependencies as the library
    lib_unit_tests.linkLibrary(alsa_lib);
    lib_unit_tests.addIncludePath(alsa_dep.path("zig-out/include"));
    lib_unit_tests.addIncludePath(vosk_dep.path(""));
    lib_unit_tests.addLibraryPath(vosk_dep.path(""));
    lib_unit_tests.linkSystemLibrary("vosk");

    const run_lib_unit_tests = b.addRunArtifact(lib_unit_tests);

    // Creates a step for unit testing the dedicated test file
    const dedicated_unit_tests = b.addTest(.{
        .root_module = b.createModule(.{
            .root_source_file = b.path("src/test.zig"),
            .target = target,
            .optimize = optimize,
            .link_libc = true,
        }),
    });

    // Link the same dependencies as the library for dedicated tests
    dedicated_unit_tests.linkLibrary(alsa_lib);
    dedicated_unit_tests.addIncludePath(alsa_dep.path("zig-out/include"));
    dedicated_unit_tests.addIncludePath(vosk_dep.path(""));
    dedicated_unit_tests.addLibraryPath(vosk_dep.path(""));
    dedicated_unit_tests.linkSystemLibrary("vosk");

    const run_dedicated_unit_tests = b.addRunArtifact(dedicated_unit_tests);

    // Creates a step for unit testing the demo application
    const exe_unit_tests = b.addTest(.{
        .root_module = b.createModule(.{
            .root_source_file = b.path("src/main.zig"),
            .target = target,
            .optimize = optimize,
            .link_libc = true,
        }),
    });

    exe_unit_tests.linkLibrary(stt_lib);
    exe_unit_tests.linkLibrary(alsa_lib);
    exe_unit_tests.addIncludePath(alsa_dep.path("zig-out/include"));
    exe_unit_tests.addIncludePath(vosk_dep.path(""));
    exe_unit_tests.addLibraryPath(vosk_dep.path(""));
    exe_unit_tests.linkSystemLibrary("vosk");

    const run_exe_unit_tests = b.addRunArtifact(exe_unit_tests);

    // Test step that runs all unit tests
    const test_step = b.step("test", "Run unit tests");
    test_step.dependOn(&run_lib_unit_tests.step);
    test_step.dependOn(&run_dedicated_unit_tests.step);
    test_step.dependOn(&run_exe_unit_tests.step);
}

fn selectVoskDependency(target: std.Target) []const u8 {
    return switch (target.cpu.arch) {
        .x86_64 => "vosk_linux_x86_64",
        .aarch64 => "vosk_linux_aarch64",
        .riscv64 => "vosk_linux_riscv64",
        else => @panic("Unsupported architecture - only x86_64, aarch64, and riscv64 are supported"),
    };
}

const ModelDownloadStep = struct {
    step: std.Build.Step,
    builder: *std.Build,

    pub fn create(builder: *std.Build) *ModelDownloadStep {
        const self = builder.allocator.create(ModelDownloadStep) catch @panic("OOM");
        self.* = .{
            .step = std.Build.Step.init(.{
                .id = .custom,
                .name = "download-model",
                .owner = builder,
                .makeFn = make,
            }),
            .builder = builder,
        };
        return self;
    }

    pub fn getOutputPath(self: *ModelDownloadStep) std.Build.LazyPath {
        var hasher = std.hash.Wyhash.init(0);
        hasher.update("https://alphacephei.com/vosk/models/vosk-model-small-en-us-0.15.zip");
        const cache_hash = hasher.final();

        var cache_dir_buf: [std.fs.max_path_bytes]u8 = undefined;
        const cache_dir = std.fmt.bufPrint(&cache_dir_buf, "{s}/o/{x}/vosk-model-small-en-us-0.15", .{ self.builder.cache_root.path.?, cache_hash }) catch @panic("path too long");

        return .{ .cwd_relative = self.builder.allocator.dupe(u8, cache_dir) catch @panic("OOM") };
    }

    fn make(step: *std.Build.Step, options: std.Build.Step.MakeOptions) anyerror!void {
        _ = options;
        const self: *ModelDownloadStep = @fieldParentPtr("step", step);

        const model_dir = "vosk-model-small-en-us-0.15";

        // Create a cache hash based on the URL
        var hasher = std.hash.Wyhash.init(0);
        hasher.update("https://alphacephei.com/vosk/models/vosk-model-small-en-us-0.15.zip");
        const cache_hash = hasher.final();

        var cache_dir_buf: [std.fs.max_path_bytes]u8 = undefined;
        const cache_dir = std.fmt.bufPrint(&cache_dir_buf, "{s}/o/{x}", .{ self.builder.cache_root.path.?, cache_hash }) catch @panic("path too long");

        const cached_model_dir = std.fmt.allocPrint(self.builder.allocator, "{s}/{s}", .{ cache_dir, model_dir }) catch @panic("OOM");
        defer self.builder.allocator.free(cached_model_dir);

        // Check if already cached
        if (std.fs.cwd().access(cached_model_dir, .{})) |_| {
            step.result_cached = true;
            return;
        } else |_| {}

        // Not cached, need to download
        std.fs.cwd().makePath(cache_dir) catch @panic("Could not create cache directory");

        const model_zip = std.fmt.allocPrint(self.builder.allocator, "{s}/model.zip", .{cache_dir}) catch @panic("OOM");
        defer self.builder.allocator.free(model_zip);

        // Download
        const download_result = std.process.Child.run(.{
            .allocator = self.builder.allocator,
            .argv = &.{ "curl", "-s", "-o", model_zip, "https://alphacephei.com/vosk/models/vosk-model-small-en-us-0.15.zip" },
        }) catch return error.DownloadFailed;
        if (download_result.term.Exited != 0) return error.DownloadFailed;

        // Extract to cache using stdlib
        var zip_file = std.fs.cwd().openFile(model_zip, .{}) catch return error.UnzipFailed;
        defer zip_file.close();

        var cache_dir_handle = std.fs.cwd().openDir(cache_dir, .{}) catch return error.UnzipFailed;
        defer cache_dir_handle.close();

        var zip_file_buffer: [4096]u8 = undefined;
        var zip_file_reader = zip_file.reader(&zip_file_buffer);

        std.zip.extract(cache_dir_handle, &zip_file_reader, .{}) catch return error.UnzipFailed;

        step.result_cached = false;
    }
};