# Colemak-DH EC Keyboard Remap for Framework Laptop
#
# TROUBLESHOOTING NOTES (for future agents):
# ==========================================
# This module applies keyboard remaps via the Framework EC (Embedded Controller).
# EC remaps are VOLATILE - lost after ~10 seconds unplugged (cold boot).
#
# If remap doesn't work at LUKS prompt after cold boot:
#
# 1. Check if service ran:
#    journalctl -b -u colemak-ec-remap
#
# 2. Enable verbose boot for debugging:
#    boot.initrd.verbose = lib.mkForce true;
#    boot.kernelParams = lib.mkForce [ "boot.shell_on_fail" ];
#
# 3. Verify EC modules loaded in initrd:
#    lsinitrd /boot/initrd | grep cros_ec
#
# 4. Check service ordering - must run BEFORE cryptsetup.target
#
# 5. If systemd initrd approach fails entirely, fallback options:
#    - FrameworkHacksPkg EFI driver: https://github.com/DHowett/FrameworkHacksPkg
#    - This runs before OS starts, survives cold boots
#
# References:
# - https://www.howett.net/posts/2021-12-framework-ec/
# - https://deck.sh/using-the-ec-to-change-the-keyboard-layout-on-the-framework-laptop/
# - https://wiki.archlinux.org/title/Framework_Laptop_13

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

let
  # All ectool remap commands
  remapCommands = ectoolPath: ''
    # Colemak-DH layout via EC for Framework 12
    # Matrix positions verified via FRAME12_KEY_MATRIX.md scan

    # Top row (QWERTY → Colemak-DH)
    ${ectoolPath} raw 0x3E0C d1,d1,b7,b8,w2b   # e(7,8) → f
    ${ectoolPath} raw 0x3E0C d1,d1,b7,b9,w4d   # r(7,9) → p
    ${ectoolPath} raw 0x3E0C d1,d1,b2,b3,w32   # t(2,3) → b
    ${ectoolPath} raw 0x3E0C d1,d1,b2,b6,w3b   # y(2,6) → j
    ${ectoolPath} raw 0x3E0C d1,d1,b7,b1,w4b   # u(7,1) → l
    ${ectoolPath} raw 0x3E0C d1,d1,b7,b2,w3c   # i(7,2) → u
    ${ectoolPath} raw 0x3E0C d1,d1,b7,b3,w35   # o(7,3) → y
    ${ectoolPath} raw 0x3E0C d1,d1,b7,b4,w4c   # p(7,4) → ;

    # Home row (QWERTY → Colemak-DH)
    ${ectoolPath} raw 0x3E0C d1,d1,b3,b4,w2d   # s(3,4) → r
    ${ectoolPath} raw 0x3E0C d1,d1,b4,b2,w1b   # d(4,2) → s
    ${ectoolPath} raw 0x3E0C d1,d1,b4,b3,w2c   # f(4,3) → t
    ${ectoolPath} raw 0x3E0C d1,d1,b1,b6,w3a   # h(1,6) → m
    ${ectoolPath} raw 0x3E0C d1,d1,b4,b6,w31   # j(4,6) → n
    ${ectoolPath} raw 0x3E0C d1,d1,b4,b5,w24   # k(4,5) → e
    ${ectoolPath} raw 0x3E0C d1,d1,b4,b9,w43   # l(4,9) → i
    ${ectoolPath} raw 0x3E0C d1,d1,b4,b8,w44   # ;(4,8) → o

    # Bottom row with angle mod (QWERTY → Colemak-DH)
    ${ectoolPath} raw 0x3E0C d1,d1,b6,b1,w22   # z(6,1) → x
    ${ectoolPath} raw 0x3E0C d1,d1,b5,b8,w21   # x(5,8) → c
    ${ectoolPath} raw 0x3E0C d1,d1,b5,b5,w23   # c(5,5) → d
    ${ectoolPath} raw 0x3E0C d1,d1,b0,b3,w1a   # b(0,3) → z
    ${ectoolPath} raw 0x3E0C d1,d1,b0,b5,w42   # n(0,5) → k
    ${ectoolPath} raw 0x3E0C d1,d1,b5,bb,w33   # m(5,b) → h
  '';

  # Smart remap: check first, only apply if needed (for systemd service)
  # Reads position (7,8) - if scancode is 0x2b (f), remap is already applied
  smartRemapScript = ectoolPath: ''
    # Read current scancode at e-key position (7,8)
    # Response has scancode at bytes 10-11; we check byte 10 for 0x2b (f)
    CURRENT=$(${ectoolPath} raw 0x3E0C d1,d0,b7,b8,w0 2>/dev/null | grep '|' | head -1 | awk '{print $11}')
    if [ "$CURRENT" = "2b" ]; then
      echo "Colemak-DH remap already active, skipping."
      exit 0
    fi

    echo "Applying Colemak-DH EC remap..."
    ${remapCommands ectoolPath}
    echo "Colemak-DH remap applied."
  '';

  # Script with retry logic for initrd
  initrdRemapScript = pkgs.writeShellScript "colemak-ec-initrd-remap" ''
    set -e
    ECTOOL="${pkgs.fw-ectool}/bin/ectool"

    # Wait for EC driver to be ready (up to 3 seconds)
    echo "Waiting for EC driver..."
    for i in $(seq 1 30); do
      if $ECTOOL version >/dev/null 2>&1; then
        echo "EC driver ready after $((i * 100))ms"
        break
      fi
      sleep 0.1
    done

    # Verify EC communication works
    if ! $ECTOOL version >/dev/null 2>&1; then
      echo "ERROR: Cannot communicate with EC after 3 seconds"
      exit 1
    fi

    echo "Applying Colemak-DH EC remap..."
    ${remapCommands "$ECTOOL"}
    echo "Colemak-DH remap applied successfully"
  '';
