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

{
  options.vpnProxy = {
    enable = lib.mkEnableOption "VPN SOCKS proxy script";
    
    vmUser = lib.mkOption {
      type = lib.types.str;
      default = "nate";
      description = "Username for VM SSH connection";
    };
    
    vmIp = lib.mkOption {
      type = lib.types.str;
      default = "192.168.122.241";
      description = "IP address of the VM";
    };
    
    socksPort = lib.mkOption {
      type = lib.types.int;
      default = 1080;
      description = "Local SOCKS proxy port";
    };
    
    vpnTestUrl = lib.mkOption {
      type = lib.types.str;
      default = "https://doc-aut.app.vasionnow.com/";
      description = "URL to test VPN connectivity";
    };
  };

  config = lib.mkIf config.vpnProxy.enable {
    home.packages = [
      (pkgs.writeShellApplication {
        name = "vpn-proxy";
        runtimeInputs = with pkgs; [ 
          openssh 
          procps 
          lsof 
          gnused 
          coreutils 
          curl 
          glib  # for gsettings
          nettools  # for ss command
        ];
        text = ''
          VM_USER="${config.vpnProxy.vmUser}"
          VM_IP="${config.vpnProxy.vmIp}"
          SOCKS_PORT=${toString config.vpnProxy.socksPort}
          PID_FILE="/tmp/vpn-socks.pid"
          VPN_TEST_URL="${config.vpnProxy.vpnTestUrl}"

          check_ssh_agent() {
              # Check if SSH agent is available
              if [ -z "$SSH_AUTH_SOCK" ]; then
                  echo "Error: SSH agent not running (SSH_AUTH_SOCK not set)"
                  return 1
              fi

              # Check if any SSH keys are loaded
              if ! ssh-add -l >/dev/null 2>&1; then
                  echo "Error: No SSH keys are loaded in the agent"
                  echo "Please add your SSH key first with: ssh-add ~/.ssh/id_ed25519"
                  echo ""
                  echo "Agent status: SSH_AUTH_SOCK=$SSH_AUTH_SOCK"
                  return 1
              fi

              # Show loaded keys for confirmation
              echo "✓ SSH agent has loaded keys:"
              ssh-add -l | sed 's/^/  /'
              return 0
          }

          start_proxy() {
              # Check if proxy is already running
              if [ -f "$PID_FILE" ]; then
                  local pid
                  pid=$(cat "$PID_FILE" 2>/dev/null)
                  if [ -n "$pid" ] && ps -p "$pid" > /dev/null 2>&1; then
                      echo "Proxy already running (PID: $pid)"
                      return 1
                  else
                      echo "Stale PID file found, cleaning up..."
                      rm -f "$PID_FILE"
                  fi
              fi
              
              # Check SSH agent and keys
              echo "Checking SSH agent..."
              if ! check_ssh_agent; then
                  return 1
              fi
              echo ""
              
              # Test SSH connectivity first
              echo "Testing SSH connection to VM..."
              if ! ssh -o ConnectTimeout=5 -o BatchMode=yes -o StrictHostKeyChecking=no "$VM_USER@$VM_IP" exit 2>/dev/null; then
                  echo "Error: Cannot connect to VM: $VM_USER @ $VM_IP"
                  echo "Make sure:"
                  echo "  1. VM is running"
                  echo "  2. SSH service is running in VM"
                  echo "  3. SSH key authentication is set up"
                  return 1
              fi
              
              # Start SSH tunnel in background
              echo "Starting SSH SOCKS proxy..."
              ssh -D $SOCKS_PORT -f -N -o ServerAliveInterval=60 -o ServerAliveCountMax=3 -o ExitOnForwardFailure=yes "$VM_USER@$VM_IP"
              local ssh_exit=$?
              
              if [ $ssh_exit -ne 0 ]; then
                  echo "Error: SSH tunnel failed to start (exit code: $ssh_exit)"
                  return 1
              fi
              
              # Wait a moment for SSH to fully establish
              sleep 1
              
              # Find and save the SSH process PID
              local pid
              pid=$(pgrep -f "ssh -D $SOCKS_PORT.*$VM_USER@$VM_IP" | head -1)
              if [ -z "$pid" ]; then
                  echo "Error: SSH process not found after startup"
                  return 1
              fi
              echo "$pid" > "$PID_FILE"
              
              # Verify port is listening
              echo "Verifying SOCKS proxy is listening on port $SOCKS_PORT..."
              local retry=0
              while [ $retry -lt 5 ]; do
                  if lsof -i :"$SOCKS_PORT" > /dev/null 2>&1 || ss -tln | grep -q ":$SOCKS_PORT "; then
                      break
                  fi
                  sleep 1
                  retry=$((retry + 1))
              done
              
              if [ $retry -eq 5 ]; then
                  echo "Error: SOCKS proxy port $SOCKS_PORT is not listening"
                  pkill -P "$pid" 2>/dev/null
                  kill "$pid" 2>/dev/null
                  rm -f "$PID_FILE"
                  return 1
              fi
              
              # Configure system proxy
              gsettings set org.gnome.system.proxy mode 'manual'
              gsettings set org.gnome.system.proxy.socks host 'localhost'
              gsettings set org.gnome.system.proxy.socks port $SOCKS_PORT
              
              echo "✓ VPN proxy started successfully on localhost:$SOCKS_PORT (PID: $pid)"
              
              # Test VPN connectivity
              echo "Testing VPN connectivity..."
              if curl -s -m 10 --socks5-hostname localhost:"$SOCKS_PORT" "$VPN_TEST_URL" > /dev/null 2>&1; then
                  echo "✓ VPN connection verified - can reach $VPN_TEST_URL"
              else
                  echo "⚠ Warning: Could not reach $VPN_TEST_URL through proxy"
                  echo "  The SOCKS proxy is running but VPN connection may not be active in the VM"
              fi
          }

          stop_proxy() {
              if [ -f "$PID_FILE" ]; then
                  local pid
                  pid=$(cat "$PID_FILE" 2>/dev/null)
                  
                  # Kill the specific SSH process
                  if [ -n "$pid" ] && ps -p "$pid" > /dev/null 2>&1; then
                      echo "Stopping VPN proxy (PID: $pid)..."
                      kill "$pid" 2>/dev/null
                      
                      # Wait for process to die
                      local retry=0
                      while [ $retry -lt 5 ] && ps -p "$pid" > /dev/null 2>&1; do
                          sleep 1
                          retry=$((retry + 1))
                      done
                      
                      # Force kill if still alive
                      if ps -p "$pid" > /dev/null 2>&1; then
                          echo "Process didn't stop gracefully, force killing..."
                          kill -9 "$pid" 2>/dev/null
                      fi
                  else
                      echo "PID $pid not found in process list"
                  fi
                  
                  rm -f "$PID_FILE"
              else
                  echo "Proxy not running (no PID file found)"
              fi
              
              # Always clean up orphaned SSH processes
              if pgrep -f "ssh -D $SOCKS_PORT" > /dev/null 2>&1; then
                  echo "Cleaning up orphaned SSH processes..."
                  pkill -f "ssh -D $SOCKS_PORT"
              fi
              
              # Always disable system proxy and clear SOCKS settings
              echo "Clearing proxy settings..."
              gsettings set org.gnome.system.proxy mode 'none'
              gsettings set org.gnome.system.proxy.socks host ""
              gsettings set org.gnome.system.proxy.socks port 0
              
              # Verify settings are cleared
              local proxy_mode
              proxy_mode=$(gsettings get org.gnome.system.proxy mode 2>/dev/null)
              if [ "$proxy_mode" = "'none'" ]; then
                  echo "✓ System proxy disabled"
              else
                  echo "⚠ Warning: Could not verify proxy was disabled (current mode: $proxy_mode)"
              fi
              
              echo "✓ VPN proxy stopped and proxy settings cleared"
              
              # Return success even if there was no PID file
              return 0
          }

          status_proxy() {
              if [ -f "$PID_FILE" ]; then
                  local pid
                  pid=$(cat "$PID_FILE" 2>/dev/null)
                  if [ -n "$pid" ] && ps -p "$pid" > /dev/null 2>&1; then
                      echo "✓ Proxy running on localhost:$SOCKS_PORT (PID: $pid)"
                      
                      # Check if port is actually listening
                      if lsof -i :"$SOCKS_PORT" > /dev/null 2>&1 || ss -tln | grep -q ":$SOCKS_PORT "; then
                          echo "✓ Port $SOCKS_PORT is listening"
                      else
                          echo "⚠ Warning: Process running but port not listening"
                      fi
                      
                      # Check system proxy settings
                      local proxy_mode
                      proxy_mode=$(gsettings get org.gnome.system.proxy mode 2>/dev/null)
                      if [ "$proxy_mode" = "'manual'" ]; then
                          echo "✓ System proxy configured"
                      else
                          echo "⚠ System proxy not configured (mode: $proxy_mode)"
                      fi
                      
                      # Quick connectivity test
                      if curl -s -m 5 --socks5-hostname localhost:"$SOCKS_PORT" "$VPN_TEST_URL" > /dev/null 2>&1; then
                          echo "✓ VPN connectivity verified"
                      else
                          echo "⚠ Cannot reach VPN resources"
                      fi
                      
                      return 0
                  else
                      echo "✗ Proxy not running (stale PID file)"
                      return 1
                  fi
              else
                  echo "✗ Proxy not running"
                  
                  # Check for orphaned processes
                  if pgrep -f "ssh -D $SOCKS_PORT" > /dev/null 2>&1; then
                      echo "⚠ Warning: Found SSH process without PID file"
                      echo "  Run 'vpn-proxy stop' to clean up"
                  fi
                  
                  return 1
              fi
          }

          case "''${1:-}" in
              start)
                  start_proxy
                  ;;
              stop)
                  stop_proxy
                  ;;
              restart)
                  stop_proxy
                  sleep 1
                  start_proxy
                  ;;
              status)
                  status_proxy
                  ;;
              test)
                  # Quick VPN test without starting/stopping
                  if curl -s -m 10 --socks5-hostname localhost:"$SOCKS_PORT" "$VPN_TEST_URL" > /dev/null 2>&1; then
                      echo "✓ VPN connection working - can reach $VPN_TEST_URL"
                      exit 0
                  else
                      echo "✗ Cannot reach $VPN_TEST_URL through proxy"
                      exit 1
                  fi
                  ;;
              *)
                  echo "Usage: vpn-proxy {start|stop|restart|status|test}"
                  exit 1
                  ;;
          esac
        '';
      })
    ];
  };
}