{-# LANGUAGE OverloadedStrings #-}

-----------------------------------------------------------------------------

-----------------------------------------------------------------------------

-- |
-- Module      : System.Taffybar.Information.Hyprland
-- Copyright   : (c) Ivan A. Malison
-- License     : BSD3-style (see LICENSE)
--
-- Maintainer  : Ivan A. Malison
-- Stability   : unstable
-- Portability : unportable
--
-- A small Hyprland "client" that provides a structured interface for issuing
-- Hyprland commands (optionally over the Hyprland command socket) and for
-- connecting to the Hyprland event socket.
--
-- This module is intended to centralize the socket/path logic so widgets can
-- share it. It is not (yet) wired into widgets.
module System.Taffybar.Information.Hyprland
  ( -- * Client
    HyprlandClient,
    HyprlandClientConfig (..),
    defaultHyprlandClientConfig,
    newHyprlandClient,
    reloadHyprlandClient,
    HyprlandClientEnv (..),
    getHyprlandClientEnv,

    -- * Hyprland Monad
    HyprlandT,
    runHyprlandT,
    askHyprlandClient,
    runHyprlandCommandRawM,
    runHyprlandCommandJsonM,

    -- * Shared Event Channel
    HyprlandEventChan (..),
    subscribeHyprlandEvents,
    buildHyprlandEventChan,

    -- * Commands
    HyprlandCommand (..),
    hyprCommand,
    hyprCommandJson,
    hyprlandCommandToSocketCommand,
    runHyprlandCommandRaw,
    runHyprlandCommandJson,

    -- * Sockets
    HyprlandSocket (..),
    hyprlandSocketName,
    hyprlandSocketPaths,
    openHyprlandSocket,
    openHyprlandEventSocket,
    withHyprlandEventSocket,

    -- * Monitor queries
    getFocusedMonitorPosition,

    -- * Errors
    HyprlandError (..),
  )
where

import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.STM.TChan
  ( TChan,
    dupTChan,
    newBroadcastTChanIO,
    writeTChan,
  )
import Control.Exception.Enclosed (catchAny)
import Control.Monad (forM, forever, void)
import Control.Monad.IO.Class (MonadIO (..))
import Control.Monad.STM (atomically)
import Control.Monad.Trans.Reader (ReaderT, ask, runReaderT)
import Data.Aeson (FromJSON (..), eitherDecode', withObject, (.:))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import Data.List (sortOn)
import Data.Ord (Down (..))
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Network.Socket as NS
import qualified Network.Socket.ByteString as NSB
import System.Directory (doesDirectoryExist, listDirectory)
import System.Environment (lookupEnv)
import System.FilePath (takeDirectory, takeFileName, (</>))
import System.IO
  ( BufferMode (LineBuffering),
    Handle,
    IOMode (ReadWriteMode),
    hClose,
    hGetLine,
    hSetBuffering,
  )
import System.Log.Logger (Priority (..), logM)
import System.Posix.Files (getFileStatus, isSocket, modificationTime)
import System.Posix.Types (EpochTime)
import System.Taffybar.Util (runCommand)
import Text.Printf (printf)

-- | Errors that can occur while talking to Hyprland.
data HyprlandError
  = HyprlandEnvMissing String
  | HyprlandSocketUnavailable HyprlandSocket [FilePath]
  | HyprlandSocketException FilePath String
  | HyprlandHyprctlFailed String
  | HyprlandJsonDecodeFailed String
  | HyprlandCommandBuildFailed String
  deriving (Show, Eq)

-- | Configuration for creating and using a 'HyprlandClient'.
data HyprlandClientConfig = HyprlandClientConfig
  { -- | Try the Hyprland command/event sockets when possible.
    useSocket :: Bool,
    -- | If the socket path is unavailable, fall back to invoking @hyprctl@.
    fallbackToHyprctl :: Bool,
    hyprctlPath :: FilePath
  }
  deriving (Show, Eq)

-- | Default 'HyprlandClientConfig'.
--
-- Uses sockets when available and falls back to @hyprctl@.
defaultHyprlandClientConfig :: HyprlandClientConfig
defaultHyprlandClientConfig =
  HyprlandClientConfig
    { useSocket = True,
      fallbackToHyprctl = True,
      hyprctlPath = "hyprctl"
    }

-- | Environment-derived values used for socket path resolution.
data HyprlandClientEnv = HyprlandClientEnv
  { instanceSignature :: String,
    runtimeDir :: Maybe FilePath
  }
  deriving (Show, Eq)

-- | Read Hyprland-specific values from the process environment.
getHyprlandClientEnv :: IO (Either HyprlandError HyprlandClientEnv)
getHyprlandClientEnv = do
  mSig <- lookupEnv "HYPRLAND_INSTANCE_SIGNATURE"
  case mSig of
    Nothing ->
      return $ Left $ HyprlandEnvMissing "HYPRLAND_INSTANCE_SIGNATURE"
    Just sig -> do
      mRuntime <- lookupEnv "XDG_RUNTIME_DIR"
      return $
        Right $
          HyprlandClientEnv
            { instanceSignature = sig,
              runtimeDir = mRuntime
            }

-- | Runtime client state used by command and event helpers.
data HyprlandClient = HyprlandClient
  { clientConfig :: HyprlandClientConfig,
    clientEnv :: Maybe HyprlandClientEnv
  }
  deriving (Show, Eq)

-- | Construct a 'HyprlandClient' by reading environment variables.
--
-- If 'HYPRLAND_INSTANCE_SIGNATURE' is not set, socket-based operations will be
-- unavailable, but command execution may still succeed via the @hyprctl@
-- fallback depending on 'HyprlandClientConfig'.
newHyprlandClient :: HyprlandClientConfig -> IO HyprlandClient
newHyprlandClient cfg = do
  envResult <- getHyprlandClientEnv
  let env = either (const Nothing) Just envResult
  pure $ HyprlandClient {clientConfig = cfg, clientEnv = env}

-- | Reload the environment-derived portions of a 'HyprlandClient'.
reloadHyprlandClient :: HyprlandClient -> IO HyprlandClient
reloadHyprlandClient client = newHyprlandClient (clientConfig client)

-- | A reader transformer that carries a 'HyprlandClient'.
type HyprlandT m a = ReaderT HyprlandClient m a

-- | Run a 'HyprlandT' computation with the provided client.
runHyprlandT :: HyprlandClient -> HyprlandT m a -> m a
runHyprlandT = flip runReaderT

-- | Get the current 'HyprlandClient' from 'HyprlandT'.
askHyprlandClient :: (Monad m) => HyprlandT m HyprlandClient
askHyprlandClient = ask

-- | Run a command in 'HyprlandT' and return the raw bytes response.
runHyprlandCommandRawM :: (MonadIO m) => HyprlandCommand -> HyprlandT m (Either HyprlandError BS.ByteString)
runHyprlandCommandRawM cmd = do
  client <- ask
  liftIO $ runHyprlandCommandRaw client cmd

-- | Run a command in 'HyprlandT' and decode the response as JSON.
runHyprlandCommandJsonM :: (MonadIO m, FromJSON a) => HyprlandCommand -> HyprlandT m (Either HyprlandError a)
runHyprlandCommandJsonM cmd = do
  client <- ask
  liftIO $ runHyprlandCommandJson client cmd

-- | Which Hyprland socket should be used for a request.
data HyprlandSocket
  = HyprlandCommandSocket
  | HyprlandEventSocket
  deriving (Show, Eq)

-- | Socket file basename for a given 'HyprlandSocket' kind.
hyprlandSocketName :: HyprlandSocket -> FilePath
hyprlandSocketName HyprlandCommandSocket = ".socket.sock"
hyprlandSocketName HyprlandEventSocket = ".socket2.sock"

-- | Candidate socket paths for a client and socket type.
hyprlandSocketPaths :: HyprlandClient -> HyprlandSocket -> [FilePath]
hyprlandSocketPaths HyprlandClient {clientEnv = Nothing} _ = []
hyprlandSocketPaths
  HyprlandClient
    { clientEnv =
        Just
          HyprlandClientEnv
            { instanceSignature = sig,
              runtimeDir = mRuntimeDir
            }
    }
  sock =
    let name = hyprlandSocketName sock
        runtimePaths =
          case mRuntimeDir of
            Nothing -> []
            Just rd -> [rd ++ "/hypr/" ++ sig ++ "/" ++ name]
        tmpPath = "/tmp/hypr/" ++ sig ++ "/" ++ name
     in runtimePaths ++ [tmpPath]

-- | Open a Hyprland unix socket, trying known and discovered locations.
openHyprlandSocket :: HyprlandClient -> HyprlandSocket -> IO (Either HyprlandError NS.Socket)
openHyprlandSocket client sock = do
  -- Prefer the instance signature from the process environment first. If that
  -- fails (e.g. Hyprland restarted and the signature changed), fall back to
  -- discovering currently-running instance sockets from the filesystem.
  let envPaths = hyprlandSocketPaths client sock
  envResult <- connectFirst envPaths
  case envResult of
    Just s -> pure (Right s)
    Nothing -> do
      discoveredPaths <- discoverHyprlandSocketPaths client sock
      discResult <- connectFirst discoveredPaths
      case discResult of
        Just s -> pure (Right s)
        Nothing ->
          if null envPaths && null discoveredPaths
            then pure $ Left $ HyprlandEnvMissing "HYPRLAND_INSTANCE_SIGNATURE"
            else pure $ Left $ HyprlandSocketUnavailable sock (envPaths ++ discoveredPaths)
  where
    connectFirst :: [FilePath] -> IO (Maybe NS.Socket)
    connectFirst [] = pure Nothing
    connectFirst (path : rest) = do
      result <- connectToSocket path
      case result of
        Left _ -> connectFirst rest
        Right s -> pure (Just s)

discoverHyprlandSocketPaths :: HyprlandClient -> HyprlandSocket -> IO [FilePath]
discoverHyprlandSocketPaths _client sock = do
  mRuntime <- lookupEnv "XDG_RUNTIME_DIR"
  let bases =
        maybe [] (\rd -> [rd </> "hypr"]) mRuntime
          ++ ["/tmp/hypr"]
  let name = hyprlandSocketName sock
  pathsWithTimes <- fmap concat $ forM bases $ \base -> do
    baseExists <- doesDirectoryExist base
    if not baseExists
      then pure []
      else do
        -- Best-effort: Hyprland instances come and go and we don't want to
        -- bring down widgets if this scan races a restart.
        sigEntries <- listDirectory base `catchAny` \_ -> pure []
        let candidates = map (\sig -> base </> sig </> name) sigEntries
        fmap concat $ forM candidates $ \p -> do
          mTime <- socketPathMTime p
          pure $ maybe [] (\t -> [(t, p)]) mTime

  -- Prefer the newest sockets first to minimize binding to a stale instance.
  pure $ map snd $ sortOn (Down . fst) pathsWithTimes

socketPathMTime :: FilePath -> IO (Maybe EpochTime)
socketPathMTime path =
  ( do
      st <- getFileStatus path
      if isSocket st
        then pure $ Just $ modificationTime st
        else pure Nothing
  )
    `catchAny` \_ -> pure Nothing

connectToSocket :: FilePath -> IO (Either HyprlandError NS.Socket)
connectToSocket path = do
  sock <- NS.socket NS.AF_UNIX NS.Stream NS.defaultProtocol
  ( do
      NS.connect sock (NS.SockAddrUnix path)
      return $ Right sock
    )
    `catchAny` \e -> do
      void $ NS.close sock `catchAny` \_ -> pure ()
      return $ Left $ HyprlandSocketException path (show e)

-- | Open the Hyprland event socket and return it as a line-buffered 'Handle'.
openHyprlandEventSocket :: HyprlandClient -> IO (Either HyprlandError Handle)
openHyprlandEventSocket client = do
  sockResult <- openHyprlandSocket client HyprlandEventSocket
  case sockResult of
    Left err -> return (Left err)
    Right sock -> do
      handle <- NS.socketToHandle sock ReadWriteMode
      hSetBuffering handle LineBuffering
      return (Right handle)

-- | Bracket-style helper for safely using the Hyprland event socket handle.
withHyprlandEventSocket :: HyprlandClient -> (Handle -> IO a) -> IO (Either HyprlandError a)
withHyprlandEventSocket client action = do
  handleResult <- openHyprlandEventSocket client
  case handleResult of
    Left err -> return (Left err)
    Right handle ->
      ( do
          result <- action handle
          hClose handle
          return (Right result)
      )
        `catchAny` \e -> do
          hClose handle
          return (Left $ HyprlandSocketException (show HyprlandEventSocket) (show e))

-- | A shared broadcast channel for Hyprland events read from the event socket.
--
-- Readers should call 'subscribeHyprlandEvents' to get their own cursor.
newtype HyprlandEventChan
  = HyprlandEventChan (TChan T.Text)

-- | Subscribe to a per-reader cursor for a shared Hyprland event channel.
subscribeHyprlandEvents :: HyprlandEventChan -> IO (TChan T.Text)
subscribeHyprlandEvents (HyprlandEventChan chan) =
  atomically $ dupTChan chan

-- | Create and maintain a shared Hyprland event broadcast channel.
buildHyprlandEventChan :: HyprlandClient -> IO HyprlandEventChan
buildHyprlandEventChan client = do
  chan <- newBroadcastTChanIO
  _ <- forkIO $ eventThread chan
  pure $ HyprlandEventChan chan
  where
    logH :: Priority -> String -> IO ()
    logH = logM "System.Taffybar.Information.Hyprland"

    retryDelayMicros :: Int
    retryDelayMicros = 1 * 1000000

    eventThread chan = forever $ do
      handleResult <- openHyprlandEventSocket client
      case handleResult of
        Left err -> do
          logH WARNING $ printf "Hyprland event socket unavailable: %s" (show err)
          threadDelay retryDelayMicros
        Right handle -> do
          -- Emit a synthetic event on (re)connect so widgets can refresh their
          -- state after Hyprland restarts without resorting to polling.
          atomically $ writeTChan chan "taffybar-hyprland-connected>>"
          let loop =
                hGetLine handle >>= (atomically . writeTChan chan . T.pack) >> loop
          loop `catchAny` \e ->
            logH WARNING $ printf "Hyprland event socket failed: %s" (show e)
          void $ hClose handle `catchAny` \_ -> pure ()
          threadDelay retryDelayMicros

-- | Minimal monitor info parsed from @hyprctl monitors -j@.
data HyprMonitorInfo = HyprMonitorInfo
  { hmX :: Int,
    hmY :: Int,
    hmFocused :: Bool
  }

instance FromJSON HyprMonitorInfo where
  parseJSON = withObject "HyprMonitorInfo" $ \o ->
    HyprMonitorInfo
      <$> o .: "x"
      <*> o .: "y"
      <*> o .: "focused"

-- | Get the global coordinates of the currently focused Hyprland monitor.
getFocusedMonitorPosition :: HyprlandClient -> IO (Maybe (Int, Int))
getFocusedMonitorPosition client = do
  result <- runHyprlandCommandJson client (hyprCommandJson ["monitors"])
  case result of
    Left _ -> return Nothing
    Right monitors ->
      case filter hmFocused monitors of
        (m : _) -> return $ Just (hmX m, hmY m)
        [] -> return Nothing

-- | Structured representation of a Hyprland command invocation.
data HyprlandCommand = HyprlandCommand
  { commandArgs :: [String],
    commandJson :: Bool
  }
  deriving (Show, Eq)

-- | Construct a non-JSON command.
hyprCommand :: [String] -> HyprlandCommand
hyprCommand args = HyprlandCommand {commandArgs = args, commandJson = False}

-- | Construct a command whose output should be JSON.
hyprCommandJson :: [String] -> HyprlandCommand
hyprCommandJson args = HyprlandCommand {commandArgs = args, commandJson = True}

-- | Encode a 'HyprlandCommand' into the bytes expected by the command socket.
hyprlandCommandToSocketCommand :: HyprlandCommand -> Either HyprlandError BS.ByteString
hyprlandCommandToSocketCommand HyprlandCommand {commandArgs = args, commandJson = isJson}
  | null args = Left $ HyprlandCommandBuildFailed "No Hyprland command provided"
  | isJson = Right $ TE.encodeUtf8 $ T.pack $ "j/" ++ unwords args
  | otherwise = Right $ TE.encodeUtf8 $ T.pack $ unwords args

-- | Run a Hyprland command, preferring the command socket if enabled.
--
-- If socket execution fails and 'fallbackToHyprctl' is enabled, @hyprctl@ will
-- be invoked as a fallback.
runHyprlandCommandRaw :: HyprlandClient -> HyprlandCommand -> IO (Either HyprlandError BS.ByteString)
runHyprlandCommandRaw
  client@HyprlandClient
    { clientConfig =
        HyprlandClientConfig
          { useSocket = useSock,
            fallbackToHyprctl = fallback,
            hyprctlPath = hyprctl
          }
    }
  cmd = do
    socketResult <-
      if useSock
        then runHyprlandCommandSocket client cmd
        else pure $ Left $ HyprlandSocketUnavailable HyprlandCommandSocket []
    case socketResult of
      Right out -> pure (Right out)
      Left sockErr ->
        if fallback
          then do
            hyprctlResult <- runHyprlandCommandHyprctl client hyprctl cmd
            pure $ case hyprctlResult of
              Right out -> Right out
              Left hyprctlErr ->
                Left $
                  HyprlandHyprctlFailed $
                    printf "%s (socket error: %s)" hyprctlErr (show sockErr)
          else pure (Left sockErr)

runHyprlandCommandSocket :: HyprlandClient -> HyprlandCommand -> IO (Either HyprlandError BS.ByteString)
runHyprlandCommandSocket client cmd = do
  sockResult <- openHyprlandSocket client HyprlandCommandSocket
  case sockResult of
    Left err -> pure (Left err)
    Right sock ->
      ( do
          cmdBytes <- case hyprlandCommandToSocketCommand cmd of
            Left e -> NS.close sock >> pure (Left e)
            Right b -> pure (Right b)
          case cmdBytes of
            Left e -> pure (Left e)
            Right b -> do
              NSB.sendAll sock b
              NS.shutdown sock NS.ShutdownSend
              resp <- recvAll sock
              NS.close sock
              pure (Right resp)
      )
        `catchAny` \e -> do
          NS.close sock
          pure $ Left $ HyprlandSocketException (show HyprlandCommandSocket) (show e)

runHyprlandCommandHyprctl :: HyprlandClient -> FilePath -> HyprlandCommand -> IO (Either String BS.ByteString)
runHyprlandCommandHyprctl client hyprctl HyprlandCommand {commandArgs = args, commandJson = isJson} = do
  mSig <- pickHyprlandInstanceSignature client
  let flags =
        ["-j" | isJson]
          ++ maybe [] (\sig -> ["-i", sig]) mSig
  result <- runCommand hyprctl (flags ++ args)
  pure $ case result of
    Left err -> Left err
    Right out -> Right $ TE.encodeUtf8 $ T.pack out

pickHyprlandInstanceSignature :: HyprlandClient -> IO (Maybe String)
pickHyprlandInstanceSignature client =
  case clientEnv client of
    Just HyprlandClientEnv {instanceSignature = sig} -> do
      -- Prefer the signature from the environment if we can actually connect
      -- to the command socket. (After a Hyprland restart the old socket path
      -- might still exist but be stale.)
      envAlive <-
        anyM canConnectSocket $
          hyprlandSocketPaths client HyprlandCommandSocket
      if envAlive
        then pure (Just sig)
        else discover
    Nothing -> discover
  where
    canConnectSocket path = do
      result <- connectToSocket path
      case result of
        Left _ -> pure False
        Right sock -> do
          void $ NS.close sock `catchAny` \_ -> pure ()
          pure True

    discover = do
      -- Use the newest discovered command socket.
      paths <- discoverHyprlandSocketPaths client HyprlandCommandSocket
      pure $ case paths of
        p : _ -> Just $ takeFileName $ takeDirectory p
        [] -> Nothing

anyM :: (Monad m) => (a -> m Bool) -> [a] -> m Bool
anyM _ [] = pure False
anyM p (x : xs) = do
  b <- p x
  if b then pure True else anyM p xs

recvAll :: NS.Socket -> IO BS.ByteString
recvAll sock = go []
  where
    go acc = do
      chunk <- NSB.recv sock 4096
      if BS.null chunk
        then return (BS.concat (reverse acc))
        else go (chunk : acc)

-- | Run a command and decode the response body as JSON.
runHyprlandCommandJson :: (FromJSON a) => HyprlandClient -> HyprlandCommand -> IO (Either HyprlandError a)
runHyprlandCommandJson client cmd = do
  raw <- runHyprlandCommandRaw client cmd
  pure $ case raw of
    Left err -> Left err
    Right out ->
      case eitherDecode' (BL.fromStrict out) of
        Left decodeErr -> Left $ HyprlandJsonDecodeFailed decodeErr
        Right a -> Right a
