{-# LANGUAGE OverloadedStrings #-}
module Main where

import Control.Monad
import Control.Applicative
import Control.Concurrent
import Control.Concurrent.STM
import Data.Monoid
import Data.Char
import System.INotify ( initINotify, EventVariety(Modify), addWatch )
import Data.Word      ( Word8 )
import Data.Text      ( Text )
import Data.Map       ( Map )
import Data.List      ( foldl' )
import qualified Data.Map as Map
import qualified Data.Traversable as Traversable
import qualified Data.Text as Text
import qualified Data.Text.IO as Text
import qualified Network.BSD as BSD

import WaitForSignal  ( waitForTermSignal )
import UTmp           ( users2, utmp_file, UtmpRecord(..), UT_Type(..) )
import FGConsole      ( monitorTTY )

data ConsoleState = ConsoleState
    { csActiveTTY :: TVar Word8
    , csUtmp :: TVar (Map Text (TVar (Maybe UtmpRecord)))
    }

newConsoleState = atomically $
    ConsoleState <$> newTVar 0 <*> newTVar Map.empty


onLogin cs start = \e -> do
    us <- UTmp.users2
    let (m,cruft) =
            foldl' (\(m,cruft) x ->
                        case utmpType x of
                         USER_PROCESS
                           -> (Map.insert (utmpTty x) x m,cruft)
                         DEAD_PROCESS | utmpPid x /= 0
                           -> (m,Map.insert (utmpTty x) x cruft)
                         _ -> (m,cruft))
                   (Map.empty,Map.empty)
                   us
    forM_ (Map.elems cruft) $ \c -> do
        putStrLn $ "cruft " ++ show (utmpTty c, utmpPid c,utmpHost c, utmpRemoteAddr c)
    newborn <- atomically $ do
        old <- readTVar (csUtmp cs) -- swapTVar (csUtmp cs) m
        newborn <- flip Traversable.mapM (m Map.\\ old)
                    $ newTVar . Just
        updated <- let upd v u = writeTVar v $ Just u
                   in Traversable.sequence $ Map.intersectionWith upd old m
        let dead = old Map.\\ m
        Traversable.mapM (flip writeTVar Nothing) dead
        writeTVar (csUtmp cs) $ (old `Map.union` newborn) Map.\\ dead
        return newborn
    let getActive = do
            tty <- readTVar $ csActiveTTY cs
            utmp <- readTVar $ csUtmp cs
            flip (maybe $ return (tty,Nothing))
                 (Map.lookup ("tty"<>tshow tty) utmp)
                $ \tuvar -> do
            tu <- readTVar tuvar
            return (tty,tu)
            
    forM_ (Map.elems newborn) $ 
        forkIO . start getActive
    -- forM_ (Map.elems dead   ) $ putStrLn . ("gone: "++) . show

onTTY outvar cs vtnum = do
    logit outvar $ "switch: " <> tshow vtnum
    atomically $ writeTVar (csActiveTTY cs) vtnum

retryWhen var pred = do
    value <- var
    if pred value then retry
                  else return value

tshow x = Text.pack . show $ x

resource :: UtmpRecord -> Text
resource u =
    case utmpTty u of
        s | Text.take 3 s == "tty" -> s
        s | Text.take 4 s == "pts/" -> "pty" <> Text.drop 4 s <> ":" <> utmpHost u
        s -> escapeR s <> ":" <> utmpHost u
 where
    escapeR s = s

textHostName = fmap Text.pack BSD.getHostName

ujid u = do
    h <- textHostName
    return $ utmpUser u <> "@" <> h <> "/" <> resource u

newCon :: (Text -> IO ()) -> STM (Word8,Maybe UtmpRecord) -> TVar (Maybe UtmpRecord) -> IO ()
newCon log activeTTY utmp = do
    ((tty,tu),u) <- atomically $
        liftM2 (,) activeTTY
                   (readTVar utmp)
    flip (maybe $ return ()) u $ \u -> do
    jid <- ujid u
    log $ status (resource u) tty tu <> " " <> jid <> "  pid=" <> tshow (utmpPid u)
            <> (if istty (resource u)
                 then " host=" <> tshow (utmpHost u)
                 else "")
            <> " session=" <> tshow (utmpSession u)
            <> " addr=" <> tshow (utmpRemoteAddr u)
    loop tty tu (Just u)
 where
    bstatus r ttynum mtu
        =  r == ttystr
        || match mtu
        where ttystr = "tty" <> tshow ttynum
              searchstr mtu = maybe ttystr utmpHost $ do
                tu <- mtu
                guard (not $ Text.null $ utmpHost tu)
                return tu
              match mtu = searchstr mtu `Text.isInfixOf` Text.dropWhile (/=':') r
    status r ttynum tu =
        if bstatus r ttynum tu
            then "Available"
            else "Away     "

    istty r = fst3 == "tty" && Text.all isDigit rst
     where
        (fst3,rst) = Text.splitAt 3 r

    loop tty tu u = do
        what <- atomically $ foldr1 orElse
            [ do (tty',tu') <- retryWhen activeTTY
                            (\ttyu -> bstatus r tty tu == uncurry (bstatus r) ttyu)
                 return $ ttyChanged tty' tu'
            , do u' <- retryWhen (readTVar utmp) (==u)
                 return $ utmpChanged u'
            ]
        what
      where
        r = maybe "" resource u

        ttyChanged tty' tu' = do
            jid <- maybe (return "") ujid u
            log $ status r tty' tu' <> " " <> jid
            loop tty' tu' u

        utmpChanged u' = maybe dead changed u'
         where
            changed u' = do
                jid0 <- maybe (return "") ujid u
                jid <- ujid u'
                log $ "changed: " <> jid0 <> " --> " <> jid
                loop tty tu (Just u')
            dead = do
                jid <- maybe (return "") ujid u
                log $ "Offline   " <> jid

logit outvar s = do
    atomically $ takeTMVar outvar
    Text.putStrLn s
    atomically $ putTMVar outvar ()


main = do
    outvar <- atomically $ newTMVar ()

    cs <- newConsoleState
    inotify <- initINotify

    -- get active tty
    mtty <- monitorTTY (onTTY outvar cs)
    atomically $ retryWhen (readTVar $ csActiveTTY cs) (==0)

    -- read utmp
    onLogin cs (newCon $ logit outvar) Modify

    -- monitor utmp
    wd <- addWatch
            inotify
            [Modify] -- [CloseWrite,Open,Close,Access,Modify,Move]
            utmp_file
            (onLogin cs (newCon $ logit outvar))

    waitForTermSignal

    putStrLn "goodbye."