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

module Snap.Extension.SessionUtil where

import Data.Word
import Snap.Iteratee hiding (map)
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

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