# UGREEN DXP4800 LED control
# Based on https://github.com/miskcoo/ugreen_leds_controller
{ config, lib, pkgs, ... }:

let
  # Disk mapping: ata port -> LED name
  # DXP4800 has bays 1-4, currently bays 2 and 4 are populated
  diskMapping = {
    # ata-2 (sdb) -> disk2
    "2" = "disk2";
    # ata-4 (sdc) -> disk4
    "4" = "disk4";
  };

  # LED colors (R G B)
  colors = {
    healthy = "0 255 0";      # Green
    activity = "0 255 0";     # Green blink
    standby = "0 0 255";      # Dim blue when sleeping
    fail = "255 0 0";         # Red
    network = "0 255 0";      # Green
    power = "255 255 255";    # White
  };

  brightness = 255;
  standbyBrightness = 30;     # Dim when in standby
  refreshInterval = "0.1";    # Seconds between activity checks
  powerCheckInterval = 30;    # Seconds between power state checks

  # Script to initialize LEDs on boot
  initLedsScript = pkgs.writeShellScript "ugreen-leds-init" ''
    set -euo pipefail
    PATH="${lib.makeBinPath [ pkgs.ugreen-leds-cli ]}:$PATH"

    # Wait for i2c device to be available
    for i in $(seq 1 30); do
      if [ -e /dev/i2c-0 ]; then
        break
      fi
      sleep 1
    done

    # Initialize power LED - solid white
    ugreen_leds_cli power -on -color ${colors.power} -brightness ${toString brightness} || true

    # Initialize network LED - will be controlled by netdevmon
    ugreen_leds_cli netdev -off || true

    # Initialize disk LEDs based on mapping
    ${lib.concatStringsSep "\n" (lib.mapAttrsToList (ata: led: ''
      ugreen_leds_cli ${led} -off || true
    '') diskMapping)}

    echo "UGREEN LEDs initialized"
  '';

  # Disk activity monitoring script
  diskMonitorScript = pkgs.writeShellScript "ugreen-diskiomon" ''
    set -euo pipefail
    PATH="${lib.makeBinPath [ pkgs.ugreen-leds-cli pkgs.coreutils pkgs.gnugrep pkgs.gawk pkgs.smartmontools pkgs.hdparm ]}:$PATH"

    # Build device -> LED mapping by checking ata ports
    declare -A devices
    declare -A diskio_data
    declare -A disk_healthy
    declare -A disk_standby

    # Discover disks based on ata port mapping
    for path in /dev/disk/by-path/pci-*-ata-*; do
      [ -e "$path" ] || continue
      # Skip partitions
      [[ "$path" == *-part* ]] && continue

      # Extract ata port number (e.g., ata-2 -> 2)
      ata_port=$(echo "$path" | grep -oP 'ata-\K[0-9]+' | head -1)

      case "$ata_port" in
        ${lib.concatStringsSep "\n        " (lib.mapAttrsToList (ata: led: ''
          ${ata})
            device=$(readlink -f "$path")
            short_name=$(basename "$device")
            devices["${led}"]="$short_name"
            echo "Mapped $short_name (ata-${ata}) -> ${led}"
            ;;'') diskMapping)}
      esac
    done

    if [ ''${#devices[@]} -eq 0 ]; then
      echo "No disks found matching ATA ports, exiting"
      exit 1
    fi

    # Set initial LED state for discovered disks
    for led in "''${!devices[@]}"; do
      device="''${devices[$led]}"

      # Check SMART health (this will wake the disk at boot, which is acceptable)
      if smartctl -H "/dev/$device" 2>/dev/null | grep -q "PASSED"; then
        disk_healthy["$led"]=1
        ugreen_leds_cli "$led" -on -color ${colors.healthy} -brightness ${toString brightness} || true
      else
        disk_healthy["$led"]=0
        ugreen_leds_cli "$led" -on -color ${colors.fail} -brightness ${toString brightness} || true
      fi

      # Initialize tracking
      diskio_data["$led"]=""
      disk_standby["$led"]=0
    done

    echo "Starting disk activity monitoring for ''${#devices[@]} disk(s)"

    # Function to update LED based on current state
    update_led() {
      local led="$1"
      local device="''${devices[$led]}"

      # Check power state without waking disk
      local power_state
      power_state=$(hdparm -C "/dev/$device" 2>/dev/null | grep -oP '(standby|active/idle|active|idle)' | head -1 || echo "unknown")

      if [[ "$power_state" == "standby" ]]; then
        if [[ "''${disk_standby[$led]}" != "1" ]]; then
          # Disk just went to standby - dim the LED
          disk_standby["$led"]=1
          ugreen_leds_cli "$led" -on -color ${colors.standby} -brightness ${toString standbyBrightness} || true
          echo "Disk $device entered standby, dimming LED"
        fi
      else
        if [[ "''${disk_standby[$led]}" == "1" ]]; then
          # Disk woke up - restore health-based color
          disk_standby["$led"]=0
          if [[ "''${disk_healthy[$led]}" == "1" ]]; then
            ugreen_leds_cli "$led" -on -color ${colors.healthy} -brightness ${toString brightness} || true
          else
            ugreen_leds_cli "$led" -on -color ${colors.fail} -brightness ${toString brightness} || true
          fi
          echo "Disk $device woke up, restoring LED"
        fi
      fi
    }

    # Background power state checker
    check_power_states() {
      while true; do
        sleep ${toString powerCheckInterval}
        for led in "''${!devices[@]}"; do
          update_led "$led"
        done
      done
    }

    # Start power state checker in background
    check_power_states &
    POWER_CHECK_PID=$!
    trap "kill $POWER_CHECK_PID 2>/dev/null || true" EXIT

    # Main activity monitoring loop
    while true; do
      for led in "''${!devices[@]}"; do
        device="''${devices[$led]}"
        stat_file="/sys/block/$device/stat"

        if [ -f "$stat_file" ]; then
          new_stat=$(cat "$stat_file" 2>/dev/null || echo "")

          if [ -n "$new_stat" ] && [ "''${diskio_data[$led]}" != "$new_stat" ]; then
            # Activity detected - disk must be awake now
            if [[ "''${disk_standby[$led]}" == "1" ]]; then
              disk_standby["$led"]=0
              if [[ "''${disk_healthy[$led]}" == "1" ]]; then
                ugreen_leds_cli "$led" -on -color ${colors.healthy} -brightness ${toString brightness} || true
              else
                ugreen_leds_cli "$led" -on -color ${colors.fail} -brightness ${toString brightness} || true
              fi
            fi

            # Trigger LED blink for activity
            if [ -e "/sys/class/leds/$led/shot" ]; then
              echo 1 > "/sys/class/leds/$led/shot" 2>/dev/null || true
            else
              ugreen_leds_cli "$led" -blink 100 100 2>/dev/null || true
              sleep 0.05
              ugreen_leds_cli "$led" -on 2>/dev/null || true
            fi
          fi

          diskio_data["$led"]="$new_stat"
        fi
      done

      sleep ${refreshInterval}
    done
  '';

  # Network activity monitoring script
  netMonitorScript = pkgs.writeShellScript "ugreen-netdevmon" ''
    set -euo pipefail
    PATH="${lib.makeBinPath [ pkgs.ugreen-leds-cli pkgs.coreutils pkgs.iproute2 ]}:$PATH"

    INTERFACE="$1"
    CHECK_INTERVAL=60

    echo "Starting network monitoring on $INTERFACE"

    # Configure LED to trigger on network activity
    led_path="/sys/class/leds/netdev"

    while true; do
      # Check if interface is up
      if ip link show "$INTERFACE" 2>/dev/null | grep -q "state UP"; then
        # Link is up - set green
        ugreen_leds_cli netdev -on -color ${colors.network} -brightness ${toString brightness} || true

        # Try to enable hardware trigger for activity indication
        if [ -e "$led_path/device_name" ]; then
          echo "$INTERFACE" > "$led_path/device_name" 2>/dev/null || true
        fi
        if [ -e "$led_path/rx" ]; then
          echo 1 > "$led_path/rx" 2>/dev/null || true
        fi
        if [ -e "$led_path/tx" ]; then
          echo 1 > "$led_path/tx" 2>/dev/null || true
        fi
      else
        # Link is down - turn off
        ugreen_leds_cli netdev -off || true
      fi

      sleep $CHECK_INTERVAL
    done
  '';

in
{
  # Load i2c-dev kernel module for LED controller communication
  boot.kernelModules = [ "i2c-dev" ];

  # Install CLI tool
  environment.systemPackages = [ pkgs.ugreen-leds-cli ];

  # LED initialization service - runs once at boot
  systemd.services.ugreen-leds-init = {
    description = "Initialize UGREEN NAS LEDs";
    wantedBy = [ "multi-user.target" ];
    after = [ "local-fs.target" "systemd-modules-load.service" ];
    serviceConfig = {
      Type = "oneshot";
      RemainAfterExit = true;
      ExecStart = "${initLedsScript}";
    };
  };

  # Disk activity monitoring service
  systemd.services.ugreen-diskiomon = {
    description = "UGREEN disk activity LED monitor";
    wantedBy = [ "multi-user.target" ];
    after = [ "ugreen-leds-init.service" "local-fs.target" ];
    requires = [ "ugreen-leds-init.service" ];
    serviceConfig = {
      Type = "simple";
      ExecStart = "${diskMonitorScript}";
      Restart = "always";
      RestartSec = "5s";
    };
  };

  # Network activity monitoring service (template for interface)
  systemd.services."ugreen-netdevmon@" = {
    description = "UGREEN network LED monitor for %i";
    after = [ "ugreen-leds-init.service" "network-online.target" ];
    requires = [ "ugreen-leds-init.service" ];
    serviceConfig = {
      Type = "simple";
      ExecStart = "${netMonitorScript} %i";
      Restart = "always";
      RestartSec = "10s";
    };
  };

  # Enable network monitoring for primary interface
  systemd.services."ugreen-netdevmon@enp2s0" = {
    wantedBy = [ "multi-user.target" ];
  };
}