{ config, lib, pkgs, aic8800-src ? null, ... }:

with lib;

let
  cfg = config.hardware.aic8800;

  # Get source from function argument, specialArgs, or throw error
  aic8800Src = if aic8800-src != null then
    "${aic8800-src}/usr/src/AIC8800"
  else if builtins.hasAttr "_module" config && 
          builtins.hasAttr "args" config._module &&
          builtins.hasAttr "aic8800-src" config._module.args then
    "${config._module.args.aic8800-src}/usr/src/AIC8800"
  else
    throw ''
      aic8800-src must be provided either:
      1. As a function argument when importing the module, or
      2. Via specialArgs in your flake.nix:
         specialArgs = { aic8800-src = inputs.aic8800-src; };
    '';

  aic8800Firmware = pkgs.stdenvNoCC.mkDerivation {
    pname = "aic8800-firmware";
    version = "2022-12-19";
    src = aic8800Src;

    installPhase = ''
      runHook preInstall
      mkdir -p $out/lib/firmware
      cp -r fw/aic8800D80 $out/lib/firmware/
      runHook postInstall
    '';

    meta = with lib; {
      description = "Firmware blobs for AIC8800 USB adapters";
      license = licenses.unfreeRedistributableFirmware;
      platforms = platforms.linux;
    };
  };

  aic8800Driver = config.boot.kernelPackages.callPackage
    ({ stdenv,
       lib,
       kernel,
       python3
     }: stdenv.mkDerivation {
      pname = "aic8800-driver";
      version = "2022-12-19";

      src = aic8800Src;

      nativeBuildInputs = (kernel.moduleBuildDependencies or [ ]) ++ [ python3 ];

      postPatch = ''
        python3 <<'EOF'
      from pathlib import Path

      header = Path('drivers/aic8800/aic8800_fdrv/aicwf_usb.h')
      lines = header.read_text().splitlines()

      define_line = '#define USB_PRODUCT_ID_AIC8800D80       0x8d80'
      if all('USB_PRODUCT_ID_AIC8800D80' not in line for line in lines):
        out_lines = []
        inserted = 0
        for line in lines:
          out_lines.append(line)
          if line.strip().startswith('#define USB_PRODUCT_ID_AIC8800D81'):
            out_lines.append(define_line)
            inserted += 1
        if inserted == 0:
          raise SystemExit('marker for USB product IDs missing')
        header.write_text('\n'.join(out_lines) + '\n')

      table = Path('drivers/aic8800/aic8800_fdrv/aicwf_usb.c')
      table_text = table.read_text()
      table_anchor = '{USB_DEVICE_AND_INTERFACE_INFO(USB_VENDOR_ID_AIC, USB_PRODUCT_ID_AIC8800D81, 0xff, 0xff, 0xff)},\n'
      table_entry = '{USB_DEVICE_AND_INTERFACE_INFO(USB_VENDOR_ID_AIC, USB_PRODUCT_ID_AIC8800D80, 0xff, 0xff, 0xff)},\n'

      if table_entry not in table_text:
        if table_anchor not in table_text:
          raise SystemExit('USB table anchor missing')
        table.write_text(table_text.replace(table_anchor, table_anchor + table_entry, 1))
      EOF
      '';

      hardeningDisable = [ "pic" "format" ];

      buildPhase = ''
        runHook preBuild
        make -C drivers/aic8800 \
          KDIR=${kernel.dev}/lib/modules/${kernel.modDirVersion}/build \
          ARCH=${stdenv.hostPlatform.linuxArch}
        runHook postBuild
      '';

      installPhase = ''
        runHook preInstall
        modDir=$out/lib/modules/${kernel.modDirVersion}/kernel/drivers/net/wireless
        install -Dm644 drivers/aic8800/aic_load_fw/aic_load_fw.ko \
          "$modDir/aic_load_fw.ko"
        install -Dm644 drivers/aic8800/aic8800_fdrv/aic8800_fdrv.ko \
          "$modDir/aic8800_fdrv.ko"
        mkdir -p $out/lib/firmware
        cp -r --no-preserve=mode,ownership fw/aic8800D80 $out/lib/firmware/
        runHook postInstall
      '';

      meta = with lib; {
        description = "UGREEN / AIC8800 USB Wi-Fi driver";
        homepage = "https://www.ugreen.com";
        license = licenses.unfreeRedistributableFirmware;
        platforms = platforms.linux;
      };
    }) { inherit (pkgs) python3; };

in
{
  options.hardware.aic8800 = {
    enable = mkEnableOption "AIC8800 USB Wi-Fi driver support";

    enableUdevRules = mkOption {
      type = types.bool;
      default = true;
      description = "Enable udev rules for automatic USB disk ejection and power management";
    };
  };

  config = mkIf cfg.enable {
    boot.extraModulePackages = [ aic8800Driver ];
    boot.kernelModules = lib.mkAfter [ "aic_load_fw" "aic8800_fdrv" ];

    hardware.firmware = [ aic8800Firmware ];

    services.udev.extraRules = mkIf cfg.enableUdevRules (lib.mkAfter ''
      # UGREEN AX900 shows up as a tiny storage device; ejecting it flips to Wi-Fi mode
      KERNEL=="sd*", ACTION=="add", SUBSYSTEM=="block", \
        ATTRS{idVendor}=="a69c", ATTRS{idProduct}=="5721", \
        SYMLINK+="aicudisk", RUN+="${pkgs.util-linux}/bin/eject /dev/%k"
      KERNEL=="sd*", ACTION=="add", SUBSYSTEM=="block", \
        ATTRS{idVendor}=="a69c", ATTRS{idProduct}=="5723", \
        SYMLINK+="tendaudisk", RUN+="${pkgs.util-linux}/bin/eject /dev/%k"
      KERNEL=="sd*", ACTION=="add", SUBSYSTEM=="block", \
        ATTRS{idVendor}=="a69c", ATTRS{idProduct}=="5724", \
        SYMLINK+="ugreenudisk", RUN+="${pkgs.util-linux}/bin/eject /dev/%k"

      # Keep the Wi-Fi interface awake to avoid round-trip spikes after autosuspend
      ACTION=="add", SUBSYSTEM=="usb", ATTR{idVendor}=="a69c", ATTR{idProduct}=="8d80", \
        TEST=="power/control", ATTR{power/control}="on"
    '');
  };
}