{-# LANGUAGE PackageImports, DeriveDataTypeable #-}
-- | Miscellaneous functions. Not really for public consumption.
module Network.Hermes.Misc where

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import System.Random(RandomGen)
import Control.Arrow
import Data.Bits
import Control.Exception
import Data.Data
import Control.Applicative
import Control.Concurrent.STM
import "monads-tf" Control.Monad.State
import System.Log.Logger
import Control.Concurrent
import Data.Generics
import Data.Maybe
import Codec.Crypto.RSA
import qualified Data.Map as M
import Data.Map(Map)

ghead :: (Data x, Typeable y) => x -> y
ghead = fromJust . gfindtype

byteStringToInteger :: B.ByteString -> Integer
byteStringToInteger = fst . B.foldl (\(acc,scale) word -> (acc + shiftL (toInteger word) scale, scale+8)) (0,0)

runTMVar :: TMVar s -> StateT s IO a -> IO a
runTMVar var act = block $ do
  initial <- atomically $ takeTMVar var
  (retVal, final) <- unblock (runStateT act initial) `onException` (atomically $ putTMVar var initial)
  atomically $ putTMVar var final
  return retVal
  
throwM :: (MonadIO m, Exception e) => e -> m a
throwM = liftIO . throwIO

modifyTVar :: TVar a -> (a -> a) -> STM ()
modifyTVar var f = writeTVar var =<< f <$> readTVar var

-- | Our notion of types: A shown Typeable.
type Type = String

showType :: Typeable a => a -> Type
showType = show . typeOf

-- | Logs any unhandled exceptions
trapForkIO :: String -> IO () -> IO ThreadId
trapForkIO mod act = forkIO (traplogging mod EMERGENCY ("trap: " ++ mod) act)

-- | Encryption stuff

rsaEncrypt :: (RandomGen g) => g -> PublicKey -> B.ByteString -> (B.ByteString,g)
rsaEncrypt g key bs = first (B.concat . BL.toChunks) $ encrypt g key (BL.fromChunks [bs])

rsaDecrypt :: PrivateKey -> B.ByteString -> B.ByteString
rsaDecrypt key = B.concat . BL.toChunks . decrypt key . BL.fromChunks . (:[])

rsaVerify :: PublicKey -> B.ByteString -> B.ByteString -> Bool
rsaVerify key msg sig = verify key (BL.fromChunks [msg]) (BL.fromChunks [sig])

rsaSign :: PrivateKey -> B.ByteString -> B.ByteString
rsaSign key msg = B.concat $ BL.toChunks $ sign key (BL.fromChunks [msg])

-- | Swap values in a Map. Returns the old value, if any.
swap :: Ord k => k -> v -> M.Map k v -> (Maybe v, M.Map k v)
swap = M.insertLookupWithKey (\_ v _ -> v)

-- | Executes an action once for each value of the TVar. May skip
-- values if it changes quickly.
listenTVar :: Eq a => TVar a -> (a -> IO ()) -> IO ThreadId
listenTVar var act = forkIO $ do
  cur <- atomically $ readTVar var
  act cur
  listen cur
  where listen cur = do
          new <- atomically $ do
            new <- readTVar var
            check (new /= cur)
            return new
          act new
          listen new

adjustWithDefault :: Ord k => a -> (a -> a) -> k -> Map k a -> Map k a
adjustWithDefault def f k m = M.alter adj k m
  where adj Nothing  = Just (f def)
        adj (Just a) = Just (f a)