module Rahm.Desktop.Swap
  ( SwapBuddy (..),
    getSwapBuddy,
    recordSwap,
    swapWindows,
    removeSwapBuddy,
  )
where

import Control.Monad (forM_, void)
import Data.Default (Default (..))
import Data.Map (Map)
import qualified Data.Map as Map
import Rahm.Desktop.Common
import Rahm.Desktop.Hooks.WindowChange
import Rahm.Desktop.Logger
import qualified Rahm.Desktop.StackSet as SS
import XMonad (ExtensionClass (..), Window, X, withWindowSet, StateExtension (PersistentExtension), windows)
import qualified XMonad.Util.ExtensibleState as XS
import Data.Maybe (maybeToList)

newtype SwapBuddy = SwapBuddy
  { swapBuddies :: Map Window Window
  }
  deriving (Read, Show)

instance Default SwapBuddy where
  def = SwapBuddy Map.empty

instance ExtensionClass SwapBuddy where
  initialValue = def
  extensionType = PersistentExtension

-- | Record a swap pairing between two windows.
-- The buddy relationship is bidirectional: w1's buddy becomes w2 and vice versa.
recordSwap :: Window -> Window -> X ()
recordSwap w1 w2 = do
  XS.modify $ \sw@(SwapBuddy m) ->
    SwapBuddy $ Map.insert w1 w2 $ Map.insert w2 w1 m
  logs Debug "Recorded swap pairing: %s <-> %s" (show w1) (show w2)

-- | Get the current (most recent) swap buddy for a window.
getSwapBuddy :: Window -> X (Maybe Window)
getSwapBuddy w = do
  SwapBuddy m <- XS.get
  return (Map.lookup w m)

-- | Swap windows between two lists.
-- Each window in w1s is swapped with the corresponding window in w2s.
-- Both windows in each pair become each other's swap buddy.
swapWindows :: [Window] -> [Window] -> X ()
swapWindows ws1 ws2 = do
  windows $ SS.swapWindows (zip ws1 ws2)
  mapM_ (uncurry recordSwap) (zip ws1 ws2)

-- | Clean up all swap buddy entries for a given window.
-- Should be called when a window is killed or moves to hidden workspace.
removeSwapBuddy :: Window -> X ()
removeSwapBuddy w = do
  SwapBuddy m <- XS.get
  let buddy = Map.lookup w m
      cleaned = foldr Map.delete (Map.delete w m) (maybeToList buddy)
  XS.put $ SwapBuddy cleaned
  logs Debug "Removed swap buddy for window: %s" (show w)