{-# LANGUAGE OverloadedStrings #-}

--------------------------------------------------------------------------------
-- Commands for controlling power

module XMonad.Internal.Command.Power
-- commands
  ( runHibernate
  , runOptimusPrompt
  , runPowerOff
  , runPowerPrompt
  , runReboot
  , runScreenLock
  , runSuspend
  , runSuspendPrompt
  , runQuitPrompt
  -- daemons
  , runAutolock
  -- functions
  , hasBattery
  , suspendPrompt
  , quitPrompt
  , powerPrompt
  , defFontPkgs
  , promptFontDep
  )
where

import Data.Internal.Dependency
import Graphics.X11.Types
import RIO
import RIO.FilePath
import qualified RIO.Map as M
import qualified RIO.Process as P
import qualified RIO.Text as T
import UnliftIO.Directory
import XMonad.Core hiding (spawn)
import XMonad.Internal.Shell
import qualified XMonad.Internal.Theme as XT
import XMonad.Prompt
import XMonad.Prompt.ConfirmPrompt

--------------------------------------------------------------------------------
-- Executables
myScreenlock :: FilePath
myScreenlock = "screenlock"

myOptimusManager :: FilePath
myOptimusManager = "optimus-manager"

myPrimeOffload :: FilePath
myPrimeOffload = "prime-offload"

--------------------------------------------------------------------------------
-- Packages

optimusPackages :: [Fulfillment]
optimusPackages = [Package AUR "optimus-manager"]

--------------------------------------------------------------------------------
-- Core commands

runScreenLock :: SometimesX
runScreenLock =
  sometimesExe
    "screen locker"
    "i3lock script"
    [Package AUR "i3lock-color"]
    False
    myScreenlock

runPowerOff :: X ()
runPowerOff = spawn "systemctl poweroff"

runSuspend :: X ()
runSuspend = spawn "systemctl suspend"

runHibernate :: X ()
runHibernate = spawn "systemctl hibernate"

runReboot :: X ()
runReboot = spawn "systemctl reboot"

--------------------------------------------------------------------------------
-- Autolock

runAutolock :: Sometimes (FIO (P.Process () () ()))
runAutolock = sometimesIO_ "automatic screen lock" "xss-lock" tree cmd
  where
    tree =
      And_ (Only_ $ sysExe [Package Official "xss-lock"] "xss-lock") $
        Only_ $
          IOSometimes_ runScreenLock
    cmd = P.proc "xss-lock" ["--ignore-sleep", "screenlock"] (P.startProcess . P.setCreateGroup True)

--------------------------------------------------------------------------------
-- Confirmation prompts

promptFontDep :: IOTree XT.FontBuilder
promptFontDep = fontTreeAlt XT.defFontFamily defFontPkgs

defFontPkgs :: [Fulfillment]
defFontPkgs = [Package Official "ttf-dejavu"]

confirmPrompt' :: T.Text -> X () -> XT.FontBuilder -> X ()
confirmPrompt' s x fb = confirmPrompt (XT.promptTheme fb) (T.unpack s) x

suspendPrompt :: XT.FontBuilder -> X ()
suspendPrompt = confirmPrompt' "suspend?" runSuspend

quitPrompt :: XT.FontBuilder -> X ()
quitPrompt = confirmPrompt' "quit?" $ io exitSuccess

sometimesPrompt :: T.Text -> (XT.FontBuilder -> X ()) -> SometimesX
sometimesPrompt n = sometimesIO n (T.append n " command") promptFontDep

-- TODO doesn't this need to also lock the screen?
runSuspendPrompt :: SometimesX
runSuspendPrompt = sometimesPrompt "suspend prompt" suspendPrompt

runQuitPrompt :: SometimesX
runQuitPrompt = sometimesPrompt "quit prompt" quitPrompt

--------------------------------------------------------------------------------
-- Nvidia Optimus

-- TODO for some reason the screen never wakes up after suspend when
-- the nvidia card is up, so block suspend if nvidia card is running
-- and warn user
isUsingNvidia :: MonadUnliftIO m => m Bool
isUsingNvidia = doesDirectoryExist "/sys/module/nvidia"

hasBattery :: MonadUnliftIO m => m (Maybe T.Text)
hasBattery = do
  ps <- fromRight [] <$> tryIO (listDirectory syspath)
  ts <- catMaybes <$> mapM readType ps
  return $
    if any (T.isPrefixOf "Battery") ts
      then Nothing
      else Just "battery not found"
  where
    readType p = either (const Nothing) Just <$> tryIO (readFileUtf8 $ syspath </> p </> "type")
    syspath = "/sys/class/power_supply"

runOptimusPrompt' :: XT.FontBuilder -> X ()
runOptimusPrompt' fb = do
  nvidiaOn <- io isUsingNvidia
  switch $ if nvidiaOn then "integrated" else "nvidia"
  where
    switch mode = confirmPrompt' (prompt mode) (cmd mode) fb
    prompt mode = T.concat ["gpu switch to ", mode, "?"]
    cmd mode =
      spawn $
        T.pack myPrimeOffload
          #!&& T.unwords [T.pack myOptimusManager, "--switch", mode, "--no-confirm"]
          #!&& "killall xmonad"

runOptimusPrompt :: SometimesX
runOptimusPrompt =
  Sometimes
    "graphics switcher"
    (\x -> xpfOptimus x && xpfBattery x)
    [s]
  where
    s = Subfeature {sfData = r, sfName = "optimus manager"}
    r = IORoot runOptimusPrompt' t
    t =
      And1 promptFontDep $
        listToAnds (socketExists "optimus-manager" [] socketName) $
          sysExe optimusPackages <$> [myOptimusManager, myPrimeOffload]
    socketName = (</> "optimus-manager") <$> getTemporaryDirectory

--------------------------------------------------------------------------------
-- Universal power prompt

data PowerMaybeAction
  = Poweroff
  | Shutdown
  | Hibernate
  | Reboot
  deriving (Eq)

instance Enum PowerMaybeAction where
  toEnum 0 = Poweroff
  toEnum 1 = Shutdown
  toEnum 2 = Hibernate
  toEnum 3 = Reboot
  toEnum _ = errorWithoutStackTrace "Main.Enum.PowerMaybeAction.toEnum: bad argument"

  fromEnum Poweroff = 0
  fromEnum Shutdown = 1
  fromEnum Hibernate = 2
  fromEnum Reboot = 3

data PowerPrompt = PowerPrompt

instance XPrompt PowerPrompt where
  showXPrompt PowerPrompt = "(P)oweroff (S)uspend (H)ibernate (R)eboot:"

runPowerPrompt :: SometimesX
runPowerPrompt = Sometimes "power prompt" (const True) [sf]
  where
    sf = Subfeature withLock "prompt with lock"
    withLock = IORoot (uncurry powerPrompt) tree
    tree = And12 (,) lockTree promptFontDep
    lockTree = Or (Only $ IOSometimes runScreenLock id) (Only $ IOConst skip)

powerPrompt :: X () -> XT.FontBuilder -> X ()
powerPrompt lock fb = mkXPrompt PowerPrompt theme comp executeMaybeAction
  where
    comp = mkComplFunFromList theme []
    theme = (XT.promptTheme fb) {promptKeymap = keymap}
    keymap =
      M.fromList $
        ((controlMask, xK_g), quit)
          : map
            (first $ (,) 0)
            [ (xK_p, sendMaybeAction Poweroff)
            , (xK_s, sendMaybeAction Shutdown)
            , (xK_h, sendMaybeAction Hibernate)
            , (xK_r, sendMaybeAction Reboot)
            , (xK_Return, quit)
            , (xK_Escape, quit)
            ]
    sendMaybeAction a = setInput (show $ fromEnum a) >> setSuccess True >> setDone True
    executeMaybeAction a = case toEnum $ read a of
      Poweroff -> runPowerOff
      Shutdown -> lock >> runSuspend
      Hibernate -> lock >> runHibernate
      Reboot -> runReboot