{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}

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

import Control.Monad
import Control.Monad.Reader
import Data.Char
import Data.Dynamic
import Data.Functor.WithIndex
import Data.List
import Data.Map qualified as M
import Data.Maybe
import Data.Set qualified as S
import Data.Time.Clock
import Data.Time.Format
import Data.Time.LocalTime
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.Error
import System.Linux.Inotify qualified as I
import System.Posix.StatVFS
import Text.Printf
import UnliftIO.Concurrent
import UnliftIO.Exception
import UnliftIO.Memoize
import UnliftIO.STM
import Witherable (ifilter)

type SensorM a =
  ReaderT (MVar (M.Map String (Memoized Dynamic))) IO a

runSensorM :: SensorM a -> IO a
runSensorM m = do
  g <- newMVar M.empty
  runReaderT m g

poll :: (Typeable a) => Sensor a -> SensorM (TVar a)
poll (Sensor id vM') = do
  gM <- ask
  g <- takeMVar gM
  vM <- case M.lookup id g of
    Just vM -> do
      putMVar gM g
      pure (fromJust . fromDynamic <$> vM)
    Nothing -> do
      -- liftIO (printf "[%s] start\n" id)
      vM <- vM'
      putMVar gM (M.insert id (toDyn <$> vM) g)
      pure vM
  runMemoized vM

data Sensor a = Sensor String (SensorM (Memoized (TVar a)))

data CpuStat = CpuStat
  { total :: Int,
    idle :: Int
  }
  deriving (Show, Eq)

cpuStat :: Sensor CpuStat
cpuStat = pollFile "cpuStat" aggregate "/proc/stat" (5 * 10 ^ 5)
  where
    aggregate fp = liftIO do
      stat <- readFile fp
      case filter ("cpu " `isPrefixOf`) (lines stat) of
        [] -> error "/proc/stat: no cpu line"
        (cpu : _) ->
          case map read (drop 1 (words cpu)) of
            xs@(_ : _ : _ : idle : _) ->
              pure $ CpuStat (sum xs) idle
            _ -> error "/proc/stat: unexpected cpu line"

cpu' :: Sensor [CpuStat]
cpu' = histogram "cpu'" [] (\b a -> take 7 (a : b)) cpuStat

data Cpu = Cpu {unCpu :: [Float]}

instance P.Pretty Cpu where
  pretty = P.diagram 3 . unCpu

cpu :: Sensor Cpu
cpu = transform "cpu" f cpu'
  where
    f xs = do
      let xs' =
            zipWith
              ( \next prev ->
                  let total = next.total - prev.total
                      idle = next.idle - prev.idle
                   in fromIntegral (total - idle) / fromIntegral total
              )
              (fromMaybe [] (initMay xs))
              (fromMaybe [] (tailMay xs))
      pure (Cpu xs')

data IoStat = IoStat {unIoStat :: Int} deriving (Show, Eq)

ioStat :: Sensor IoStat
ioStat = poll' "ioStat" aggregate (5 * 10 ^ 5)
  where
    aggregate = liftIO do
      fmap (IoStat . sum) . mapM (aggregate' . ("/sys/block" </>))
        =<< listDirectory "/sys/block"
    aggregate' fp = do
      stat <- readFile (fp </> "stat")
      case words stat of
        [_, _, read', _, _, _, write', _, _, _, _, _, _, _, _, _, _] ->
          pure (read read' + read write')
        _ -> error (printf "%s/stat: malformed" fp)

io' :: Sensor [IoStat]
io' = histogram "io'" [] (\b a -> take 7 (a : b)) ioStat

data Io = Io {unIo :: [Float]} deriving (Show, Eq)

instance P.Pretty Io where
  pretty = P.diagram 3 . unIo

io :: Sensor Io
io = transform "io" f io'
  where
    f (map unIoStat -> xs) = do
      let xs' =
            zipWith
              (-)
              (fromMaybe [] (initMay xs))
              (fromMaybe [] (tailMay xs))
          x' = max 1 (fromMaybe 0 (maximumMay xs'))
      pure . Io $
        map (\x -> fromIntegral x / fromIntegral x') xs'

data NetStat = NetStat {unNetStat :: Int} deriving (Show, Eq)

netStat :: Sensor NetStat
netStat = poll' "netStat" aggregate (5 * 10 ^ 5)
  where
    aggregate = liftIO do
      is <- listDirectory "/sys/class/net"
      NetStat . sum <$> mapM (aggregate' . ("/sys/class/net" </>)) is
    aggregate' fp = do
      (+)
        <$> (read <$> readFile (fp </> "statistics/rx_bytes"))
        <*> (read <$> readFile (fp </> "statistics/tx_bytes"))

net' :: Sensor [NetStat]
net' = histogram "net'" [] (\b a -> take 7 (a : b)) netStat

data Net = Net {unNet :: [Float]} deriving (Show, Eq)

instance P.Pretty Net where
  pretty = P.diagram 3 . unNet

net :: Sensor Net
net = transform "net" f net'
  where
    f (map unNetStat -> xs) = do
      let xs' =
            zipWith
              (-)
              (fromMaybe [] (initMay xs))
              (fromMaybe [] (tailMay xs))
          x' = max 1 (fromMaybe 0 (maximumMay xs'))
      pure $ Net $ map (\x -> fromIntegral x / fromIntegral x') xs'

data MemStat = MemStat {unMemStat :: Float} deriving (Show, Eq)

memStat :: Sensor MemStat
memStat = pollFile "memStat" aggregate "/proc/meminfo" (5 * 10 ^ 5)
  where
    aggregate fp = liftIO do
      meminfo <- readFile fp
      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 $ MemStat (1 - avail / total)
        (Nothing, _) -> error (printf "%s: MemTotal missing" fp)
        (_, Nothing) -> error (printf "%s: MemAvail missing" fp)

mem' :: Sensor [MemStat]
mem' = histogram "mem'" [] (\b a -> take 7 (a : b)) memStat

data Mem = Mem {unMem :: [Float]} deriving (Show, Eq)

instance P.Pretty Mem where
  pretty = P.diagram 3 . unMem

mem :: Sensor Mem
mem = transform "mem" (\(map unMemStat -> xs) -> pure (Mem xs)) mem'

data BatStat = BatStat {unBatStat :: Float} deriving (Show, Eq)

batStat :: Sensor BatStat
batStat = poll' "batStat" aggregate (5 * 10 ^ 5)
  where
    aggregate = liftIO do
      fmap (BatStat . product) . mapM (aggregate' . ("/sys/class/power_supply" </>))
        =<< listDirectory "/sys/class/power_supply"
    aggregate' fp =
      choice 1 $
        [ (/)
            <$> (read <$> readFile (fp </> "charge_now"))
            <*> (read <$> readFile (fp </> "charge_full")),
          (/)
            <$> (read <$> readFile (fp </> "energy_now"))
            <*> (read <$> readFile (fp </> "energy_full"))
        ]

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

bat' :: Sensor [BatStat]
bat' = histogram "bat'" [] (\b a -> take 7 (a : b)) batStat

data Bat = Bat {unBat :: [Float]} deriving (Show, Eq)

instance P.Pretty Bat where
  pretty = P.diagram 1 . unBat

bat :: Sensor Bat
bat = transform "bat" (\(map unBatStat -> xs) -> pure (Bat xs)) bat'

data Snd = Snd {unSnd :: Float} deriving (Show, Eq)

instance P.Pretty Snd where
  pretty = P.diagram 1 . (: []) . unSnd

snd :: Sensor Snd
snd = Sensor "snd" $ memoizeMVar do
  stateT <- newTVarIO (Snd 1)
  void $ forkIO $ forever do
    x <- liftIO aggregate
    atomically do writeTVar stateT (Snd x)
    threadDelay (2 * 10 ^ 6)
  pure stateT
  where
    aggregate = do
      (/ 153) . read <$> [sh|pamixer --get-volume|]

data Disk = Disk {unDisk :: [Float]} deriving (Show, Eq)

instance P.Pretty Disk where
  pretty = P.diagram 1 . unDisk

disk :: Sensor Disk
disk = Sensor "disk" $ memoizeMVar do
  stateT <- newTVarIO (Disk [])
  void $ forkIO $ forever do
    x <- liftIO aggregate
    atomically do
      readTVar stateT >>= \(Disk xs) ->
        writeTVar stateT (Disk (x : xs))
    threadDelay (2 * 10 ^ 6)
  pure stateT
  where
    aggregate = do
      stat <- statVFS "/"
      pure $
        (fromIntegral stat.statVFS_bfree)
          / fromIntegral (stat.statVFS_bfree + stat.statVFS_bavail)

currentTime :: Sensor UTCTime
currentTime = Sensor "currentTime" $ memoizeMVar do
  stateT <- newTVarIO =<< liftIO getCurrentTime
  void $ forkIO $ forever do
    x <- liftIO getCurrentTime
    atomically . writeTVar stateT =<< liftIO getCurrentTime
    threadDelay (((10 ^ 12) - (fromEnum (utctDayTime x) `mod` (10 ^ 12))) `div` (10 ^ 6))
  pure stateT

data Date = Date {unDate :: String} deriving (Show, Eq)

instance P.Pretty Date where
  pretty = P.color P.White . P.pretty . unDate

date :: Sensor Date
date = transform "date" f currentTime
  where
    f x = do
      z <- liftIO getCurrentTimeZone
      pure . Date . formatTime defaultTimeLocale "%b %e" $
        utcToLocalTime z x

data Time = Time {unTime :: String} deriving (Show, Eq)

instance P.Pretty Time where
  pretty = P.color P.White . P.pretty . unTime

time :: Sensor Time
time = transform "time" f currentTime
  where
    f x = do
      z <- liftIO getCurrentTimeZone
      pure . Time . formatTime defaultTimeLocale "%R" $
        utcToLocalTime z x

data WmName = WmName {unWmName :: String} deriving (Show, Eq)

instance P.Pretty WmName where
  pretty = P.color P.White . P.pretty . unWmName

wmName' :: Sensor WmName
wmName' = watchXPropertyChanges "wmName'" as aggregate
  where
    as = ["_NET_ACTIVE_WINDOW"]
    aggregate dpy = liftIO do
      let root = X.rootWindowOfScreen (X.defaultScreenOfDisplay dpy)
      netActiveWindow <- X.internAtom dpy "_NET_ACTIVE_WINDOW" False
      X.getWindowProperty32 dpy netActiveWindow root >>= \case
        Just (win : _) -> do
          netWmName <- X.internAtom dpy "_NET_WM_NAME" False
          fmap (WmName . strip . head) . X.wcTextPropertyToTextList dpy
            =<< do
              let loop 0 = do
                    error "failed to read _NET_WM_NAME"
                  loop n = do
                    catch
                      (X.getTextProperty dpy (fromIntegral win) netWmName)
                      ( \(e :: IOError) -> do
                          if isUserError e
                            then threadDelay 1000 >> loop (n - 1)
                            else throwIO e
                      )
              loop 16
        _ -> pure (WmName "")
    strip = reverse . dropWhile isSpace . reverse . dropWhile isSpace

wmName :: Sensor String
wmName = transform "wmName" (\(WmName x) -> pure x) wmName'

data WmWorkspaces = WmWorkspaces {unWmWorkspaces :: [Workspace]}
  deriving (Show, Eq)

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

instance P.Pretty WmWorkspaces where
  pretty =
    P.Col . intersperse (P.pretty " ") . map P.pretty . unWmWorkspaces

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

wmWorkspaces :: Sensor WmWorkspaces
wmWorkspaces = watchXPropertyChanges "wmWorkspaces" as aggregate
  where
    as =
      [ "_NET_CLIENT_LIST",
        "_NET_CURRENT_DESKTOP",
        "_NET_DESKTOP_NAMES",
        "_NET_WM_DESKTOP"
      ]
    aggregate dpy = liftIO do
      let root = X.rootWindowOfScreen (X.defaultScreenOfDisplay dpy)
      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
      currentDesktop <- fmap (fromIntegral . head) <$> X.getWindowProperty32 dpy netCurrentDesktop root
      occupiedDesktops <-
        fmap (S.unions . catMaybes)
          . mapM
            ( \win -> do
                fmap (S.singleton . fromIntegral . head)
                  <$> X.getWindowProperty32 dpy netWmDesktop win
            )
          . map fromIntegral
          . fromMaybe []
          =<< X.getWindowProperty32 dpy netClientList root
      fmap
        ( WmWorkspaces
            . 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

watchFile :: String -> (FilePath -> SensorM a) -> FilePath -> Sensor a
watchFile id aggregate fp = Sensor id $ memoizeMVar do
  i <- liftIO do
    i <- I.init
    void (I.addWatch i fp I.in_CLOSE)
    pure i
  stateT <- newTVarIO =<< aggregate fp
  void $ forkIO $ forever do
    void (liftIO (I.getEvent i))
    -- liftIO (printf "[%s] aggregate\n" id)
    x <- aggregate fp
    atomically (writeTVar stateT x)
  pure stateT

pollFile :: String -> (FilePath -> SensorM a) -> FilePath -> Int -> Sensor a
pollFile id aggregate fp delay = Sensor id $ memoizeMVar do
  stateT <- newTVarIO =<< aggregate fp
  void $ forkIO $ forever do
    -- liftIO (printf "[%s] aggregate\n" id)
    x <- aggregate fp
    threadDelay delay
    atomically (writeTVar stateT x)
  pure stateT

poll' :: String -> SensorM a -> Int -> Sensor a
poll' id aggregate delay = Sensor id $ memoizeMVar do
  stateT <- newTVarIO =<< aggregate
  void $ forkIO $ forever do
    -- liftIO (printf "[%s] aggregate\n" id)
    x <- aggregate
    threadDelay delay
    atomically (writeTVar stateT x)
  pure stateT

watchXPropertyChanges ::
  String ->
  [String] ->
  (X.Display -> SensorM a) ->
  Sensor a
watchXPropertyChanges id as' aggregate = Sensor id $ memoizeMVar do
  (dpy, as) <- liftIO do
    dpy <- X.openDisplay ""
    let root = X.rootWindowOfScreen (X.defaultScreenOfDisplay dpy)
    X.selectInput dpy root X.propertyChangeMask
    as <- S.fromList <$> mapM (\a -> X.internAtom dpy a False) as'
    pure (dpy, as)
  stateT <- newTVarIO =<< aggregate dpy
  void $ forkIO $ forever $ do
    as' <- liftIO $ X.allocaXEvent $ \ev -> do
      X.nextEvent dpy ev
      e <- X.getEvent ev
      let loop es = do
            isTimeout <- X.waitForEvent dpy 0
            if isTimeout
              then pure (S.fromList (map (X.ev_atom) es))
              else do
                X.nextEvent dpy ev
                e <- X.getEvent ev
                loop (e : es)
      loop [e]
    unless (S.null (as `S.intersection` as')) do
      atomically . writeTVar stateT =<< aggregate dpy
  pure stateT

histogram ::
  (Eq a, Typeable a) =>
  String ->
  b ->
  (b -> a -> b) ->
  Sensor a ->
  Sensor b
histogram id def acc sensor = transform' id def f sensor
  where
    f b a = pure (acc b a)

transform' ::
  (Eq a, Typeable a) =>
  String ->
  b ->
  (b -> a -> SensorM b) ->
  Sensor a ->
  Sensor b
transform' id def fM sensor = Sensor id $ memoizeMVar do
  xT <- poll sensor
  x <- readTVarIO xT
  x'T <- newTVarIO x
  stateT <- newTVarIO =<< fM def x
  void $ forkIO $ forever $ do
    (def, x) <- atomically do
      x' <- readTVar x'T
      x <- readTVar xT
      checkSTM (x /= x')
      def <- readTVar stateT
      writeTVar x'T x
      pure (def, x)
    atomically . writeTVar stateT =<< fM def x
  pure stateT

transform ::
  (Eq a, Typeable a) =>
  String ->
  (a -> SensorM b) ->
  Sensor a ->
  Sensor b
transform id fM sensor = Sensor id $ memoizeMVar do
  xT <- poll sensor
  x <- readTVarIO xT
  x'T <- newTVarIO x
  stateT <- newTVarIO =<< fM x
  void $ forkIO $ forever $ do
    x <- atomically do
      x' <- readTVar x'T
      x <- readTVar xT
      checkSTM (x /= x')
      writeTVar x'T x
      pure x
    atomically . writeTVar stateT =<< fM x
  pure stateT