-- IRC connection maintainer

-- The point of this program is to maintain a connection to an IRC server.
-- When the connection to the IRC server is lost, it should be restarted.
-- It should accept one connection from one client at a time, and to that
-- client all messages from the IRC server should be forwarded. Likewise,
-- all messages from the client should be forwarded unchanged to the IRC
-- server. The client should respond to PING.

module Main where

import System
-- We are working with sockets via Handles
import System.IO                 (Handle,hSetBuffering,hGetContents,hClose
                                 ,BufferMode(LineBuffering),hPutStrLn)
-- Need to ignore some signals
import System.Posix.Signals      (installHandler, sigPIPE, Handler(Ignore))
-- We're using threads
import Control.Concurrent        (forkIO,threadDelay,ThreadId,killThread)
-- We need to store Handles in MVars for access between threads
import Control.Concurrent.MVar   (newMVar,MVar,swapMVar,readMVar,modifyMVar_)
-- We need a Chan for outputting messages atomically to stdout
import Control.Concurrent.Chan   (newChan,Chan,writeChan,readChan)
-- The MVar Handles will be passed in the state monad
import Control.Monad.State       (evalStateT,StateT,get,gets)
-- Working with sockets
import Network                   (withSocketsDo,connectTo,PortNumber,sClose
                                 ,PortID(PortNumber),accept,sClose,listenOn
                                 ,Socket)
-- Nice abstractions
import Control.Arrow             ((>>>),(&&&))
-- Monad transformer functions
import Control.Monad.Trans       (liftIO)
-- Monad funky stuff
import Control.Monad             (liftM2,forever,foldM_,when)
-- Applicative abstractions
import Control.Applicative       ((<$>))
-- Matching IRC messages
import Text.Regex                (mkRegex,matchRegex)
-- Maybe abstractions
import Data.Maybe                (isJust,fromJust)
-- List stuff
import Data.List                 (isPrefixOf)

-- TODO: make this not crap
main :: IO ()
main = do
  args <- getArgs
  case args of
    [host,port,listen] -> go (Settings "" "" host (fromInteger (read port)) (fromInteger (read listen)) False)
    [nick,user,host,port,listen] ->
        go (Settings nick 
                     user
                     host 
                     (fromInteger (read port))
                     (fromInteger (read listen))
                     True)
    _ -> error "Expected: [<nick> <user>] <IRCd address> <IRCd port> <bouncer port>"
  return ()

go :: Settings -> IO ()
go settings = 
  -- Need to initialise sockets subsystem.
  withSocketsDo $ do
  -- Writing data to closed connections will not end the program.
  installHandler sigPIPE Ignore Nothing
  -- Create some MVars to hold Handles for the client and the server.
  clientHandle <- newMVar Nothing
  clients <- newMVar []
  serverHandle <- newMVar Nothing
  -- We don't want messages becoming mixed together, so we'll have a
  -- channel to write to and the main thread will read from it and print
  -- the messages to stdout.
  outChan <- newChan
  -- Simple utility to start the threads with the state.
  let state = State clientHandle clients serverHandle outChan settings
      start = forkIO . flip evalStateT state
  -- Start the connection to the IRC server and maintain it.
  start startCon
  -- Start listening for a connection.
  start listenClients
  -- Print all messages to stdout.
  forever $ readChan outChan >>= putStrLn

-- Listening for clients.
listenClients :: Maintainer ()
listenClients = do
  -- Get the port to listen on.
  port <- gets $ stateSettings >>> settingsListen
  -- Start listening for a client.
  socket <- io $ listenOn (PortNumber port)
  echo "Listening for client connections..."
  run <- flip evalStateT <$> get
  io $ forever $ run $ accepts socket

accepts :: Socket -> Maintainer ()
accepts socket = do
  (h,host,_) <- io $ accept socket -- This blocks.
  run <- flip evalStateT <$> get
  io $ forkIO $ run $ do
    -- IRC sends messages by line.
    io $ hSetBuffering h LineBuffering
    echo $ "Client " ++ host ++ " connected."
    -- Make an MVar for the handle
    hv <- io $ newMVar (Just h)
    -- Add it to the clients list
    clients <- gets stateClients
    io $ modifyMVar_ clients (return . ((h,hv):))
    -- If any errors occur, we ought to delete the client from the 
    -- client list.
    handleErrors (deleteClient h) $ do
      -- We can now recieve messages from the server
      -- Start proxying all messages indiscriminately to IRC server.
      lines <- lines <$> (io $ hGetContents h)
      -- Start timeout thread. The connection is closed if the client fails
      -- to respond to PING in 20 seconds.
      v <- io $ newMVar True
      timeoutThread v h
      server <- gets stateServerH
      mapM_ (echo <> (checkPing v h $ proxyTo server)) lines
      -- When the above line completes, that means the connection is closed,
      echo $ "Client " ++ host ++ " disconnected."
      return ()
  return ()

-- Check to see if a message is a PING reply. If so we ought to capture
-- it. Clients only recieve pings from the bouncer, never the server.
checkPing :: MVar Bool -> Handle -> (String -> Maintainer ()) -> String 
          -> Maintainer ()
checkPing v h f msg | prefix "PONG" msg = io $ setMVar v True
                    | otherwise = do io $ setMVar v True
                                     f msg
                    where prefix = flip isPrefixOf

