{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE Arrows #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}

module Sensor
  ( cpu,
    io,
    net,
    mem,
    bat,
    Sensor.snd,
    disk,
    date,
    time,
    wmName,
    wmWorkspaces,
  )
where

import Control.Arrow
import Control.DeepSeq
import Control.Monad
import Control.Monad.Reader
import Data.Char
import Data.Dynamic
import Data.Functor.WithIndex
import Data.List
import Data.Maybe
import Data.Sensor qualified as S
import Data.Set qualified as S
import Data.Time.Clock
import Data.Time.Format
import Data.Time.LocalTime
import GHC.Generics (Generic)
import Graphics.X11 qualified as X
import Graphics.X11.Xlib.Extras qualified as X
import Pretty qualified as P
import Process.Shell (sh)
import Safe
import System.Directory
import System.FilePath
import System.IO (hGetContents)
import System.Posix.StatVFS
import System.Posix.Types (Fd (..))
import Text.Printf
import UnliftIO
import UnliftIO.Concurrent
import Witherable (ifilter)
import Prelude hiding (readFile)

data CpuStatData = CpuStatData
  { used :: Int,
    total :: Int
  }
  deriving (Eq, Show, Generic, NFData)

data CpuStat = CpuStat deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m CpuStat CpuStatData where
  aggregate _ = forever do S.yield =<< parse <* sleep
    where
      parse = readFile "/proc/stat" $ \stat -> do
        case filter (["cpu"] `isPrefixOf`) (words <$> lines stat) of
          [] -> error "/proc/stat: no cpu line"
          ((drop 1 . fmap read -> (xs@(_ : _ : _ : idle : _))) : _) -> do
            let total = sum xs
                used = total - idle
            pure (CpuStatData used total)
          _ -> error "/proc/stat: unexpected cpu line"

readFile :: (MonadUnliftIO m, NFData a) => FilePath -> (String -> m a) -> m a
readFile fp f =
  withFile fp ReadMode $ \h ->
    evaluate =<< f =<< (liftIO (hGetContents h))

sleep :: (MonadIO m) => S.AggregateT s m ()
sleep = threadDelay (5 * 10 ^ 5)

cpuStat :: (S.MonadSensor m) => S.Sensor m () CpuStatData
cpuStat = S.sensor CpuStat

cpu :: (S.MonadSensor m) => S.Sensor m () (P.Diagram [Float])
cpu = step <$$> diagram 3 cpuStat
  where
    step xs =
      zipWith
        ( \next prev ->
            let CpuStatData {used, total} =
                  CpuStatData (next.used - prev.used) (next.total - prev.total)
             in if total == 0 then 0 else fromIntegral used / fromIntegral total
        )
        (fromMaybe [] (tailMay xs))
        (fromMaybe [] (initMay xs))

diagram :: (Monad m) => Int -> S.Sensor m () a -> S.Sensor m () (P.Diagram [a])
diagram n sf = S.feedbackS [] $ proc ((), xs) -> do
  x <- sf -< ()
  returnA -< (P.diagram n (reverse (x : xs)), take (2 * (n - 1) + 1) (x : xs))

data IoStat = IoStat deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m IoStat Int where
  aggregate _ = forever do S.yield =<< parse <* sleep
    where
      parse = liftIO do
        fmap sum . mapM (parse1 . ("/sys/block" </>))
          =<< listDirectory "/sys/block"
      parse1 :: FilePath -> IO Int
      parse1 fp = do
        readFile (fp </> "stat") $ \stat -> do
          case words stat of
            [_, _, read', _, _, _, write, _, _, _, _, _, _, _, _, _, _] ->
              pure (read read' + read write)
            _ -> error (printf "%s/stat: malformed" fp)

ioStat :: (S.MonadSensor m) => S.Sensor m () Int
ioStat = S.sensor IoStat

io :: (S.MonadSensor m) => S.Sensor m () (P.Diagram [Float])
io = diagram 3 (maxS (rate ioStat))

maxS :: (Monad m, Num a, Ord a, Integral a) => S.Sensor m () a -> S.Sensor m () Float
maxS sf = S.feedbackS 0 $ proc ((), xMax) -> do
  x <- fi <$> sf -< ()
  let xMax' = max xMax x
  returnA -< (if xMax <= 0 then 0 else x / xMax', xMax')

rate :: (Monad m, Num a) => S.Sensor m () a -> S.Sensor m () a
rate sf = S.feedbackS Nothing $ proc ((), x') -> do
  x <- sf -< ()
  returnA -< (maybe 0 ((-) x) x', Just x)

data NetStat = NetStat deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m NetStat Int where
  aggregate _ = forever do S.yield =<< parse <* sleep
    where
      parse = liftIO do
        is <- listDirectory "/sys/class/net"
        sum <$> mapM (parse1 . ("/sys/class/net" </>)) is
      parse1 :: FilePath -> IO Int
      parse1 fp = do
        (+)
          <$> (readFile (fp </> "statistics/rx_bytes") (pure . read))
          <*> (readFile (fp </> "statistics/tx_bytes") (pure . read))

netStat :: (S.MonadSensor m) => S.Sensor m () Int
netStat = S.sensor NetStat

net :: (S.MonadSensor m) => S.Sensor m () (P.Diagram [Float])
net = diagram 3 (maxS (rate netStat))

data MemStat = MemStat deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m MemStat Float where
  aggregate _ = forever do S.yield =<< parse <* sleep
    where
      parse = readFile "/proc/meminfo" $ \meminfo -> do
        case foldl
          ( \(total, avail) xs ->
              case xs of
                ["MemTotal:", v, "kB"] -> (Just (read v), avail)
                ["MemAvailable:", v, "kB"] -> (total, Just (read v))
                _ -> (total, avail)
          )
          (Nothing, Nothing)
          (map words (lines meminfo)) of
          (Just total, Just avail) -> pure (1 - avail / total)
          (Nothing, _) -> error "/proc/meminfo: MemTotal missing"
          (_, Nothing) -> error "/proc/meminfo: MemAvail missing"

memStat :: (S.MonadSensor m) => S.Sensor m () Float
memStat = S.sensor MemStat

mem :: (S.MonadSensor m) => S.Sensor m () (P.Diagram [Float])
mem = diagram 3 memStat

data BatStat = BatStat deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m BatStat (Float, Bool) where
  aggregate _ = forever do S.yield =<< parse <* sleep
    where
      parse = liftIO do
        fmap combine
          . mapM (parse1 . ("/sys/class/power_supply" </>))
          =<< listDirectory "/sys/class/power_supply"
      combine abs =
        ( product (map fst abs),
          or (map Prelude.snd abs)
        )
      parse1 :: FilePath -> IO (Float, Bool)
      parse1 fp =
        liftM2
          (,)
          ( choice 1 $
              [ (/)
                  <$> (readFile (fp </> "charge_now") (pure . read))
                  <*> (readFile (fp </> "charge_full") (pure . read)),
                (/)
                  <$> (readFile (fp </> "energy_now") (pure . read))
                  <*> (readFile (fp </> "energy_full") (pure . read))
              ]
          )
          (choice False [readFile (fp </> "status") (pure . (==) "charging" . map toLower)])

bat :: (S.MonadSensor m) => S.Sensor m () P.Doc
bat =
  S.sensor BatStat >>= \(value, charging) ->
    if charging
      then
        if
          | value >= 0.9 -> pure (P.pretty "\xf0085")
          | value >= 0.8 -> pure (P.pretty "\xf008b")
          | value >= 0.7 -> pure (P.pretty "\xf008a")
          | value >= 0.6 -> pure (P.pretty "\xf089e")
          | value >= 0.5 -> pure (P.pretty "\xf0089")
          | value >= 0.4 -> pure (P.pretty "\xf089d")
          | value >= 0.3 -> pure (P.pretty "\xf0088")
          | value >= 0.2 -> pure (P.pretty "\xf0087")
          | value >= 0.1 -> pure (P.pretty "\xf0086")
          | otherwise -> pure (P.pretty "\xf089c")
      else
        if
          | value >= 0.9 -> pure (P.pretty "\xf0082")
          | value >= 0.8 -> pure (P.pretty "\xf0081")
          | value >= 0.7 -> pure (P.pretty "\xf0080")
          | value >= 0.6 -> pure (P.pretty "\xf008f")
          | value >= 0.5 -> pure (P.pretty "\xf008e")
          | value >= 0.4 -> pure (P.pretty "\xf007d")
          | value >= 0.3 -> pure (P.pretty "\xf007c")
          | value >= 0.2 -> pure (P.pretty "\xf007b")
          | value >= 0.1 -> pure (P.pretty "\xf007a")
          | otherwise ->
              (\b -> if b then P.pretty "\xf008e" else (P.color P.Red) (P.pretty "\xf0083"))
                <$> blink

data Blink = Blink deriving (Show)

blink :: (S.MonadSensor m) => S.Sensor m () Bool
blink = S.sensor Blink

instance (S.MonadSensor m) => S.Aggregate m Blink Bool where
  aggregate _ = loop True
    where
      loop b = S.yield b >> sleep >> loop (not b)

choice :: a -> [IO a] -> IO a
choice def [] = pure def
choice def (x : xs) = x `catch` (\(_ :: SomeException) -> choice def xs)

data Snd = Snd deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m Snd Float where
  aggregate _ = forever do S.yield =<< parse <* sleep
    where
      parse = liftIO do (/ 153) . read <$> [sh|pamixer --get-volume|]

snd :: (S.MonadSensor m) => S.Sensor m () String
snd = do
  value <- S.sensor Snd
  return
    ( if
        | 0.4 < value -> "\xf028"
        | 0.0 < value -> "\xf027"
        | otherwise -> "\xf026"
    )

data DiskStat = DiskStat deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m DiskStat Float where
  aggregate _ = forever do S.yield =<< parse <* sleep
    where
      parse = liftIO do
        stat <- statVFS "/"
        pure $
          (fi stat.statVFS_bfree)
            / fi (stat.statVFS_bfree + stat.statVFS_bavail)

diskStat :: (S.MonadSensor m) => S.Sensor m () Float
diskStat = S.sensor DiskStat

disk :: (S.MonadSensor m) => S.Sensor m () (P.Diagram [Float])
disk = diagram 1 diskStat

data CurrentTime = CurrentTime deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m CurrentTime UTCTime where
  aggregate _ = forever do
    x <- liftIO getCurrentTime
    S.yield =<< liftIO getCurrentTime
    threadDelay
      (((10 ^ 12) - (fromEnum (utctDayTime x) `mod` (10 ^ 12))) `div` (10 ^ 6))

data CurrentTimeZone = CurrentTimeZone deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m CurrentTimeZone TimeZone where
  aggregate _ = forever do
    S.yield =<< liftIO getCurrentTimeZone
    sleep

currentTimeZone :: (S.MonadSensor m) => S.Sensor m () TimeZone
currentTimeZone = S.sensor CurrentTimeZone

currentTime :: (S.MonadSensor m) => S.Sensor m () UTCTime
currentTime = S.sensor CurrentTime

date :: (S.MonadSensor m) => S.Sensor m () String
date =
  ((formatTime defaultTimeLocale "%b %e" .) . utcToLocalTime)
    <$> currentTimeZone
    <*> currentTime

time :: (S.MonadSensor m) => S.Sensor m () String
time =
  ((formatTime defaultTimeLocale "%R" .) . utcToLocalTime)
    <$> currentTimeZone
    <*> currentTime

wmName :: (S.MonadSensor m) => S.Sensor m () String
wmName = activeWindow >>= maybe (pure "") wmNameOf

data ActiveWindow = ActiveWindow deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m ActiveWindow (Maybe X.Window) where
  aggregate _ = do
    bracket
      ( liftIO do
          dpy <- X.openDisplay ""
          let root = X.rootWindowOfScreen (X.defaultScreenOfDisplay dpy)
          X.selectInput dpy root X.propertyChangeMask
          X.sync dpy False
          pure (dpy, root)
      )
      (\(dpy, _) -> liftIO (X.closeDisplay dpy))
      ( \(dpy, root) -> do
          netActiveWindow <- liftIO (X.internAtom dpy "_NET_ACTIVE_WINDOW" False)
          let waitForEvent = do
                liftIO (X.pending dpy) >>= \case
                  0 -> do
                    liftIO do threadWaitRead (Fd (X.connectionNumber dpy))
                    waitForEvent
                  _ -> do
                    e <-
                      liftIO
                        ( X.allocaXEvent $ \ev -> do
                            X.nextEvent dpy ev
                            X.getEvent ev
                        )
                    when (X.ev_atom e /= netActiveWindow) waitForEvent

              readActiveWindow = do
                liftIO (X.getWindowProperty32 dpy netActiveWindow root) >>= \case
                  Just (0 : _) -> S.yield Nothing
                  Just (win : _) -> S.yield (Just (fi win))
                  _ -> S.yield Nothing
          forever do
            readActiveWindow
            waitForEvent
      )

activeWindow :: (S.MonadSensor m) => S.Sensor m () (Maybe X.Window)
activeWindow = S.sensor ActiveWindow

data WmNameOf = WmNameOf X.Window deriving (Show)

instance (S.MonadSensor m) => S.Aggregate m WmNameOf String where
  aggregate (WmNameOf win) = do
    bracket
      ( liftIO do
          dpy <- X.openDisplay ""
          X.selectInput dpy win X.propertyChangeMask
          X.sync dpy False
          pure dpy
      )
      (liftIO . X.closeDisplay)
      ( \dpy -> do
          netWmName <- liftIO (X.internAtom dpy "_NET_WM_NAME" False)
          let waitForEvent = do
                liftIO (X.pending dpy) >>= \case
                  0 -> do
                    liftIO do threadWaitRead (Fd (X.connectionNumber dpy))
                    waitForEvent
                  _ -> do
                    e <-
                      liftIO
                        ( X.allocaXEvent $ \ev -> do
                            X.nextEvent dpy ev
                            X.getEvent ev
                        )
                    when (X.ev_atom e /= netWmName) waitForEvent

              readWmName = do
                S.yield
                  =<< fmap (strip . head) . liftIO . X.wcTextPropertyToTextList dpy
                  =<< liftIO (X.getTextProperty dpy win netWmName)
          forever do
            readWmName
            waitForEvent
      )

wmNameOf :: (S.MonadSensor m) => X.Window -> S.Sensor m () String
wmNameOf = S.sensor . WmNameOf

data WmWorkspaces = WmWorkspaces deriving (Show)

data Workspace
  = Active String
  | Inactive String
  deriving (Eq, Typeable, Show)

instance (S.MonadSensor m) => S.Aggregate m WmWorkspaces [Workspace] where
  aggregate WmWorkspaces = do
    bracket
      ( liftIO do
          dpy <- X.openDisplay ""
          let root = X.rootWindowOfScreen (X.defaultScreenOfDisplay dpy)
          X.selectInput dpy root X.propertyChangeMask
          X.sync dpy False
          pure (dpy, root)
      )
      (\(dpy, _) -> liftIO (X.closeDisplay dpy))
      ( \(dpy, root) -> do
          netClientList <- liftIO $ X.internAtom dpy "_NET_CLIENT_LIST" False
          netCurrentDesktop <- liftIO $ X.internAtom dpy "_NET_CURRENT_DESKTOP" False
          netDesktopNames <- liftIO $ X.internAtom dpy "_NET_DESKTOP_NAMES" False
          netWmDesktop <- liftIO $ X.internAtom dpy "_NET_WM_DESKTOP" False
          let waitForEvent = do
                liftIO (X.pending dpy) >>= \case
                  0 -> do
                    liftIO (threadWaitRead (Fd (X.connectionNumber dpy)))
                    waitForEvent
                  _ -> do
                    e <-
                      liftIO
                        ( X.allocaXEvent $ \ev -> do
                            X.nextEvent dpy ev
                            X.getEvent ev
                        )
                    when
                      ( ((not .) . elem)
                          (X.ev_atom e)
                          [ netClientList,
                            netCurrentDesktop,
                            netDesktopNames,
                            netWmDesktop
                          ]
                      )
                      waitForEvent

              readWmWorkspaces = do
                S.yield
                  =<< liftIO do
                    currentDesktop <-
                      fmap (fi . head)
                        <$> X.getWindowProperty32 dpy netCurrentDesktop root
                    occupiedDesktops <-
                      fmap (S.unions . catMaybes)
                        . mapM
                          ( \win -> do
                              fmap (S.singleton . fi . head)
                                <$> X.getWindowProperty32 dpy netWmDesktop win
                          )
                        . map fi
                        . fromMaybe []
                        =<< X.getWindowProperty32 dpy netClientList root
                    fmap
                      ( ifilter
                          ( \i _ ->
                              Just i == currentDesktop
                                || i `S.member` occupiedDesktops
                          )
                          . imap
                            ( \i ->
                                if Just i == currentDesktop
                                  then Active
                                  else Inactive
                            )
                      )
                      . X.wcTextPropertyToTextList dpy
                      =<< X.getTextProperty dpy root netDesktopNames
          forever do
            readWmWorkspaces
            waitForEvent
      )

wmWorkspaces :: (S.MonadSensor m) => S.Sensor m () [Workspace]
wmWorkspaces = S.sensor WmWorkspaces

instance P.Pretty [Workspace] where
  pretty =
    P.Col . intersperse (P.pretty " ") . map P.pretty

instance P.Pretty Workspace where
  pretty (Active s) = P.color P.White (P.pretty s)
  pretty (Inactive s) = P.color P.Cyan (P.pretty s)

(<$$>) :: (Functor f, Functor g) => (a -> b) -> f (g a) -> f (g b)
f <$$> x = fmap f <$> x

infixl 4 <$$>

fi :: (Integral a, Num b) => a -> b
fi = fromIntegral

strip :: String -> String
strip =
  reverse . dropWhile isSpace . reverse . dropWhile isSpace