{-# LANGUAGE FlexibleInstances #-}

-- | Tablet mode support for pen displays
--
-- This module provides tablet mode functionality including:
--   - Screen rotation (0°, 90°, 180°, 270°)
--   - Coordinate transformation for pen/finger devices
--   - Keyboard disable/enable in tablet mode
--   - State management via XMonad ExtensibleState
module Rahm.Desktop.TabletMode
  ( -- * Configuration
    TabletModeConfig (..),
    defaultConfig,
    setTabletConfig,

    -- * State
    TabletModeState (..),

    -- * Operations
    setTabletMode,
    toggleTabletMode,
    getTabletMode,
    setRotation,

    -- * Convenience
    enableTabletMode,
    disableTabletMode,

    -- * Types
    RotationAngle (..),
    tabletModeXMobarString,
    tabletModeCommandHandler,
    syntheticKeyStartDrag,
    syntheticKeyStartResize,
  )
where

import Control.Concurrent (forkIO)
import Control.Monad (forM_, unless, when)
import Data.Char (toLower)
import Data.List (intercalate, isInfixOf)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, isJust)
import Data.Typeable (Typeable)
import Graphics.X11.Xlib hiding (hWaitReady)
import qualified Graphics.X11.Xlib.Extras
import Rahm.Desktop.Logger (LogLevel (Debug, Info, Warn), logs)
import Rahm.Desktop.Server
import System.Directory (createDirectoryIfMissing)
import System.Environment (lookupEnv)
import System.FilePath ((</>))
import qualified System.IO as IO
import System.Process (callProcess, readProcessWithExitCode)
import XMonad
  ( Atom,
    Display,
    ExtensionClass (..),
    StateExtension (PersistentExtension),
    Window,
    X,
    ask,
    io,
    spawn,
  )
import qualified XMonad.Util.ExtensibleState as XS
import XMonad.Util.Run (runProcessWithInput)

syntheticKeyStartDrag :: KeySym
syntheticKeyStartDrag = 0x1100000000000011

syntheticKeyStartResize :: KeySym
syntheticKeyStartResize = 0x1100000000000111

-- | Rotation angle for the display
data RotationAngle
  = -- | Normal orientation (0°)
    Rot0
  | -- | Rotated 90° clockwise
    Rot90
  | -- | Rotated 180° (inverted)
    Rot180
  | -- | Rotated 270° clockwise (90° counter-clockwise)
    Rot270
  deriving (Eq, Show, Typeable, Enum, Bounded, Read)

circularSucc :: (Bounded a, Enum a) => a -> a
circularSucc x = toEnum $ (fromEnum x + 1) `mod` (fromEnum (maxBound `asTypeOf` x) + 1)

-- | Configuration for tablet mode
data TabletModeConfig = TabletModeConfig
  { -- | Display name (e.g., "eDP-1")
    configDisplay :: String,
    -- | Keyboard device name (Nothing to disable keyboard control)
    configKeyboard :: Maybe String
  }
  deriving (Eq, Show, Typeable, Read)

-- | Default configuration matching the user's hardware
defaultConfig :: TabletModeConfig
defaultConfig =
  TabletModeConfig
    { configDisplay = "eDP-1",
      configKeyboard = Just "AT Translated Set 2 keyboard"
    }

-- | Tablet mode state
data TabletModeState = TabletModeState
  { stateConfig :: TabletModeConfig,
    tabletModeEnabled :: Bool,
    tabletRotation :: RotationAngle,
    lazyHasTabletMode :: Maybe Bool
  }
  deriving (Eq, Show, Read)

instance ExtensionClass TabletModeState where
  initialValue =
    TabletModeState
      { stateConfig = defaultConfig,
        tabletModeEnabled = False,
        tabletRotation = Rot0,
        lazyHasTabletMode = Nothing
      }
  extensionType = PersistentExtension

hasTabletMode :: X Bool
hasTabletMode = do
  cached <- lazyHasTabletMode <$> XS.get
  case cached of
    Just b -> return b
    Nothing -> do
      hasTouch <- io $ do
        output <- runProcessWithInput "xinput" ["list"] ""
        let hasTouchDevice = any (isInfixOf "touch" . map toLower) (lines output)
        return hasTouchDevice
      XS.modify (\st -> st {lazyHasTabletMode = Just hasTouch})
      return hasTouch

-- | Update the tablet mode configuration
setTabletConfig :: TabletModeConfig -> X ()
setTabletConfig config = XS.modify (\s -> s {stateConfig = config})

-- | Check if tablet mode is currently enabled
getTabletMode :: X Bool
getTabletMode = tabletModeEnabled <$> XS.get

-- | Enable or disable tablet mode
setTabletMode :: Bool -> X ()
setTabletMode enable = do
  state <- XS.get
  let config = stateConfig state

  when enable $
    spawn "on-tablet-mode.sh 1"

  if enable
    then do
      logs Info "Enabling tablet mode"
      setRotation Rot180
      maybeDisableKeyboard config
      XS.put state {tabletModeEnabled = True, tabletRotation = Rot180}
    else do
      logs Info "Disabling tablet mode"
      setRotation Rot0
      maybeEnableKeyboard config
      XS.put state {tabletModeEnabled = False, tabletRotation = Rot0}

  unless enable $
    spawn "on-tablet-mode.sh 0"

