{-#OPTIONS -cpp #-}
{-#LANGUAGE OverloadedStrings,MultiParamTypeClasses,DeriveGeneric,BangPatterns #-}
{-|
Module      :  Parry.Server
Copyright   :  (c) Pierre-Étienne Meunier 2014
License     :  GPL-3
Maintainer  :  pierre-etienne.meunier@lif.univ-mrs.fr
Stability   :  experimental
Portability :  All

Tools to build synchronization servers. For instance, to write a simple
server with just a web interface on port 8000, you would use:

> import Control.Concurrent
> import Parry.Server
> import Parry.WebUI
>
> main::IO ()
> main=do
>   state<-initState initial
>   _<-forkIO $ webUI 8000 state
>   server (defaultConfig public) state
-}

module Parry.Server (
  -- * Jobs on the server side
  Exhaustive(..),
  Result(..),
  -- * Server's internal state
  initState,
  stateFromFile,
  saveThread,
  State(..),
  -- * Server configuration and functions
  Config(..),defaultConfig,server
  ) where
import Control.Concurrent
import Control.Exception as E
import Control.Monad
import Control.Concurrent.MSem as Sem
import Network
import System.IO
import System.Directory
import Data.List
import Data.Time.Format()
import Data.Time.Clock.POSIX
#ifdef UNIX
import System.Posix.Signals
#endif

import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy.Char8 as LB
import qualified Data.Map as M
import qualified Data.Set as S
import GHC.Generics
import Data.Binary

import Codec.Crypto.RSA.Pure

import Parry.Protocol
import Parry.Util

-- | The class of jobs and job results that Parry can deal with. For
-- efficiency and to keep types simple, jobs and results are stored in
-- a single type.
class Exhaustive j where
  -- | Indication of the depth of a job in the explored tree. The server sends
  -- the least deep jobs first, as an optimization of network use.
  depth :: j->Int
  -- | Number of times a job has been killed. When a job is killed,
  -- either because it must be reshared, or because the client itself
  -- was killed, it is scheduled to be re-executed by the server.
  killed::j->Int
  -- | Called each time a job needs to be killed. For better resharing,
  -- this function must verify @killed (kill j) >= killed j@.
  kill::j->j

-- | The class of results, and how to combine them in the server state.
class Result j r where
  -- | A function to tell how to combine job results. That function will be
  -- called on the hostname of the reporting client, with the finished job it
  -- sent, and the current result from the server state.
  addResult::HostName->r->j->r

-- | This type is exposed mostly for writing alternative user interfaces.
-- Other operations must be done using the functions in this module, or
-- the correction of the protocol can be lost.
data State j r=State {
  -- | Available jobs
  jobs::S.Set (Int,j),
  -- | Map from the machine id to its hostname, its current job, its
  -- starting time, the last time we heard from it.
  ongoing::M.Map Integer (HostName,PublicKey,j,Double,Double),
  -- | Set of unemployed machines
  unemployed::S.Set Integer,
  -- | The results.
  results::r,
  -- | The smallest available machine id. In a run of the server, it
  -- is guaranteed that are never assigned the same.
  newId::Integer,
  -- | Total number of jobs killed from the beginning (for benchmarking purposes).
  killings::Int,
  -- | Number of jobs finished (for benchmarking purposes).
  solved::Integer,
  -- | The list of authorized RSA public keys.
  authorizedKeys::[PublicKey]
  } deriving (Show,Read,Generic)

instance (Binary j,Binary r)=>Binary (State j r)

answer::(Exhaustive j,Result j r,Eq j,Ord j,Binary j)=>
        Double->String->State j r->ClientMessage j
        ->(State j r,ServerMessage j)
answer t host st (GetJob num key)=
  case M.lookup num (ongoing st) of
    Just (ho,key0,j0,t0,_)->

      if ho==host && key0==key then
        (st { ongoing=M.insert num (ho,key,j0,t0,t) (ongoing st) },
         Job (not $ S.null $ unemployed st) j0)
      else
        (st,Die)
    Nothing->

      if S.null (jobs st) then
        if M.null (ongoing st) then
          (st,Finished)
        else
          (st { unemployed=S.insert num (unemployed st) },Die)

      else
        let ((_,h),nextJobs)=S.deleteFindMin (jobs st)
            shareIt=killed h>0
        in
         (st { jobs=nextJobs,
               unemployed=S.delete num (unemployed st),
               ongoing=M.insert num (host,key,h,t,t) (ongoing st) },
          Job shareIt h)

answer t host st j@(NewJobs {})=
  case M.lookup (clientId j) (ongoing st) of
    Nothing->(st,Die)
    Just (ho,key,j0,t0,_)->
      if host==ho && j0==currentJob j then
        (st { jobs=foldl' (\s x->S.insert (depth x,x) s) (jobs st) (newJobs j),
              ongoing=M.insert (clientId j)
                      (host,key,nextJob j,t0,t)
                      (ongoing st),
              results=foldl' (addResult host) (results st) (jobResults j) }, Ack)
      else
        (st, Die)

answer _ _ st j@(JobDone {})=
  case M.lookup (clientId j) (ongoing st) of
    Nothing->(st,Die);
    Just (host,_,j0,_,_)->
      if j0==currentJob j then
        (st { ongoing=M.delete (clientId j) (ongoing st),
              results=foldl' (addResult host) (results st) (jobResults j),
              solved=solved st+1 }, Ack)
      else
        (st,Die)

answer t host st (Alive num)=
  case M.lookup num (ongoing st) of
    Nothing->(st,Die);
    Just (ho,key,j,t0,_)->
      if ho==host && (S.null (unemployed st) || (not $ S.null $ jobs st)) then
        (st { ongoing=M.insert num (ho,key,j,t0,t) (ongoing st) },Ack)
      else
        (st,Die)

reply::(Binary j,Exhaustive j,Result j r,Ord j)=>
       MVar (State j r) -> Handle -> HostName -> IO ()
reply state rhandle host=(do
  l<-B.hGetLine rhandle
  if l==B.pack "Hello" then
    modifyMVar_ state $ \st->do
      LB.hPutStrLn rhandle $ encode16l $ encode $ newId st
      return $ st { newId=newId st+1 }

    else do
    st<-withMVar state return
    sig<-B.hGetLine rhandle
    let dec=LB.fromStrict $ decode16 l
        msg=decode dec
        num=case msg of
          GetJob x _->x
          JobDone x _ _->x
          NewJobs x _ _ _ _->x
          Alive x->x
        key=case msg of
          GetJob _ pub->
            if any (==pub) (authorizedKeys st) then
              Just pub
            else
              Nothing
          _->(case M.lookup num $ ongoing st of
                 Just (_,pub,_,_,_)->Just pub
                 Nothing->Nothing)

    message<-case key of
      Nothing->return Die
      Just pub->
        case verify pub dec (LB.fromStrict $ decode16 sig) of
          Right True->do
            t<-getPOSIXTime
            modifyMVar state $ \st0->
              let (!a,!b)=answer (realToFrac t) host st0 msg in
              return (a,b)
          _->return Die
    LB.hPutStrLn rhandle (encode16l $ encode message)
  )

cleanupThread::(Ord j,Exhaustive j)=>MVar (State j r)->IO ()
cleanupThread state=do {
  t_<-getPOSIXTime;
  let { t=realToFrac t_ };
  modifyMVar_ state $ \st0->do {
    -- Find machines that have not given news for more than 10 minutes.
    let { (a,b)=M.partition (\(_,_,_,_,t1)->(t-t1) > 600) (ongoing st0);
          st=st0 { jobs=
                      M.foldl' (\set (_,_,job,t0,_)->
                                 S.insert (depth job,
                                           if t-t0 > 3600 then
                                             kill job
                                           else job) set)
                      (jobs st0)
                      a,
                   ongoing=b }
        };
    return st
    };
  -- Sleep 30 seconds, and clean again.
  threadDelay 30000000;
  cleanupThread state
  }

-- | Starts the synchronization server.
server::(Ord j, Binary j, Exhaustive j, Result j r)=>
        Config->MVar (State j r)->IO ()
server config state=withSocketsDo $ do {
#ifdef UNIX
  installHandler sigPIPE Ignore Nothing;
#endif
  threads<-Sem.new $ maxThreads config;
  _<-forkIO $ cleanupThread state;
  forever $ do {
    E.catch (bracket (listenOn (port config)) sClose $
             \sock->forever $ do {
               bracket (do { (s,a,_)<-accept sock; wait threads; return (s,a) })
               (\(s,_)->do { signal threads; hClose s})
               (\(s,a)->reply state s a)
               })
    (\e->let _=e::SomeException in appendFile (logFile config) (show e++"\n"));
    threadDelay 100000;
    };
  }

-- | Server configuration
data Config=Config {
  -- | The network port the synchronization server will listen on.
  port::PortID,
  -- | The maximal number of simultaneous threads that can be launched.
  maxThreads::Int,
  -- | Log file
  logFile::FilePath
  }

-- | Default server configuration, matching the client. Note that you
-- must provide your own public key for signing the messages.

defaultConfig::Config
defaultConfig=Config { port=PortNumber 5129, maxThreads=20, logFile="parry.err" }

-- | Creates a valid server state from an initial job.
initState::(Exhaustive j,Ord j,Result j r)=>[j]->r->IO (MVar (State j r))
initState initial r0=
  newMVar $ State { jobs=foldl (\s j->S.insert (depth j,j) s) S.empty initial,
                    ongoing=M.empty,
                    results=r0,
                    unemployed=S.empty,
                    newId=0,
                    killings=0,solved=0,
                    authorizedKeys=[] };

-- | Reads initial state from a file, or calls 'initState' if the file does
-- not exist.
stateFromFile::(Binary r,Binary j,Exhaustive j, Result j r,Ord j)=>FilePath->[j]->r->IO (MVar (State j r))
stateFromFile f initial r0=do {
  e<-doesFileExist f;
  state<-if e then do {
    st<-decodeFile f;
    newMVar st
    } else initState initial r0;
  return state
  }

-- | Saves state to the given file with the given periodicity, in
-- microseconds. This function does not return, so calling it inside a
-- 'Control.Concurrent.forkIO' is probably the best thing to do.
saveThread::(Binary r,Binary j)=>FilePath->Int->MVar (State j r)->IO ()
saveThread f del state=
  let { save=do {
           e<-doesFileExist f;
           if e then renameFile f (f++".last") else return ();
           withMVar state $ \st->LB.writeFile f $ encode st;
           threadDelay del;
           save
           }}
  in save