import Control.Monad
import Control.Exception
import Control.Concurrent
import Control.Concurrent.STM
import Network
import System.IO
import System.Time
import System.Locale (defaultTimeLocale)
import System.Process
import Data.Typeable
import Data.List
import Data.Char

import Cmdline


-- encapsulates a self-managing connection
data Con = Con (IO (Handle, Handle, Handle))
               (TMVar Handle) (TVar Bool)  -- outgoing (aka "stdin")
               (TMVar Handle) (TVar Bool)  -- incoming (aka "stdout")
               (TMVar Handle) (TVar Bool)  -- aux-in   (aka "stderr")
               (TVar Bool)  -- user-flag: connection enabled?
               (TVar Bool)  -- user-flag: auto-reconnect?

mkcon :: IO (Handle,Handle,Handle) -> IO Con
mkcon mk = do ti  <- atomically newEmptyTMVar
              to  <- atomically newEmptyTMVar
              te  <- atomically newEmptyTMVar
              tic <- atomically (newTVar False)
              toc <- atomically (newTVar False)
              tec <- atomically (newTVar False)
              tc  <- atomically (newTVar True)
              ta  <- atomically (newTVar True)
              let con = Con mk ti tic to toc te tec tc ta
              forkIO (con_manage con)
              return con

con_connect mk = loop 1
  where
  loop d = handle (\_-> retry)
                  (do (i,o,e) <- mk
                      ok <- hIsOpen i
                      if ok then return (i,o,e)
                            else retry)
    where
    d' = min 512 (2*d)
    retry = do threadDelay (d*1000000)
               loop d'

con_manage con@(Con mk ti tic to toc te tec tc ta)
 = do -- precondition: connection closed, handle TMVars empty
      -- wait until we should (re-)connect
      atomically $ do c <- readTVar tc
                      unless c retry
      (i,o,e) <- con_connect mk
      atomically $ do putTMVar ti i
                      putTMVar to o
                      putTMVar te e
                      writeTVar tic True
                      writeTVar toc True
                      writeTVar tec True
      -- wait until con. becomes disabled or fails
      atomically $ do c  <- readTVar tc
                      ic <- readTVar tic
                      oc <- readTVar toc
                      ec <- readTVar tec
                      when (c && ic && oc && ec) retry
                      -- disable connection if reconnect is off
                      a  <- readTVar ta
                      writeTVar tc (c&&a)
      hClose i
      hClose o
      hClose e
      -- throw out the old handles
      atomically $ do takeTMVar ti
                      takeTMVar to
                      takeTMVar te
                      writeTVar tic False
                      writeTVar toc False
                      writeTVar tec False
      -- wait at least one second before reconnecting
      threadDelay 1000000
      con_manage con

recv, recvaux :: Con -> IO String
recv    = con_access (\_ o _ -> o) hGetLine
recvaux = con_access (\_ _ e -> e) hGetLine

send :: Con -> String -> IO ()
send c s = con_access (\i _ _ -> i)
                      (\h -> hPutStrLn h s >> hFlush h)
                      c

con_access sel op con@(Con _ ti tic to toc te tec _ _)
 = do h <- atomically $ do c <- readTVar tv
                           unless c retry
                           takeTMVar tmv
      handleJust ioErrors
        (\_-> do -- mark this handle as closed
                 atomically $ do writeTVar tv False
                                 putTMVar tmv h
                 -- retry the operation
                 con_access sel op con)
        (do x <- op h
            atomically (putTMVar tmv h)
            return x)
  where
  (tmv,tv) = sel (ti,tic) (to,toc) (te,tec)

open  (Con _ _ tic _ _ _ _ tc _) = do atomically (writeTVar tc True)
                                      -- block until the handles are open
				      atomically $ do open <- readTVar tic
				                      unless open retry
close (Con _ _ tic _ _ _ _ tc _) = do atomically (writeTVar tc False)
                                      -- block until the handles are closed
				      atomically $ do open <- readTVar tic
				                      when open retry
auto  (Con _ _ _ _ _ _ _ _ ta) b = atomically (writeTVar ta b)
enab  (Con _ _ _ _ _ _ _ tc _) b = atomically (writeTVar tc b)


opts = [ parameter ["s", "status-prefix"]
                   "prepend STR to connection status messages (empty=disabled)"
                   "%"
       , parameter ["i", "incoming-prefix"]
                   "prepend STR to server lines"
                   ""
       , parameter ["c", "command-prefix"]
                   "take client lines prefixed by STR as control commands (empty=disabled)"
                   "/"
       , parameter ["o", "outgoing-prefix"]
                   "take client lines prefixed by STR as outgoing messages"
                   ""
       , parameter ["l", "logfile"]
                   "write log messages to FILE"
                   "ircmeister.log"
       ]