-- | Toggle tablet mode on/off
toggleTabletMode :: X ()
toggleTabletMode = do
  enabled <- getTabletMode
  setTabletMode (not enabled)

-- | Rotation flags for xrandr
rotationFlags :: RotationAngle -> String
rotationFlags Rot0 = "normal"
rotationFlags Rot90 = "right"
rotationFlags Rot180 = "inverted"
rotationFlags Rot270 = "left"

-- | Coordinate transformation matrices for each rotation
-- Format: [a, b, tx, c, d, ty, p, q, r] for matrix:
-- | a  b  tx |
-- | c  d  ty |
-- | p  q  r  |
rotationMatrix :: RotationAngle -> String
rotationMatrix Rot0 = "1 0 0 0 1 0 0 0 1"
rotationMatrix Rot90 = "0 1 0 -1 0 1 0 0 1"
rotationMatrix Rot180 = "-1 0 1 0 -1 1 0 0 1"
rotationMatrix Rot270 = "0 -1 1 1 0 0 0 0 1"

rotationNext :: X ()
rotationNext = do
  curRotation <- tabletRotation <$> XS.get
  setRotation (circularSucc curRotation)

-- | Set the display rotation and device transformations
setRotation :: RotationAngle -> X ()
setRotation angle = do
  state <- XS.get
  let config = stateConfig state
      display = configDisplay config
      matrix = rotationMatrix angle
      xrandrFlag = rotationFlags angle

  logs Info "Setting rotation to %s" (show angle)

  -- Rotate display
  io $ spawn $ unwords ["xrandr", "--output", display, "--rotate", xrandrFlag]

  -- Apply coordinate transformations to devices
  let devices =
        [ "Wacom HID 527E Pen stylus",
          "Wacom HID 527E Pen eraser",
          "Wacom HID 527E Finger touch"
        ]

  forM_ devices $ \dev -> do
    let matrixStr = unwords (map show matrix)
    logs Debug "Setting %s matrix: %s" dev matrix
    io $
      spawn $
        unwords
          [ "xinput",
            "set-prop",
            "\"" ++ dev ++ "\"",
            "\"Coordinate Transformation Matrix\"",
            matrix
          ]

  -- Update state
  XS.put state {tabletRotation = angle}

-- | Disable keyboard if configured
maybeDisableKeyboard :: TabletModeConfig -> X ()
maybeDisableKeyboard config = do
  when (isJust (configKeyboard config)) $ do
    let keyboard = fromMaybe "" (configKeyboard config)
    logs Info "Disabling keyboard: %s" keyboard
    io $ spawn $ unwords ["xinput", "disable", "\"" ++ keyboard ++ "\""]

-- | Enable keyboard if configured
maybeEnableKeyboard :: TabletModeConfig -> X ()
maybeEnableKeyboard config = do
  when (isJust (configKeyboard config)) $ do
    let keyboard = fromMaybe "" (configKeyboard config)
    logs Info "Enabling keyboard: %s" keyboard
    io $ spawn $ unwords ["xinput", "enable", "\"" ++ keyboard ++ "\""]

tabletModeXMobarString :: X String
tabletModeXMobarString = do
  doesHaveTabletMode <- hasTabletMode
  if doesHaveTabletMode
    then do
      tabletMode <- getTabletMode
      return $
        ( if tabletMode
            then
              concat
                [ "<action=`rdectl tablet-mode off`><fc=#ffcc66> \61706 </fc></action>",
                  "<fc=#a0a0a0><action=onboard> \61724 </action></fc>",
                  "<fc=#f0f0a0><action=`sleep 0.2 && rdectl wheel`> \11096 </action></fc>",
                  "<fc=#f08080><action=`rdectl shift . _`> ✗ </action></fc>",
                  "<fc=#808080><action=`rofi -monitor -4 -display-run Execute -show run -normal-window`>\61728 </action></fc>",
                  "<fc=#80f080><action=`rdectl next-rotation`> \58183 </action></fc>",
                  "<fc=#80f080><action=`sleep 0.2 && rdectl keysym " ++ show syntheticKeyStartDrag ++ "`> \983486 </action></fc>",
                  "<fc=#80f080><action=`sleep 0.2 && rdectl keysym " ++ show syntheticKeyStartResize ++ "`> \985704 </action></fc>"
                ]
            else "<action=`rdectl tablet-mode on`><fc=#66cc66> \61705 </fc></action>"
        )
          ++ "  "
    else return ""

data TabletModeCommand where
  SetTabletMode :: Bool -> TabletModeCommand
  NextRotation :: TabletModeCommand

tabletModeCommandHandler :: CommandHandler TabletModeCommand
tabletModeCommandHandler =
  CommandHandler
    ( \s ->
        case words s of
          ["tablet-mode", "on"] -> Just $ SetTabletMode True
          ["tablet-mode", "off"] -> Just $ SetTabletMode False
          ["next-rotation"] -> Just NextRotation
          _ -> Nothing
    )
    ( \case
        SetTabletMode mode -> setTabletMode mode
        NextRotation -> rotationNext
    )

-- | Enable tablet mode (rotate 180°, disable keyboard)
enableTabletMode :: X ()
enableTabletMode = setTabletMode True

-- | Disable tablet mode (rotate 0°, enable keyboard)
disableTabletMode :: X ()
disableTabletMode = setTabletMode False