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

let
  cfg = config.services.hyperspace;

  hyperspacePkg = pkgs.stdenv.mkDerivation {
    name = "hyperspace-pods-${cfg.version}";
    src = pkgs.fetchurl {
      url = "https://github.com/hyperspaceai/aios-cli/releases/download/v${cfg.version}/aios-cli-x86_64-unknown-linux-gnu.tar.gz";
      hash = cfg.packageHash;
    };
    sourceRoot = ".";
    installPhase = ''
      mkdir -p $out/libexec $out/bin
      cp -r * $out/libexec/
      chmod +x $out/libexec/aios-cli
      ln -s $out/libexec/aios-cli $out/bin/hyperspace
    '';
  };
in {
  options.services.hyperspace = {
    enable = lib.mkEnableOption "Hyperspace Pods P2P AI cluster agent";

    version = lib.mkOption {
      type = lib.types.str;
      default = "5.45.30";
      description = "Hyperspace CLI version to download.";
    };

    packageHash = lib.mkOption {
      type = lib.types.str;
      default = "sha256-f6fJ8t3exqtYwUD5j+WvD+Hm0oN/Eef0X+R9Rj23dE0=";
      description = ''
        SRI hash of the hyperspace release tarball (sha256-<base64>).
        Must be updated when version changes. Generate with:
        nix store prefetch-file --hash-algo sha256 \\
          https://github.com/hyperspaceai/aios-cli/releases/download/v{version}/aios-cli-x86_64-unknown-linux-gnu.tar.gz
      '';
    };

    user = lib.mkOption {
      type = lib.types.str;
      default = "ai-worker";
      description = "System user to run the Hyperspace agent.";
    };

    apiPort = lib.mkOption {
      type = lib.types.port;
      default = 8080;
      description = "OpenAI-compatible API port (configurable via --api-port).";
    };

    profile = lib.mkOption {
      type = lib.types.str;
      default = "auto";
      description = ''
        Agent profile. Options: auto (auto-detect hardware), full (all capabilities),
        inference (GPU inference only), embedding (CPU embedding only),
        relay (lightweight relay), storage (storage + memory).
      '';
    };

    autoStart = lib.mkOption {
      type = lib.types.bool;
      default = true;
      description = "Start the agent automatically on boot.";
    };

    openFirewall = lib.mkOption {
      type = lib.types.bool;
      default = true;
      description = "Open P2P mesh (4001 TCP+UDP, 30301 TCP) and API port in the firewall.";
    };

    extraArgs = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Extra arguments to pass to 'hyperspace start'.";
    };
  };

  config = lib.mkIf cfg.enable {
    systemd.services.hyperspace = {
      description = "Hyperspace Pods P2P AI Cluster Agent";
      after = [ "network.target" "network-online.target" ];
      wants = [ "network-online.target" ];
      wantedBy = lib.mkIf cfg.autoStart [ "multi-user.target" ];

      path = with pkgs; [ bash coreutils ];

      serviceConfig = {
        Type = "simple";
        User = cfg.user;
        Group = cfg.user;
        WorkingDirectory = "${hyperspacePkg}/libexec";
        ExecStart = "${hyperspacePkg}/bin/hyperspace start --profile ${cfg.profile} --api-port ${toString cfg.apiPort} ${lib.escapeShellArgs cfg.extraArgs}";
        Restart = "on-failure";
        RestartSec = 5;

        # AMD MI50 (ROCm) device access
        DeviceAllow = [ "/dev/kfd rw" "/dev/dri rw" ];

        # Supplementary groups for GPU/accelerator access
        SupplementaryGroups = [ "video" "render" ];

        # Hardening
        NoNewPrivileges = true;
        ProtectHome = "tmpfs";
        ProtectSystem = "strict";
        PrivateTmp = true;
        PrivateDevices = false; # Needs /dev/kfd and /dev/dri
      };

      environment = {
        HSA_OVERRIDE_GFX_VERSION = "9.0.6";
        HOME = "/home/${cfg.user}";
      };
    };

    # Firewall ports for P2P mesh (libp2p 4001, chain 30301) and API
    networking.firewall.allowedTCPPorts = lib.mkIf cfg.openFirewall [ 4001 30301 cfg.apiPort ];
    networking.firewall.allowedUDPPorts = lib.mkIf cfg.openFirewall [ 4001 ];

    # Add GPU/accelerator groups to the service user (persistent beyond service restarts)
    users.users = lib.mkIf (cfg.user == "ai-worker") {
      ai-worker = {
        extraGroups = [ "video" "render" ];
      };
    };

    # ROCm override for AMD MI50 (gfx906) compatibility
    environment.variables.HSA_OVERRIDE_GFX_VERSION = "9.0.6";
  };
}