in
{
  options.colemakEc.enable = lib.mkEnableOption "Colemak-DH EC keyboard remap";

  config = lib.mkIf config.colemakEc.enable {
    # ===========================================
    # INITRD: Apply remap BEFORE LUKS decryption
    # ===========================================

    # Enable systemd in initrd (required for this approach)
    boot.initrd.systemd.enable = true;

    # Add kernel modules for EC communication
    boot.initrd.availableKernelModules = [ "cros_ec" "cros_ec_lpcs" ];

    # Add fw-ectool and the remap script to initrd
    boot.initrd.systemd.storePaths = [ pkgs.fw-ectool initrdRemapScript ];

    # Systemd service that runs before cryptsetup
    # Key: wantedBy sysinit.target + requiredBy on cryptsetup units ensures
    # this service runs before LUKS password prompt
    boot.initrd.systemd.services.colemak-ec-remap = {
      description = "Apply Colemak-DH keyboard remap to EC";
      wantedBy = [ "sysinit.target" ];
      before = [ "cryptsetup.target" "systemd-ask-password-console.service" ];
      after = [ "systemd-modules-load.service" ];
      wants = [ "systemd-modules-load.service" ];

      unitConfig = {
        DefaultDependencies = false;
      };

      serviceConfig = {
        Type = "oneshot";
        RemainAfterExit = true;
        ExecStart = initrdRemapScript;
        StandardOutput = "journal+console";
        StandardError = "journal+console";
      };
    };

    # Make cryptsetup units wait for our remap service
    # This adds After= and Wants= dependencies to each LUKS device's cryptsetup unit
    boot.initrd.systemd.services."systemd-cryptsetup@" = {
      after = [ "colemak-ec-remap.service" ];
      wants = [ "colemak-ec-remap.service" ];
    };

    # ===========================================
    # SYSTEMD: Apply remap after boot (backup)
    # ===========================================
    # This ensures the remap is applied even if the EC state was reset,
    # and provides a service that can be restarted manually if needed.

    systemd.services.colemak-ec-remap = {
      description = "Apply Colemak-DH remap to EC keyboard";
      after = [ "multi-user.target" ];
      wantedBy = [ "multi-user.target" ];

      serviceConfig = {
        Type = "oneshot";
        RemainAfterExit = true;
        ExecStart = "${pkgs.writeShellScript "colemak-dh-remap" ''
          ${smartRemapScript "${pkgs.fw-ectool}/bin/ectool"}
        ''}";
      };
    };

    # Also run after resume from suspend/hibernate
    powerManagement.resumeCommands = ''
      systemctl start colemak-ec-remap
    '';
  };
}