{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings          #-}

module Snap.SessionUtil where

import Data.Word
import Snap.Types
import System.Random

import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B

import Data.Map (Map)
import qualified Data.Map as M

{-|
    Sets a cookie in both the request and the response.  This modifies the
    list of cookies in the request, so that later attempts to get cookies will
    find this one even within the same request.
-}
setCookie :: MonadSnap m => Cookie -> m ()
setCookie cookie = modifyResponse (addResponseCookie cookie)

{-|
    Retrieves a cookie, looking first in the response map, and if not there,
    then in the request.  This ensures that the most recently set cookie is
    retrieved.
-}
lookupCookie :: MonadSnap m => ByteString -> m (Maybe Cookie)
lookupCookie name = do
    rsp <- getResponse
    maybe (getCookie name) (return . Just) (getResponseCookie name rsp)

{-|
    Clears a cookie.  This involves setting the cookie to the empty string,
    with an expiration time in the past.
-}
clearCookie :: MonadSnap m => ByteString -> m ()
clearCookie name = setCookie cookie
    where cookie = Cookie name "" Nothing Nothing Nothing

{-|
    If there is another path component in the request path, pop it off, and
    pass it as a parameter to the handler.
-}
popPathTo :: MonadSnap m => (ByteString -> m a) -> m a
popPathTo handler = do
    req <- getRequest
    let (x,y) = B.break (== '/') (rqPathInfo req)
    if B.null x
        then pass
        else localRequest (\r -> r {rqPathInfo = B.drop 1 y}) (handler x)

{-|
    Ensure that we're at the top level of a request, and expect that it be a
    directory.  As with standard HTTP behavior, if a path to a directory is
    given and the request URI doesn't end in a slash, then the user is
    redirected to a path ending in a slash.
-}
ifTopDir :: MonadSnap m => m a -> m a
ifTopDir handler = ifTop $ do
    req <- getRequest
    if (B.last (rqURI req) /= '/')
        then redirect $ rqURI req `B.append` "/" `B.append` rqQueryString req
        else handler

{-|
    Ensure that we're at the top level of a request, and expect that it be a
    file.  If a trailing slash is given, we pass on the request.
-}
ifTopFile :: MonadSnap m => m a -> m a
ifTopFile handler = ifTop $ do
    req <- getRequest
    if (B.last (rqURI req) == '/') then pass else handler

{-|
    Session keys are 64-bit integers with standard numeric type classes.
-}
newtype SessionKey = K Word64
    deriving (Eq, Ord, Enum, Bounded, Num, Real, Integral)

instance Random SessionKey where
    randomR (l, h) g = let (r, g') = randomR (toInteger l, toInteger h) g
                       in  (fromInteger r, g')
    random           = randomR (minBound, maxBound)

instance Show SessionKey where
    show (K n) = show n

instance Read SessionKey where
    readsPrec n s = map (\(a,t) -> (K a, t)) (readsPrec n s)

{-|
    Generates a random key that is not already used in the given map.  Though
    not technically speaking guaranteed to terminate, this should be fast in
    practice.
-}
uniqueKey :: (Random k, Ord k) => Map k a -> IO k
uniqueKey m = do k <- randomIO
                 if M.member k m then uniqueKey m else return k