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

let
  cfg = config.modelDownloads;

  script = pkgs.writeShellScript "download-models" ''
    set -eu
    mkdir -p "${cfg.dir}"

    for url in ${lib.concatStringsSep " " (map lib.escapeShellArg cfg.urls)}; do
      name="$(basename "$url")"
      out="${cfg.dir}/$name"

      if [ -f "$out" ]; then
        echo "OK: $name already exists"
        continue
      fi

      echo "Downloading $name"
      tmp="$out.part"
      ${pkgs.curl}/bin/curl -L --fail --retry 3 --retry-delay 2 -o "$tmp" "$url"
      mv -f "$tmp" "$out"
    done
  '';
in
{
  options.modelDownloads = {
    enable = lib.mkEnableOption "Download model files into a persistent directory";

    dir = lib.mkOption {
      type = lib.types.str;
      default = "/data/ai/models";
    };

    urls = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "List of direct download URLs (filenames are taken from the URL).";
    };
  };

  config = lib.mkIf cfg.enable {
    home.packages = [ pkgs.curl ];

    home.activation.modelDownloads = lib.hm.dag.entryAfter [ "writeBoundary" ] ''
      ${script}
    '';
  };
}