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

# Aim — ML experiment tracking server.
#
# Auth model: none on the server itself.
#   - Web UI: protected by Caddy → Authelia two_factor (admins only).
#   - Python SDK: connects via aim://192.168.1.100:53800 directly on LAN;
#     the firewall restricts port 53800 to 192.168.1.0/24 only.
#
# Two containers share one volume:
#   aim-server  — tracking protocol server (aim:// SDK connections)
#   aim-ui      — web UI + REST API (behind Caddy)
#
# Volume layout:
#   <dataDir>/aim/  → /aim  in both containers (shared .aim data store)

let
  cfg     = config.homey.aim;
  dataDir = config.homey.storage.mountPoint;
  domain  = homeyConfig.domain;
in
{
  options.homey.aim = {
    enable = lib.mkEnableOption "Aim experiment tracking server";

    image = lib.mkOption {
      type    = lib.types.str;
      default = "aimstack/aim:latest";
    };

    uiPort = lib.mkOption {
      type        = lib.types.port;
      default     = 43800;
      description = "Host port for the Aim web UI (bound to 127.0.0.1, Caddy proxy).";
    };

    serverPort = lib.mkOption {
      type        = lib.types.port;
      default     = 53800;
      description = "Host port for the Aim tracking server (aim:// protocol, LAN-accessible).";
    };
  };

  config = lib.mkIf cfg.enable {
    # -------------------------------------------------------------------------
    # Tracking server — aim:// SDK protocol, LAN only
    # -------------------------------------------------------------------------
    virtualisation.oci-containers.containers.aim-server = {
      image = cfg.image;
      # Bound to 0.0.0.0 so LAN training scripts can reach it directly.
      # Port 53800 is restricted to 192.168.1.0/24 via extraCommands below.
      ports = [ "0.0.0.0:${toString cfg.serverPort}:53800" ];
      volumes  = [ "${dataDir}/aim:/aim" ];
      extraOptions = [ "--network=homey" ];
      cmd = [ "aim" "server" "--host" "0.0.0.0" "--port" "53800" "/aim" ];
    };

    systemd.services."podman-aim-server" = {
      after    = lib.mkAfter [ "mnt-data.mount" "podman-homey-network.service" ];
      requires = lib.mkAfter [ "mnt-data.mount" "podman-homey-network.service" ];
    };

    # -------------------------------------------------------------------------
    # Web UI — Caddy proxy, Authelia-protected
    # -------------------------------------------------------------------------
    virtualisation.oci-containers.containers.aim-ui = {
      image = cfg.image;
      ports = [ "127.0.0.1:${toString cfg.uiPort}:43800" ];
      volumes  = [ "${dataDir}/aim:/aim" ];
      extraOptions = [ "--network=homey" ];
      cmd = [ "aim" "ui" "--host" "0.0.0.0" "--port" "43800" "/aim" ];
    };

    systemd.services."podman-aim-ui" = {
      after    = lib.mkAfter [ "mnt-data.mount" "podman-homey-network.service" "podman-aim-server.service" ];
      requires = lib.mkAfter [ "mnt-data.mount" "podman-homey-network.service" ];
    };

    # -------------------------------------------------------------------------
    # Firewall — restrict tracking port to LAN subnet only
    # -------------------------------------------------------------------------
    networking.firewall.extraCommands = lib.mkAfter ''
      iptables -A nixos-fw -p tcp --dport ${toString cfg.serverPort} -s 192.168.1.0/24 -j nixos-fw-accept
    '';

    # -------------------------------------------------------------------------
    # Caddy virtual host — UI behind Authelia forward_auth
    # -------------------------------------------------------------------------
    homey.caddy.virtualHosts = [{
      subdomain = "aim";
      port      = cfg.uiPort;
      auth      = true;
    }];

    # -------------------------------------------------------------------------
    # Authelia — two_factor + deny, admins only (priority 27–28)
    # -------------------------------------------------------------------------
    homey.authelia.accessControlRules = [
      { priority = 27; domain = [ "aim.${domain}" ]; policy = "two_factor"; subject = [ "group:admins" ]; }
      { priority = 28; domain = [ "aim.${domain}" ]; policy = "deny"; }
    ];

    # -------------------------------------------------------------------------
    # Storage directory
    # -------------------------------------------------------------------------
    homey.storage.extraDirs = [
      { path = "aim"; mode = "0755"; }
    ];

    # -------------------------------------------------------------------------
    # Backup — the .aim binary store holds all run data
    # -------------------------------------------------------------------------
    homey.backup.extraPaths = [ "${dataDir}/aim" ];

    # -------------------------------------------------------------------------
    # Uptime Kuma monitor
    # -------------------------------------------------------------------------
    homey.monitoring.monitors = [{
      name     = "Aim";
      url      = "https://aim.${domain}";
      interval = 60;
    }];
  };
}