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

module Snap.SessionUtil where

import Data.Word
import Snap.Iteratee (enumBS)
import Snap.Types hiding (redirect)
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 = do
    modifyRequest  req
    modifyResponse (updateHeaders (M.update cleanRsp "Set-Cookie"))
    modifyResponse (addCookie cookie)
  where
    req | B.null (cookieValue cookie)
            = \r -> r { rqCookies = cleanReq (rqCookies r) }
        | otherwise
            = \r -> r { rqCookies = cookie : cleanReq (rqCookies r) }
    cleanReq = filter ((/= cookieName cookie) . cookieName)
    cleanRsp = Just . filter (not . ((cookieName cookie `B.append` "=") `B.isPrefixOf`))

{-|
    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)

{-|
    Temporarily here, in the hopes that the type of the original will be fixed
    in a future version of Snap.
-}
redirect :: MonadSnap m => ByteString -> m a
redirect target = do
    r <- getResponse
    finishWith
        $ setResponseCode 302
        $ setContentLength 0
        $ modifyResponseBody (const $ enumBS "")
        $ setHeader "Location" target r

{-|
    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