{ config, lib, pkgs, ... }:

let
  cfg = config.programs.llama-cpp;

  defaultArgsBash =
    lib.concatStringsSep " " (map lib.escapeShellArg cfg.defaultArgs);

  globalEosTextBash =
    lib.escapeShellArg (lib.concatStringsSep "\n" cfg.eosStrings);

  presetBash =
    lib.concatMapStringsSep "\n" (p: ''
      PRESET_PATHS[${lib.escapeShellArg p.name}]=${lib.escapeShellArg p.path}
      PRESET_ALIAS[${lib.escapeShellArg p.name}]=${lib.escapeShellArg (p.alias or p.name)}
      PRESET_EOS_TEXT[${lib.escapeShellArg p.name}]=${lib.escapeShellArg (lib.concatStringsSep "\n" (p.eosStrings or []))}
    '') cfg.presets;

  template = builtins.readFile ./llama-serve.sh;

  scriptText = lib.replaceStrings
    [
      "@server_bin@"
      "@default_args@"
      "@global_eos_text@"
      "@preset_bash@"
    ]
    [
      "${cfg.package}/bin/${cfg.serverBinary}"
      defaultArgsBash
      globalEosTextBash
      presetBash
    ]
    template;

  wrapper = pkgs.writeShellScriptBin cfg.wrapperName scriptText;

in {
  options.programs.llama-cpp = {
    enable = lib.mkEnableOption "llama.cpp wrapper for llama-server";

    package = lib.mkPackageOption pkgs "llama-cpp" { };

    serverBinary = lib.mkOption {
      type = lib.types.str;
      default = "llama-server";
    };

    wrapperName = lib.mkOption {
      type = lib.types.str;
      default = "llama-serve";
    };

    defaultArgs = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Args always passed to llama-server before user args.";
    };

    eosStrings = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
    };

    presets = lib.mkOption {
      type = lib.types.listOf (lib.types.submodule ({ ... }: {
        options = {
          name = lib.mkOption { type = lib.types.str; };
          path = lib.mkOption { type = lib.types.str; };
          alias = lib.mkOption { type = lib.types.nullOr lib.types.str; default = null; };
          eosStrings = lib.mkOption { type = lib.types.listOf lib.types.str; default = [ ]; };
        };
      }));
      default = [ ];
    };
  };

  config = lib.mkIf cfg.enable {
    home.packages = [
      cfg.package
      wrapper
    ];
  };
}