main = do  Args param [hst,prt',cmd]
             <- stdargs name version cpr desc opts argnames
           let prt = read prt' :: Int
               p   = param :: Typeable a => String -> a
           
           lgc  <- atomically newTChan
           forkIO (logger p lgc)
           let log n s = atomically (writeTChan lgc (n,s))
           log 0 (name++" v"++version++" started.")
           
           ctrl <- atomically newTChan
           clt <- mkcon (start p log cmd)
           srv <- mkcon (connect p log clt hst prt)
           
           forkIO (sinhandler p log clt srv)
           forkIO (cinhandler p log ctrl clt srv)
           forkIO (cerrhandler p log clt)           
           forkIO (console p log ctrl)
           
           control p log ctrl clt srv
           log 0 "exiting ircmeister."
    where
    name      = "ircmeister"
    version   = "1.0"
    cpr       = "by pesco, 2008"
    desc      = "connect client to server! :P"
    argnames  = ["server", "port", "client"]


control p log ctrl clt srv = loop
  where
  loop = do l <- atomically (readTChan ctrl)
            let cmd'       = dropWhile isSpace l
                (cmd,arg') = break isSpace cmd'
                arg        = dropWhile isSpace arg'
            unless (cmd=="die") $
              do handle cmd arg
                 loop
  
  handle :: String -> String -> IO ()
  handle "recon"   "on"  = auto srv True
  handle "recon"   "off" = auto srv False
  handle "recon"   "now" = close srv >> open srv
  handle "restart" "on"  = auto clt True
  handle "restart" "off" = auto clt False
  handle "restart" "now" = do status "restarting"
                              close clt
                              open clt
                              status "restarted"
  handle "connect" ""    = open srv
  handle "start" ""      = open clt
  handle "disconnect" "" = close srv
  handle "stop" ""       = close clt
  handle "<"       msg   = do log 2 ("< "++msg)
                              send srv msg
  handle ">"       msg   = do log 2 ("> "++msg)
                              send clt msg
  handle cmd arg         = log 0 ("unrecognized command: "++cmd++" "++arg)
  
  status
    | not (null spre) = \s -> send clt (spre++s)
    | otherwise       = \_ -> return ()
  spre = p "status-prefix"


console p log ctrl = loop
  where
  loop = do l <- getLine
            unless (null l) $
              do log 2 ("console cmd: "++l)
                 atomically (writeTChan ctrl l)
            loop

sinhandler p log clt srv = loop
  where
  loop = do l <- recv srv
            log 2 ("> "++l)
            send clt (ipre++l)
            loop
  ipre = p "incoming-prefix"

cinhandler p log ctrl clt srv = loop
  where
  loop = do l <- recv clt
            handle l
            loop
  handle l
    | opre `isPrefixOf` l = do let msg = drop no l
                               log 2 ("< "++msg)
                               send srv msg
    | cpre `isPrefixOf` l = do let cmd = drop nc l
                               log 2 ("client cmd: "++cmd)
                               atomically (writeTChan ctrl cmd)
    | otherwise           = log 2 ("unrecognized client line: "++l)  
  opre = p "outgoing-prefix"
  no = length opre
  cpre = p "command-prefix"
  nc = length cpre

cerrhandler p log clt = loop
  where
  loop = do x <- recvaux clt
            unless (null x) $
              do (n,msg) <- handle x
                 log n msg
            loop
  handle (c:' ':msg)
    | isDigit c = return (ord c-ord '0', msg)
  handle msg = return (1, msg)


connect p log clt hst prt
  = do log 0 ("connecting to "++hst++", port "++show prt)
       status ("connecting "++hst++" "++show prt)
       h <- connectTo hst (PortNumber (fromIntegral prt))
       log 0 ("connection established.")
       status ("connected "++hst++" "++show prt)
       return (h,h,h)
  where
  status
    | not (null spre) = \s -> send clt (spre++s)
    | otherwise       = \_ -> return ()
  spre = p "status-prefix"


start p log cmd
  = do log 1 ("starting client "++cmd)
       (i,o,e,p) <- handle (\e -> do log 0 (show e)
                                     throwIO e)
                           (runInteractiveCommand cmd)
       log 1 "client started."
       forkIO (waitloop p)
       return (i,o,e)
  where
  waitloop p = do threadDelay 1000000
                  x <- getProcessExitCode p
                  maybe (waitloop p) (\_ -> return ()) x


logger p lgc
  = do lgf <- openFile (p "logfile") AppendMode
       loop lgf
  where
  loop lgf = do (n',x) <- atomically (readTChan lgc)
                let n = max 0 (min 9 n')
                ts <- timestamp
                let p = show n ++ " " ++ ts ++ " "
                    y = fmt p x
                hPutStr lgf y
                hFlush lgf
                when (n==0) (putStr y)
                loop lgf

fmt p s = unlines (y : ys)
  where
  y = p ++ x
  ys = map (indent++) xs
  indent = map (const ' ') p
  (x:xs) = if null s then [""] else lines s

timestamp
  = do t <- getClockTime
       let u  = toUTCTime t
           fs = "%Y-%m-%d %T"
           ts = formatCalendarTime defaultTimeLocale fs u
       return ts