{-# LANGUAGE CPP               #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies      #-}
{-# LANGUAGE ExistentialQuantification #-}
module Main where

import System.Directory
import Control.Monad
import System.Posix.Signals
import System.Posix.Types
import System.Posix.Process
import Data.Maybe
import Data.Char
import ConfigFiles

import System.INotify
#ifndef NOUTMP
import UTmp 
-- UTmp is inconvenient for the profiling build due to Template Haskell
-- causing ghc to report:
--
--     Dynamic linking required, but this is a non-standard build (eg. prof).
--         You need to build the program twice: once the normal way, and then
--             in the desired way using -osuf to set the object file suffix.
--
-- ... so -DNOUTMP was introduced to remove TemplateHaskell from the build. 
-- However, a full-featured profiling build can be built using the bp
-- script which makes a non-profiling binary available during the build in
-- the manner in which the error message attempted (and failed) to communicate.
#endif
import FGConsole
#ifdef HAXML
import XMPPServer
#else
import XMPP
#endif
import ControlMaybe
import Data.HList
import Control.Exception hiding (catch)
import LocalPeerCred
import System.Posix.User
import qualified Data.Set as Set
import           Data.Set as Set (Set,(\\))
import qualified Data.Map as Map
import           Data.Map as Map (Map)

import Control.Concurrent.STM
import Control.Concurrent
import Control.Monad.Trans.Maybe
import Control.Monad.IO.Class

import ByteStringOperators
import qualified Data.ByteString.Lazy.Char8 as L
import           Data.ByteString.Lazy.Char8 as L (ByteString,putStrLn)
import qualified Prelude
import           Prelude hiding (putStrLn)
import System.Environment
import qualified Text.Show.ByteString as L
import Network.Socket (Family(AF_INET,AF_INET6))
import           Holumbus.Data.MultiMap as MM (MultiMap)
import qualified Holumbus.Data.MultiMap as MM

type RefCount = Int

type JabberResource = L.ByteString
type JabberName = L.ByteString
data JabberUser = JabberUser JabberName Peer
  deriving (Eq,Ord,Show)

splitResource :: JID -> Maybe (JabberUser,Maybe JabberResource)
splitResource (JID Nothing  _ _  ) = Nothing
splitResource (JID (Just n) p r  ) = Just (JabberUser n p, r)

unsplitResource (JabberUser n p) r = JID (Just n) p r

data PresenceState = PresenceState 
    { hostname :: Peer -- ByteString, TODO: remove this, its always LocalHost now
    , currentTTY    :: TVar ByteString
    , activeUsers   :: TVar (Set (UserName, Tty, ProcessID))
    , subscriberMap :: TVar (Map JID (RefCount,TChan Presence)) -- UNUSED as yet
    , localSubscriber :: TMVar (RefCount,TChan Presence) -- TODO: rename this, its not just locals
                                                         -- ... or make a seperate channel for remotes
    , rosterChannel :: TMVar (RefCount,TChan RosterEvent)
    , remoteUsers :: TVar (Map Peer (RefCount,TVar (MultiMap JabberUser (JabberResource,JabberShow))))
    , outGoingConnections :: TVar (Map Peer (TChan OutBoundMessage, ThreadId)) 
    }


getJabberUserForId muid =
             maybe (return "nobody") 
                   (\(uid,_) -> 
                        handle (\(SomeException _) -> 
                                return . L.append "uid." . L.pack . show $ uid)
                          $ do
                             user <- fmap userName $ getUserEntryForID uid
                             return (L.pack user)
                        )
                   muid

data ClientSession = ClientSession {
    localhost :: Peer, -- ByteString,
    unix_uid :: (IORef (Maybe (UserID,L.ByteString))),
    unix_resource :: (IORef (Maybe L.ByteString)),
    chans :: TVar [RefCountedChan],
    presence_state :: PresenceState
}

addJid modify user jid = do
    newitem <- parseHostNameJID jid
    let cmp jid = do
            -- putStrLn $ "Comparing "<++>bshow jid
            olditem <- parseHostNameJID jid
            if olditem==newitem then return Nothing
                                else return $ Just jid
    modify user cmp (Just jid)

instance JabberClientSession ClientSession where
    data XMPPClass ClientSession = ClientSessions PresenceState

    newSession (ClientSessions state) sock = do
        -- muid <- getLocalPeerCred sock
        addr <- getPeerName sock
        muid <- getLocalPeerCred' addr
        L.putStrLn $ "CLIENT SESSION: open " <++> bshow muid
        uid_ref <- newIORef muid
        res_ref <- newIORef Nothing
        chans <- atomically $ newTVar []
        return $ ClientSession (hostname state) uid_ref res_ref chans state

    setResource  s  resource    = do
        -- TODO: handle resource = empty string 
        us <- fmap Set.toList . readTVarIO $ activeUsers (presence_state s)
        muid <- readIORef (unix_uid s)
        let tailOf3 (_,a,b) = (a,b)
        (t,pid) <- case muid of
                      Just (uid,inode) -> identifyTTY (map tailOf3 us) uid inode
                      Nothing          -> return (Nothing,Nothing)
        let rsc = t `mplus` fmap ( ("pid."<++>) . bshow ) pid
            rsc' = maybe resource id rsc
        writeIORef (unix_resource s) (Just rsc')
        L.putStrLn $ "CLIENT SESSION: resource " <++> rsc'

    getJID       s             = do
        let host = localhost s
        user <- readIORef (unix_uid s) >>= getJabberUserForId

        rsc <- readIORef (unix_resource s)
        -- let jid = user <++> "@" <++> host <++?> "/" <++$> rsc
        L.putStrLn $ "CLIENT SESSION: jid " <++> L.show (JID (Just user) host rsc)
        return (JID (Just user) host rsc)

    closeSession session  = do
        atomically $ do
            cs <- readTVar (chans session)
            forM_ cs $ \(RefCountedChan c) -> do
                unsubscribeToChan c
        L.putStrLn "CLIENT SESSION: close"

    subscribe session Nothing = do
        let tmvar = localSubscriber (presence_state session)
        atomically $ do
            cs <- readTVar (chans session)
            writeTVar (chans session) (RefCountedChan tmvar:cs)
            subscribeToChan tmvar
    subscribe session (Just jid) = do -- UNUSED as yet
        let tvar = subscriberMap (presence_state session)
        atomically $ subscribeToMap tvar jid

    subscribeToRoster session = do
        let rchan = rosterChannel . presence_state $ session
        atomically $ do
            cs <- readTVar (chans session)
            writeTVar (chans session) (RefCountedChan rchan:cs)
            subscribeToChan rchan
            

    forCachedPresence s action = do
        jid <- getJID s 
        L.putStrLn $ "forCachedPresence "<++> bshow jid
        withJust (name jid) $ \user -> do
        let parseHostNameJID' str = do
                handle (\(SomeException _) -> return Nothing)
                       (fmap Just . parseHostNameJID $ str)
        buddies <- do
            buddies <- ConfigFiles.getBuddies user
            fmap catMaybes (mapM parseHostNameJID' buddies)
        remotes <- readTVarIO . remoteUsers . presence_state $ s
        forM_ buddies $ \buddy -> do
            L.putStrLn $ "forCachedPresence buddy = "<++> bshow buddy
            let mjids = fmap snd $ Map.lookup (peer buddy) remotes
            jids <- maybe (return MM.empty) readTVarIO mjids
            L.putStrLn $ "forCachedPresence jids = "<++> bshow jids
            withJust (splitResource buddy) $ \(buddyU,_) -> do
            forM_ (Set.toList . MM.lookup buddyU $ jids) $ \(rsc,status) -> do
            let p = Presence buddy status
            L.putStrLn $ "cached presence: " <++> bshow p
            action p

    addSolicited s jid = do
        user <- readIORef (unix_uid s) >>= getJabberUserForId
        addJid ConfigFiles.modifySolicited user jid
        let rchan = rosterChannel . presence_state $ s
        atomically $ do
            isempty <- isEmptyTMVar rchan
            when (not isempty) $ do
            (_,ch) <- readTMVar rchan
            writeTChan ch (RequestedSubscription user jid)
            

    getMyBuddies s = do
        user <- readIORef (unix_uid s) >>= getJabberUserForId
        ConfigFiles.getBuddies user
    getMySubscribers s = do
        user <- readIORef (unix_uid s) >>= getJabberUserForId
        ConfigFiles.getSubscribers user
    getMyOthers s = do
        user <- readIORef (unix_uid s) >>= getJabberUserForId
        ConfigFiles.getOthers user
    getMyPending s = do
        user <- readIORef (unix_uid s) >>= getJabberUserForId
        ConfigFiles.getPending user
    getMySolicited s = do
        user <- readIORef (unix_uid s) >>= getJabberUserForId
        ConfigFiles.getSolicited user


tupleToJID (user,tty,pid) = jid user LocalHost tty

data PeerSession = PeerSession {
    announced :: TVar (MultiMap JabberUser (JabberResource,JabberShow)),
    peer_name :: Peer,
    peer_global :: PresenceState
}
instance JabberPeerSession PeerSession where
    data XMPPPeerClass PeerSession = PeerSessions PresenceState

    newPeerSession (PeerSessions state) sock = do
        me <- fmap (RemotePeer . withoutPort) (getPeerName sock)
        L.putStrLn $ "PEER SESSION: open "<++>showPeer me
        let remotes = remoteUsers state
        jids <- atomically $ getRefFromMap remotes me (newTVar MM.empty) return
        return $ PeerSession jids me state

    closePeerSession session  = do
        L.putStrLn $ "PEER SESSION: close "<++>showPeer (peer_name session)
        let offline jid = Presence jid Offline
        unrefFromMap (remoteUsers . peer_global $ session) (peer_name session)
            $ do
                js <- fmap (MM.toAscList) (readTVarIO . announced $ session)
                forM_ js $ \(u,rs) -> do
                forM_ (Set.toList rs) $ \(rsc,_) -> do
                announcePresence session . offline $ unsplitResource u (Just rsc)

    peerSessionFactory session = PeerSessions (peer_global session)

    peerAddress session = peer_name session

    userStatus session user = do
        let state = peer_global session
        (tty,users) <- atomically $ do
            tty <- readTVar $ currentTTY state
            users <- readTVar $ activeUsers state
            return (tty,users)
        let jids = Set.filter (\(name,tty,pid) -> name ==user) users
            ps = map (\jid -> Presence (tupleToJID jid) (matchResource' tty jid)) . Set.toList $ jids
        if null ps
            then return [Presence (JID { name=Just user, peer=LocalHost, resource=Nothing }) Offline]
            else return ps

    announcePresence session (Presence jid status) = do
        (greedy,subs) <- atomically $ do
            subs <- readTVar $ subscriberMap (peer_global session)
            greedy <- fmap snd $ readTMVar $ localSubscriber (peer_global session)
            return (greedy,subs)
        update_presence (Just greedy) (fmap snd subs) [jid] (const status)
        liftIO . atomically $ do
            jids <- readTVar . announced $ session
            withJust (splitResource jid) $ \(u,rsc) -> do
            let match (r',_) = (rsc==Nothing || Just r'==rsc)
            writeTVar (announced session) 
                $ case status of
                    Offline -> MM.deleteElemIf u match jids
                    stat    -> maybe jids (\r -> MM.insert u (r,stat) jids) rsc

    getBuddies     _ user = ConfigFiles.getBuddies user
    getSubscribers _ user = ConfigFiles.getSubscribers user


data RefCountedChan = forall a. RefCountedChan (TMVar (RefCount,TChan a))

subscribeToChan :: TMVar (RefCount, TChan a) -> STM (TChan a)
subscribeToChan tmvar =
            (do (cnt,chan) <- takeTMVar tmvar
                putTMVar tmvar (cnt+1,chan)
                chan' <- dupTChan chan
                return chan' )
          `orElse`
            (do chan <- newTChan
                putTMVar tmvar (1,chan)
                return chan )
unsubscribeToChan :: TMVar (RefCount,TChan a) -> STM ()
unsubscribeToChan tmvar = do
    isEmpty <- isEmptyTMVar tmvar
    when (not isEmpty) $ do
    (cnt,chan) <- takeTMVar tmvar
    when (cnt>1) (putTMVar tmvar (cnt-1,chan))

getRefFromMap tvar key newObject copyObject = do
    subs <- readTVar tvar
    let mbobject       = Map.lookup key subs
    (object,subs') <- 
        do case mbobject of
            Nothing -> do
              newobject <- newObject
              return (newobject, Map.insert key (1,newobject) subs)
            Just (cnt,object) -> do
              object' <- copyObject object
              return (object', Map.insert key (cnt+1,object) subs)
    writeTVar tvar subs'
    return object

unrefFromMap tvar key finalizer = do
    vanished <- atomically $ do
        omap <- readTVar tvar
        let (r,omap') = Map.updateLookupWithKey unref key omap
        writeTVar tvar omap'
        return (isNothing r)
    when vanished finalizer
 where
    unref key (cnt,object) =
        if cnt==1 then Nothing
                  else Just (cnt-1,object)
    

subscribeToMap tvar jid =
    getRefFromMap tvar jid newTChan dupTChan

matchResource tty jid = maybe Away (avail . (==tty)) $ resource jid
 where
    avail True  = Available
    avail False = Away

matchResource' tty (_,rsc,_) = avail  (rsc==tty)
 where
    avail True  = Available
    avail False = Away

sendPresence chan jid status = 
    (liftIO  . atomically . writeTChan chan . Presence jid $ status) :: MaybeT IO ()

lookupT jid subscribers = MaybeT . return $ Map.lookup jid subscribers 

update_presence locals_greedy subscribers state getStatus =
    forM_ state $ \jid -> do
      let status = getStatus jid 
      runMaybeT $ do
        chan <- lookupT jid subscribers
        sendPresence chan jid status
      runMaybeT $ do
        chan <- MaybeT . return $ locals_greedy
        sendPresence chan jid status
      putStrLn $ bshow jid <++> " " <++> bshow status

newPresenceState hostname = atomically $ do
    tty  <- newTVar ""
    us   <- newTVar (Set.empty)
    subs <- newTVar (Map.empty)
    locals_greedy <- newEmptyTMVar
    rchan <- newEmptyTMVar
    remotes <- newTVar (Map.empty)
    server_connections <- newServerConnections
    return $ PresenceState hostname tty us subs locals_greedy rchan remotes server_connections

sendProbes state jid = do
    L.putStrLn $ "sending probes for " <++> bshow jid
    withJust (name jid) $ \user -> do
    let parseHostNameJID' str = do
            handle (\(SomeException _) -> return Nothing)
                   (fmap Just . parseHostNameJID $ str)
    buddies <- do
        buddies <- ConfigFiles.getBuddies user
        fmap catMaybes (mapM parseHostNameJID' buddies)
    L.putStrLn $ "buddies for " <++> bshow jid <++> " = " <++> bshow buddies
    remotes <- readTVarIO (remoteUsers state)
    forM_ buddies $ \buddy -> do
        let mjids = fmap snd $ Map.lookup (peer buddy) remotes
        jids <- maybe (return MM.empty) readTVarIO mjids
        withJust (splitResource buddy) $ \(buddyU,_) -> do
        let noinfo = not (MM.member buddyU jids)
        when noinfo $ do
        L.putStrLn $ "sendMessage " <++> bshow (PresenceProbe jid buddy)
        sendMessage (outGoingConnections state) (PresenceProbe jid buddy) (peer buddy) 
        return ()

track_login host state e = do
#ifndef NOUTMP
    us <- UTmp.users
#else
    let us = []
#endif
    let toJabberId host (user,tty,pid) = 
            if L.take 3 tty == "tty"
              then Just (user,tty,pid) -- (jid user host tty)
              else Nothing
        new_users = Set.fromList $ mapMaybe (toJabberId host) us
    (tty,known_users,subs,locals_greedy) <- atomically $ do
            tty <- readTVar                $ currentTTY state
            st  <- flip swapTVar new_users $ activeUsers state
            xs  <- readTVar                $ subscriberMap state
            locals_greedy <- tryReadTMVar $ localSubscriber state
            return (tty,st,fmap snd xs,fmap snd locals_greedy)
    let arrivals   = map tupleToJID . Set.toList $ new_users \\ known_users
        departures = map tupleToJID . Set.toList $ known_users \\ new_users
    update_presence locals_greedy subs departures $ const Offline
    update_presence locals_greedy subs arrivals   $ matchResource tty
    forM_ arrivals
          $ sendProbes state

on_chvt state vtnum = do
    let tty = L.snoc "tty" $ intToDigit (fromIntegral vtnum)
    L.putStrLn $ "VT switch: " <++> tty
    (users,subs,locals_greedy) <- atomically $ do
        us   <- readTVar $ activeUsers state
        subs <- readTVar $ subscriberMap state
        writeTVar (currentTTY state) tty
        locals_greedy <- tryReadTMVar $ localSubscriber state
        return (us,fmap snd subs,fmap snd locals_greedy)
    update_presence locals_greedy 
                    subs 
                    (map tupleToJID . Set.toList $ users) 
                    $ matchResource tty

start :: Network.Socket.Family -> IO ()
start ip4or6 = do
    let host = LocalHost
    global_state <- newPresenceState host
    let dologin e = track_login host global_state e
        dologin :: t -> IO ()

    chan <- atomically $ subscribeToChan (localSubscriber global_state)
    remotes <- forkIO $ seekRemotePeers (PeerSessions global_state) chan (outGoingConnections global_state)

    installHandler sigUSR1 (Catch (dologin (userError "signaled"))) Nothing
    -- installHandler sigTERM (CatchOnce (dologin (userError "term signaled"))) Nothing
    mtty <- monitorTTY (on_chvt global_state)
    inotify <- initINotify
#ifndef NOUTMP
    wd <- addWatch
            inotify
            [CloseWrite] -- [Open,Close,Access,Modify,Move]
            utmp_file
            dologin
#endif
    clients <- listenForXmppClients ip4or6 (ClientSessions global_state) 5222 HNil
    peers   <- listenForRemotePeers ip4or6 (PeerSessions global_state) 5269 HNil

    threadDelay 1000 -- wait a moment to obtain current tty
    dologin ()
    putStrLn "\nHit enter to terminate...\n"
    getLine
    killThread remotes
    quitListening clients
    quitListening peers
    -- threadDelay 1000
    putStrLn "closed listener."
    unmonitorTTY mtty
    putStrLn "unhooked tty monitor."
#ifndef NOUTMP
    removeWatch wd
#endif
    putStrLn "Normal termination."

sendUSR1 pid = do
    signalProcess sigUSR1 pid

getStartupAction []     = throw (userError "pid file?") >> return (Right "")
getStartupAction (p:ps) = do
    handle onEr $
        ( do
            pid <- fmap CPid (readFile p >>= readIO)
            -- signal pid
            return (Left pid) )
 where
    onEr (SomeException _) = do
        pid <- getProcessID
        putStrLn $ "starting pid = " <++> bshow pid
        handle (\(SomeException _) -> getStartupAction ps)
            (do 
                writeFile p (show pid)
                putStrLn $ "writing " <++> bshow p
                -- start daemon
                return (Right p) )

runOnce ps run notify = getStartupAction ps >>= doit
 where
    doit (Left pid      ) = notify pid
    doit (Right pidfile ) = do 
        run
        removeFile pidfile

getOptions [] = Map.empty

getOptions (('-':opt_name):xs) =
    if xs/=[] && xs!!0!!0/='-'
      then Map.insert (L.pack opt_name) (L.pack (xs!!0)) (getOptions (tail xs))
      else Map.insert (L.pack opt_name) "" (getOptions xs)

getOptions (x0:xs) = getOptions xs


main = do
    opts <- fmap getOptions getArgs
    let use_ip4 = if isJust (Map.lookup "4" opts) then AF_INET else AF_INET6
    runOnce ["/var/run/presence.pid","/tmp/presence.pid"] (start use_ip4) sendUSR1