{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

module Main where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Control.Monad.Identity
import Data.Bits
import Data.List
import Data.Map qualified as M
import Data.Maybe
import GHC.Ptr (Ptr)
import Graphics.X11 qualified as X
import Graphics.X11.Xft qualified as X
import Graphics.X11.Xlib.Extras qualified as X
import Options.Applicative qualified as O
import Pretty qualified as P
import Sensor
import Ui hiding (Env)
import Ui qualified
import UnliftIO.STM

data Args = Args
  { icons :: Bool,
    spacing :: Int
  }

args :: O.ParserInfo Args
args = O.info (Args <$> iconsArg <*> spacingArg) O.idm

iconsArg :: O.Parser Bool
iconsArg = O.switch (O.long "icons")

spacingArg :: O.Parser Int
spacingArg = O.option O.auto (O.long "spacing" <> O.value 0)

data Env = Env
  { dpy :: X.Display,
    win :: X.Window,
    swidth :: Int,
    sheight :: Int,
    wwidth :: Int,
    wheight :: Int,
    wleft :: Int,
    gc :: X.GC,
    pixm :: X.Pixmap,
    fnt :: X.XftFont,
    drw :: X.XftDraw,
    cmap :: X.Colormap,
    vis :: X.Visual,
    xcolors :: M.Map P.XColor String
  }

type Colors = M.Map (P.Intensity, P.Color) X.XftColor

data State = State
  { dirty :: Bool,
    ui :: Ui TVar P.Doc
  }

main :: IO ()
main = do
  O.execParser args >>= \args -> do
    bracket (createWindow args) destroyWindow $ \(env, stateT) ->
      withColors env (run env stateT)

data LastRun = LastRun
  { pUi :: Ui Identity P.Doc,
    lUi :: Ui'
  }

run :: Env -> TVar State -> Colors -> IO ()
run env stateT colors = do
  let loop = (loop . Just =<<) . go
  loop Nothing
  where
    go Nothing = do
      pUi <- atomically do
        state@State {..} <- readTVar stateT
        writeTVar stateT state {dirty = False}
        pollUi ui
      lUi <- paint env colors Nothing pUi
      pure (LastRun pUi lUi)
    go (Just (LastRun {pUi = pUi', lUi = lUi'})) = do
      (pUi, dirty) <- atomically do
        state@State {..} <- readTVar stateT
        pUi <- pollUi ui
        checkSTM (dirty || pUi /= pUi')
        writeTVar stateT state {dirty = False}
        pure (pUi, dirty)
      lUi <- paint env colors (if dirty then Nothing else Just lUi') pUi
      pure (LastRun pUi lUi)

processEvents :: Ptr X.XEvent -> Env -> TVar State -> IO ()
processEvents ev env@(Env {..}) stateT = do
  timeOut <- X.waitForEvent dpy 1_000_000_000
  unless timeOut do
    processEvent ev env stateT
    processEvents ev env stateT

processEvent :: Ptr X.XEvent -> Env -> TVar State -> IO ()
processEvent ev (Env {..}) stateT = do
  X.nextEvent dpy ev
  e <- X.getEvent ev
  if
    | X.ExposeEvent {} <- e -> atomically do
        state <- readTVar stateT
        writeTVar stateT state {dirty = True}
    | otherwise -> pure ()

paint ::
  Env ->
  Colors ->
  Maybe Ui' ->
  Ui Identity P.Doc ->
  IO Ui'
paint Env {..} colors Nothing ui = do
  let env' = Ui.Env {..}
  ui' <- layOutUi env' ui
  renderUi env' colors ui'
  X.copyArea dpy pixm win gc 0 0 (fi wwidth) (fi wheight) 0 0
  X.sync dpy False
  pure ui'
paint Env {..} colors (Just ui') ui = do
  let env' = Ui.Env {..}
  ui'' <- layOutUi env' ui
  renderUi' env' colors ui' ui''
  X.copyArea dpy pixm win gc 0 0 (fi wwidth) (fi wheight) 0 0
  X.sync dpy False
  pure ui''

destroyWindow :: (Env, TVar State) -> IO ()
destroyWindow (Env {..}, _) = do
  X.destroyWindow dpy win

createWindow :: Args -> IO (Env, TVar State)
createWindow args = do
  dpy <- X.openDisplay ""
  let scrn = X.defaultScreen dpy
      scr = X.defaultScreenOfDisplay dpy
      root = X.defaultRootWindow dpy
      trueColor = 4
  Just vinfo <- X.matchVisualInfo dpy scrn 32 trueColor
  let cls = X.inputOutput
      dpth = X.visualInfo_depth vinfo
      vis = X.visualInfo_visual vinfo
      vmsk = X.cWColormap .|. X.cWBorderPixel .|. X.cWBackingPixel .|. X.cWOverrideRedirect
      swidth = fi (X.displayWidth dpy scrn)
      sheight = fi (X.displayHeight dpy scrn)
      wwidth = swidth - 16
      wspac = args.spacing
      wleft = wspac
  fnt <- X.xftFontOpen dpy scr "IosevkaTerm Nerd Font:size=14"
  wheight <- X.xftfont_height fnt
  cmap <- X.createColormap dpy root vis X.allocNone
  win <- X.allocaSetWindowAttributes $ \attr -> do
    X.set_colormap attr cmap
    X.set_border_pixel attr 0
    X.set_background_pixel attr 0
    X.set_override_redirect attr True
    X.createWindow dpy root (fi wleft) 0 (fi wwidth) (fi wheight) 0 dpth cls vis vmsk attr
  atom <- X.internAtom dpy "ATOM" True
  wmState <- X.internAtom dpy "_NET_WM_STATE" False
  wmStateSticky <- X.internAtom dpy "_NET_WM_STATE_STICKY" False
  wmStateAbove <- X.internAtom dpy "_NET_WM_STATE_ABOVE" False
  wmWindowType <- X.internAtom dpy "_NET_WM_WINDOW_TYPE" False
  wmWindowTypeDock <- X.internAtom dpy "_NET_WM_WINDOW_TYPE_DOCK" False
  wmStrut <- X.internAtom dpy "_NET_WM_STRUT" False
  wmStrutPartial <- X.internAtom dpy "_NET_WM_STRUT_PARTIAL" False
  X.changeProperty32 dpy win wmState atom X.propModeReplace $
    [ fi wmStateAbove,
      fi wmStateSticky
    ]
  X.changeProperty32 dpy win wmWindowType atom X.propModeReplace $
    [ fi wmWindowTypeDock
    ]
  X.changeProperty32 dpy win wmStrut atom X.propModeReplace $
    [0, 0, fi (wheight + wspac), 0]
  X.changeProperty32 dpy win wmStrutPartial atom X.propModeReplace $
    [0, 0, fi (wheight + wspac), 0, 0, 0, 0, 0, fi wleft, fi wwidth, 0, 0]
  pixm <- X.createPixmap dpy win (fi wwidth) (fi wheight) dpth
  gc <- X.createGC dpy win
  drw <- X.xftDrawCreate dpy pixm vis cmap
  X.mapWindow dpy win
  let dirty = True
  ui <-
    runSensorM . initUi . Ui $
      intercalate [lit " "] $
        [ [sens wmWorkspaces],
          [sens wmName, fill],
          [lit (if args.icons then "\xf4bc  " else "cpu "), sens cpu],
          [lit (if args.icons then "\xf035b  " else "mem "), sens mem],
          [lit (if args.icons then "\xf0a0  " else "disk "), sens disk],
          [lit (if args.icons then "\xf1638  " else "io "), sens io],
          [lit (if args.icons then "\xf0200  " else "net "), sens net],
          [lit (if args.icons then "\xf028  " else "snd "), sens Sensor.snd],
          [lit (if args.icons then "\xf240  " else "bat "), sens bat],
          [lit " ", sens date, lit ", ", sens time]
        ]
  xcolors <- do
    X.rmInitialize
    rdb <- X.rmGetStringDatabase (X.resourceManagerString dpy)
    M.fromList
      <$> mapM
        ( \(n, def) -> do
            (toEnum n,) . fromMaybe def
              <$> X.rmGetResource rdb ("astatusbar*color" ++ show n)
        )
        [ (0 {- black -}, "rgb:0/0/0"),
          (1 {- red3 -}, "rgb:205/0/0"),
          (2 {- green3 -}, "rgb:0/205/0"),
          (3 {- yellow3 -}, "rgb:205/205/0"),
          (4 {- blue2 -}, "rgb:0/0/238"),
          (5 {- meganta3 -}, "rgb:205/0/205"),
          (6 {- cyan3 -}, "rgb:0/205/205"),
          (7 {- gray90 -}, "rgb:229/229/229"),
          (8 {- gray50 -}, "rgb:127/127/127"),
          (9 {- red -}, "rgb:255/0/0"),
          (10 {- green -}, "rgb:0/255/0"),
          (11 {- yellow -}, "rgb:255/255/0"),
          (12, "rgb:5c/5c/ff"),
          (13 {- meganta -}, "rgb:255/0/255"),
          (14 {- cyan -}, "rgb:0/255/255"),
          (15 {- white -}, "rgb:255/255/255")
        ]
  let env = Env {..}
  stateT <- newTVarIO State {..}
  void $ forkIO $ X.allocaXEvent $ \ev -> forever do
    processEvents ev env stateT
  pure (env, stateT)

withColors :: Env -> (Colors -> IO a) -> IO a
withColors Env {..} act = do
  X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Vivid, P.Black)) $ \black -> do
    X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Vivid, P.Red)) $ \red -> do
      X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Vivid, P.Green)) $ \green -> do
        X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Vivid, P.Yellow)) $ \yellow -> do
          X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Vivid, P.Blue)) $ \blue -> do
            X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Vivid, P.Magenta)) $ \magenta -> do
              X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Vivid, P.Cyan)) $ \cyan -> do
                X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Vivid, P.White)) $ \white -> do
                  X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Dull, P.Black)) $ \dullBlack -> do
                    X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Dull, P.Red)) $ \dullRred -> do
                      X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Dull, P.Green)) $ \dullGreen -> do
                        X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Dull, P.Yellow)) $ \dullYellow -> do
                          X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Dull, P.Blue)) $ \dullBlue -> do
                            X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Dull, P.Magenta)) $ \dullMagenta -> do
                              X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Dull, P.Cyan)) $ \dullCyan -> do
                                X.withXftColorName dpy vis cmap (xcolors M.! P.toXColor (P.Dull, P.White)) $ \dullWhite -> do
                                  let colors =
                                        M.fromList
                                          [ ((P.Vivid, P.Black), black),
                                            ((P.Vivid, P.Red), red),
                                            ((P.Vivid, P.Green), green),
                                            ((P.Vivid, P.Yellow), yellow),
                                            ((P.Vivid, P.Blue), blue),
                                            ((P.Vivid, P.Magenta), magenta),
                                            ((P.Vivid, P.Cyan), cyan),
                                            ((P.Vivid, P.White), white),
                                            ((P.Dull, P.Black), dullBlack),
                                            ((P.Dull, P.Red), dullRred),
                                            ((P.Dull, P.Green), dullGreen),
                                            ((P.Dull, P.Yellow), dullYellow),
                                            ((P.Dull, P.Blue), dullBlue),
                                            ((P.Dull, P.Magenta), dullMagenta),
                                            ((P.Dull, P.Cyan), dullCyan),
                                            ((P.Dull, P.White), dullWhite)
                                          ]

                                  act colors

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

clamp :: (Ord a) => a -> a -> a -> a
clamp mi ma = max mi . min ma