-- Simply delete a client from the client list according to the handle.
-- The handle is unique, so it is a good identifier.
deleteClient handle = do
    clients <- gets stateClients
    io $ modifyMVar_ clients (return . (filter ((==handle) . fst)))

timeoutThread :: MVar Bool -> Handle -> Maintainer ThreadId
timeoutThread v h = do
  run <- flip evalStateT <$> get
  io $ forkIO $ run $ handleErrors (return ()) (io pingClose) where
    pingClose = do  -- Assume no data has been recieved.
      setMVar v False
      wait 120 -- Wait 20 seconds.
      -- Check; if no data is recieved, send a PING.
      check $ do hPutStrLn h "PING :bouncer"
                 wait 60 -- Wait another 10 secs.
                 -- No reply, close the connection.
                 check $ do hClose h
        where -- Checks to see if any data has been received.
              check d = do anyData <- readMVar v 
                           if anyData
                              -- Data has been receieved, restart the timeout.
                              then pingClose 
                              -- Data has not been received, do something.
                              else d

-- Connect to the IRC server, send necessary information, and reconnect
-- if disconnected.
startCon :: Maintainer ()
startCon = handleErrors
    (do io $ wait 10; startCon) -- Wait 10 seconds and try again.
    (do
      -- Get connection settings.
      host <- gets $ stateSettings >>> settingsHost
      port <- gets $ stateSettings >>> settingsPort
      -- Connect to IRC server...
      h <- io $ connectTo host (PortNumber port)
      io $ hSetBuffering h LineBuffering -- IRC sends messages by line.
      register <- gets $ stateSettings >>> settingsRegister
      when register $ sendRegistration h
      -- Update the MVar containing the server's Handle
      serverH <- gets stateServerH; io $ swapMVar serverH (Just h)
      -- Start proxying all messages to a client (if there is one)
      lines <- lines <$> (io $ hGetContents h)
      mapM_ (echo <> (handleIRC h $ proxyToAll)) lines
      -- When the above line completes, that means the connection is closed,
      -- so we will need to clear the Handle...
      io $ swapMVar serverH Nothing
      -- and re-connect.
      startCon
    )

-- Checks for messages that must be replied to, like PING.
handleIRC :: Handle -> (String -> Maintainer ()) -> String -> Maintainer ()
handleIRC h f msg | isJust ping  = do irc pong; echo pong
                  | otherwise   = f msg -- Forward it to another handler.
    where ping = do [ps] <- matchRegex (mkRegex "^PING (.*)") msg; return ps
          pong = "PONG " ++ fromJust ping
          irc = io . hPutStrLn h

-- Send basic IRC registration information.
sendRegistration :: Handle -> Maintainer ()
sendRegistration h = do
  nick <- gets $ stateSettings >>> settingsNick
  user <- gets $ stateSettings >>> settingsUser
  -- Any errors, abort the whole function.
  handleErrors (return ()) $ do
    irc $ "NICK " ++ nick
    irc $ "USER " ++ user ++ " 0 * :" ++ user
        where irc = io . hPutStrLn h

-- Forward a message to all clients.
proxyToAll :: String -> Maintainer ()
proxyToAll line = do
  clients <- gets stateClients
  clients' <- io $ readMVar clients
  mapM (flip proxyTo line . snd) clients'
  return ()

-- Forward a message to a client.
proxyTo :: SHandle -> String -> Maintainer ()
proxyTo sh line = do
  -- Read whatever value is in the MVar.
  h <- io $ readMVar sh
  case h of
    Nothing -> return () -- No handle, discard the line.
    -- If there is a Handle in there, the connection is *probably* up,
    -- and just in case it is not up, we catch any exceptions...
    -- We simply write the line to the Handle.
    Just h' -> io $ catch (hPutStrLn h' line) (const $ return ())

handleErrors :: Maintainer () -> Maintainer () -> Maintainer ()
handleErrors restart thunk = do 
  run <- flip evalStateT <$> get
  -- Catch any crazy socket errors.
  io $ catch (run thunk) (const $ run restart)

-- Writes a message to the output channel, to be displayed in stdout.
echo :: String -> Maintainer ()
echo msg = do ch <- gets stateOutput
              io $ writeChan ch msg

-- The State type will hold the MVar Handles and a Chan for sending
-- messages to be printed to stdout.
type Maintainer a = StateT State IO a

{-
data State = State
    { stateClientH  :: SHandle
    , stateServerH  :: SHandle
    , stateOutput   :: Chan String 
    , stateSettings :: Settings }
-}

data State = State
    { stateClientH  :: SHandle
    , stateClients  :: MVar [(Handle,SHandle)]
    , stateServerH  :: SHandle
    , stateOutput   :: Chan String 
    , stateSettings :: Settings }

-- Handle that can be shared
type SHandle = MVar (Maybe Handle)

-- This program ought to send necessary messages to the IRC server such
-- as NICK, USER and PONG, in case no client is connected.

data Settings = Settings
    { settingsNick     :: String
    , settingsUser     :: String
    , settingsHost     :: String
    , settingsPort     :: PortNumber
    , settingsListen   :: PortNumber 
    , settingsRegister :: Bool }
                deriving Show

defSettings = Settings "lojbot" "lojbot" "irc.freenode.net" 6667 6667 True

-- Utility
io = liftIO

-- Nice monad utility
(<>) = liftM2 (>>)

-- Waiting utility
wait = threadDelay . (*1000000)

-- MVar utility
setMVar v val = modifyMVar_ v (return . const val) >> return ()