{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Snap.Util.FileServe.Tests
  ( tests ) where

import           Blaze.ByteString.Builder
import           Control.Monad
import           Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as S
import qualified Data.ByteString.Lazy.Char8 as L
import           Data.IORef
import qualified Data.Map as Map
import           Data.Maybe
import           Data.Monoid
import           Prelude hiding (take)
import           Test.Framework
import           Test.Framework.Providers.HUnit
import           Test.HUnit hiding (Test, path)

import           Snap.Internal.Http.Types
import           Snap.Internal.Types
import           Snap.Util.FileServe
import           Snap.Iteratee

tests :: [Test]
tests = [ testFooBin
        , testFooTxt
        , testFooHtml
        , testFooBinBinBin
        , test404s
        , testFsSingle

        , testFsCfgA
        , testFsCfgB
        , testFsCfgC
        , testFsCfgD
        , testFsCfgFancy

        , testRangeOK
        , testRangeBad
        , testMultiRange
        , testIfRange ]


expect404 :: IO Response -> IO ()
expect404 m = do
    r <- m
    assertBool "expected 404" (rspStatus r == 404)


expect302 :: ByteString -> IO Response -> IO ()
expect302 p m = do
    r <- m
    assertBool "expected 302" (rspStatus r == 302)
    assertEqual "redir location"
                (Just p)
                (getHeader "location" r)


getBody :: Response -> IO L.ByteString
getBody r = do
    let benum = rspBodyToEnum $ rspBody r
    liftM (toLazyByteString . mconcat) (runIteratee consume >>= run_ . benum)


runIt :: Snap a -> Request -> Iteratee ByteString IO (Request, Response)
runIt m rq = runSnap m d d rq
  where
    d = const $ return ()


go :: Snap a -> ByteString -> IO Response
go m s = do
    rq <- mkRequest s
    liftM snd (run_ $ runIt m rq)


goIfModifiedSince :: Snap a -> ByteString -> ByteString -> IO Response
goIfModifiedSince m s lm = do
    rq <- mkRequest s
    let r = setHeader "if-modified-since" lm rq
    liftM snd (run_ $ runIt m r)


goIfRange :: Snap a -> ByteString -> (Int,Int) -> ByteString -> IO Response
goIfRange m s (start,end) lm = do
    rq <- mkRequest s
    let r = setHeader "if-range" lm $
            setHeader "Range"
                       (S.pack $ "bytes=" ++ show start ++ "-" ++ show end)
                       rq
    liftM snd (run_ $ runIt m r)


goRange :: Snap a -> ByteString -> (Int,Int) -> IO Response
goRange m s (start,end) = do
    rq' <- mkRequest s
    let rq = setHeader "Range"
                       (S.pack $ "bytes=" ++ show start ++ "-" ++ show end)
                       rq'
    liftM snd (run_ $ runIt m rq)


goMultiRange :: Snap a -> ByteString -> (Int,Int) -> (Int,Int) -> IO Response
goMultiRange m s (start,end) (start2,end2) = do
    rq' <- mkRequest s
    let rq = setHeader "Range"
                       (S.pack $ "bytes=" ++ show start ++ "-" ++ show end
                                 ++ "," ++ show start2 ++ "-" ++ show end2)
                       rq'
    liftM snd (run_ $ runIt m rq)


goRangePrefix :: Snap a -> ByteString -> Int -> IO Response
goRangePrefix m s start = do
    rq' <- mkRequest s
    let rq = setHeader "Range"
                       (S.pack $ "bytes=" ++ show start ++ "-")
                       rq'
    liftM snd (run_ $ runIt m rq)


goRangeSuffix :: Snap a -> ByteString -> Int -> IO Response
goRangeSuffix m s end = do
    rq' <- mkRequest s
    let rq = setHeader "Range"
                       (S.pack $ "bytes=-" ++ show end)
                       rq'
    liftM snd (run_ $ runIt m rq)


mkRequest :: ByteString -> IO Request
mkRequest uri = do
    enum <- newIORef $ SomeEnumerator returnI
    return $ Request "foo" 80 "foo" 999 "foo" 1000 "foo" False Map.empty
                     enum Nothing GET (1,1) [] "" pathPart "/"
                     (S.concat ["/",uri]) queryPart Map.empty

  where
    (pathPart, queryPart) = breakQuery uri

    breakQuery s = (a, S.drop 1 b)
      where
        (a,b) = S.break (=='?') s


fs :: Snap ()
fs = do
    x <- serveDirectory "data/fileServe"
    return $! x `seq` ()


fsSingle :: Snap ()
fsSingle = do
    x <- serveFile "data/fileServe/foo.html"
    return $! x `seq` ()


fsCfg :: DirectoryConfig Snap -> Snap ()
fsCfg cfg = do
    x <- serveDirectoryWith cfg "data/fileServe"
    return $! x `seq` ()



------------------------------------------------------------------------------
testFooBin :: Test
testFooBin = testCase "fileServe/foo.bin" $ do
    r1 <- go fs "foo.bin"
    checkProps "fileServe/foo.bin/default" r1
    let !lm = fromJust $ getHeader "last-modified" r1

    go fs "foo.bin?blah=blah" >>= checkProps "fileServe/foo.bin/query-string"

    -- check last modified stuff
    r2 <- goIfModifiedSince fs "foo.bin" lm
    assertEqual "foo.bin 304" 304 $ rspStatus r2

    r3 <- goIfModifiedSince fs "foo.bin" "Wed, 15 Nov 1995 04:58:08 GMT"
    checkProps "fileServe/foo.bin/ifModifiedSince" r3

  where
    checkProps name r = do
        b <- getBody r

        assertEqual (name ++ "/contents") "FOO\n" b
        assertEqual (name ++ "/content-type")
                    (Just "application/octet-stream")
                    (getHeader "content-type" r)

        assertEqual (name ++ "/size") (Just 4) (rspContentLength r)

        assertBool  (name ++ "/last-modified")
                    (isJust $ getHeader "last-modified" r)

        assertEqual (name ++ "/accept-ranges")
                    (Just "bytes")
                    (getHeader "accept-ranges" r)


------------------------------------------------------------------------------
testFooTxt :: Test
testFooTxt = testCase "fileServe/foo.txt" $ do
    go fs "foo.txt" >>= checkProps "fileServe/foo.txt/default"
    go fs "foo.txt?blah=blah" >>= checkProps "fileServe/foo.txt/query"

  where
    checkProps name r = do
        b <- getBody r

        assertEqual (name ++ "/contents") "FOO\n" b
        assertEqual (name ++ "/content-type")
                    (Just "text/plain")
                    (getHeader "content-type" r)

        assertEqual (name ++ "/size") (Just 4) (rspContentLength r)

        assertBool  (name ++ "/last-modified")
                    (isJust $ getHeader "last-modified" r)

        assertEqual (name ++ "/accept-ranges")
                    (Just "bytes")
                    (getHeader "accept-ranges" r)
    
    

------------------------------------------------------------------------------
testFooHtml :: Test
testFooHtml = testCase "fileServe/foo.html" $ do
    go fs "foo.html" >>= checkProps "fileServe/foo.html/default"
    go fs "foo.html?bar=bar" >>= checkProps "fileServe/foo.html/query"

  where
    checkProps name r = do
        b <- getBody r

        assertEqual (name ++ "/contents") "FOO\n" b
        assertEqual (name ++ "/content-type")
                    (Just "text/html")
                    (getHeader "content-type" r)

        assertEqual (name ++ "/size") (Just 4) (rspContentLength r)

        assertBool  (name ++ "/last-modified")
                    (isJust $ getHeader "last-modified" r)

        assertEqual (name ++ "/accept-ranges")
                    (Just "bytes")
                    (getHeader "accept-ranges" r)


------------------------------------------------------------------------------
testFooBinBinBin :: Test
testFooBinBinBin = testCase "fileServe/foo.bin.bin.bin" $ do
    go fs "foo.bin.bin.bin" >>= checkProps "fileServe/foo.bin.bin.bin"

  where
    checkProps name r = do
        b <- getBody r

        assertEqual (name ++ "/contents") "FOO\n" b
        assertEqual (name ++ "/content-type")
                    (Just "application/octet-stream")
                    (getHeader "content-type" r)

        assertEqual (name ++ "/size") (Just 4) (rspContentLength r)

        assertBool  (name ++ "/last-modified")
                    (isJust $ getHeader "last-modified" r)

        assertEqual (name ++ "/accept-ranges")
                    (Just "bytes")
                    (getHeader "accept-ranges" r)


------------------------------------------------------------------------------
test404s :: Test
test404s = testCase "fileServe/404s" $ do
    expect404 $ go fs "jfldksjflksd"
    expect404 $ go fs "dummy/../foo.txt"
    expect404 $ go fs "/etc/password"

    coverMimeMap


------------------------------------------------------------------------------
printName :: FilePath -> Snap ()
printName c = writeBS $ snd $ S.breakEnd (=='/') $ S.pack c


cfgA, cfgB, cfgC, cfgD :: DirectoryConfig Snap
cfgA = DirectoryConfig {
         indexFiles      = []
       , indexGenerator  = const pass
       , dynamicHandlers = Map.empty
       , mimeTypes       = defaultMimeTypes
       , preServeHook    = const $ return ()
       }

cfgB = DirectoryConfig {
         indexFiles      = ["index.txt", "altindex.html"]
       , indexGenerator  = const pass
       , dynamicHandlers = Map.empty
       , mimeTypes       = defaultMimeTypes
       , preServeHook    = const $ return ()
       }

cfgC = DirectoryConfig {
         indexFiles      = ["index.txt", "altindex.html"]
       , indexGenerator  = printName
       , dynamicHandlers = Map.empty
       , mimeTypes       = defaultMimeTypes
       , preServeHook    = const $ return ()
       }

cfgD = DirectoryConfig {
         indexFiles      = []
       , indexGenerator  = const pass
       , dynamicHandlers = Map.fromList [ (".txt", printName) ]
       , mimeTypes       = defaultMimeTypes
       , preServeHook    = const $ return ()
       }


testFsCfgA :: Test
testFsCfgA = testCase "fileServe/cfgA" $ do
    let gooo = go (fsCfg cfgA)

    -- Named file in the root directory
    gooo "foo.bin" >>= checkProps "cfgA1/1" "application/octet-stream"
    gooo "foo.bin?blah=blah" >>=
         checkProps "cfgA1/2" "application/octet-stream"

    -- Missing file in the root directory
    expect404 $ gooo "bar.bin"

    -- Named file in a subdirectory
    gooo "mydir2/foo.txt" >>= checkProps "cfgA1/subdir/1" "text/plain" 
    gooo "mydir2/foo.txt?z=z" >>= checkProps "cfgA1/subdir/2" "text/plain"
   
    -- Missing file in a subdirectory
    expect404 $ gooo "mydir2/bar.txt"

    -- Request for directory with no trailing slash
    expect302 "/mydir1/" $ gooo "mydir1"

    -- Request for directory with no trailing slash, with query param
    expect302 "/mydir1/?z=z" $ gooo "mydir1?z=z"

    -- Request for directory with trailing slash, no index
    expect404 $ gooo "mydir1/"
    expect404 $ gooo "mydir2/"

    -- Request file with trailing slash
    expect404 $ gooo "foo.html/"
    expect404 $ gooo "mydir2/foo.txt/"

  where
    checkProps name ct r = do
        b <- getBody r

        assertEqual (name ++ "/contents") "FOO\n" b
        assertEqual (name ++ "/content-type")
                    (Just ct)
                    (getHeader "content-type" r)

        assertEqual (name ++ "/size") (Just 4) (rspContentLength r)

        assertBool  (name ++ "/last-modified")
                    (isJust $ getHeader "last-modified" r)

        assertEqual (name ++ "/accept-ranges")
                    (Just "bytes")
                    (getHeader "accept-ranges" r)
    
testFsCfgB :: Test
testFsCfgB = testCase "fileServe/cfgB" $ do
    let gooo = go (fsCfg cfgB)

    -- Request for root directory with index
    rB1 <- gooo "mydir1/"
    bB1 <- getBody rB1

    assertEqual "B1" "INDEX\n" bB1
    assertEqual "B1 content-type"
                (Just "text/plain")
                (getHeader "content-type" rB1)

    -- Request for root directory with index, query
    rB2 <- gooo "mydir1/?z=z"
    bB2 <- getBody rB2

    assertEqual "B2" "INDEX\n" bB2
    assertEqual "B2 content-type"
                (Just "text/plain")
                (getHeader "content-type" rB2)


    -- Request for root directory with alternate index
    rB3 <- gooo "mydir3/"
    bB3 <- getBody rB3

    assertEqual "B3" "ALTINDEX\n" bB3
    assertEqual "B3 content-type"
                (Just "text/html")
                (getHeader "content-type" rB3)

    -- Request for root directory with no index
    expect404 $ gooo "mydir2/"


testFsCfgC :: Test
testFsCfgC = testCase "fileServe/cfgC" $ do
    let gooo = go (fsCfg cfgC)

    -- Request for root directory with index
    rC1 <- gooo "mydir1/"
    bC1 <- getBody rC1

    assertEqual "C1" "INDEX\n" bC1
    assertEqual "C1 content-type"
                (Just "text/plain")
                (getHeader "content-type" rC1)

    -- Request for root directory with index, query
    rC2 <- gooo "mydir1/?z=z"
    bC2 <- getBody rC2

    assertEqual "C2" "INDEX\n" bC2
    assertEqual "C2 content-type"
                (Just "text/plain")
                (getHeader "content-type" rC2)

    -- Request for root directory with generated index
    rC3 <- gooo "mydir2/"
    bC3 <- getBody rC3

    assertEqual "C3" "mydir2" bC3


testFsCfgD :: Test
testFsCfgD = testCase "fileServe/cfgD" $ do
    -- Request for file with dynamic handler
    rD1 <- go (fsCfg cfgD) "mydir2/foo.txt"
    bD1 <- getBody rD1

    assertEqual "D1" "foo.txt" bD1


testFsCfgFancy :: Test
testFsCfgFancy = testCase "fileServe/cfgFancy" $ do
    -- Request for directory with autogen index
    rE1 <- go (fsCfg fancyDirectoryConfig) "mydir2/"
    bE1 <- S.concat `fmap` L.toChunks `fmap` getBody rE1

    assertBool "autogen-sub-index" $
        "Directory Listing: /mydir2/" `S.isInfixOf` bE1
    assertBool "autogen-sub-parent" $
        "<a href='../'" `S.isInfixOf` bE1
    assertBool "autogen-sub-file" $
        "<a href='foo.txt'" `S.isInfixOf` bE1


    -- Request for directory with autogen index
    rE2 <- go (fsCfg fancyDirectoryConfig) "mydir2/?z=z"
    bE2 <- S.concat `fmap` L.toChunks `fmap` getBody rE2

    assertBool "autogen-sub-index" $
        "Directory Listing: /mydir2/" `S.isInfixOf` bE2
    assertBool "autogen-sub-parent" $
        "<a href='../'" `S.isInfixOf` bE2
    assertBool "autogen-sub-file" $
        "<a href='foo.txt'" `S.isInfixOf` bE2



testFsSingle :: Test
testFsSingle = testCase "fileServe/Single" $ do
    r1 <- go fsSingle "foo.html"
    b1 <- getBody r1

    assertEqual "foo.html" "FOO\n" b1
    assertEqual "foo.html content-type"
                (Just "text/html")
                (getHeader "content-type" r1)

    assertEqual "foo.html size" (Just 4) (rspContentLength r1)



testRangeOK :: Test
testRangeOK = testCase "fileServe/range/ok" $ do
    r1 <- goRange fsSingle "foo.html" (1,2)
    assertEqual "foo.html 206" 206 $ rspStatus r1
    b1 <- getBody r1

    assertEqual "foo.html partial" "OO" b1
    assertEqual "foo.html partial size" (Just 2) (rspContentLength r1)
    assertEqual "foo.html content-range"
                (Just "bytes 1-2/4")
                (getHeader "Content-Range" r1)

    r2 <- goRangeSuffix fsSingle "foo.html" 3
    assertEqual "foo.html 206" 206 $ rspStatus r2
    b2 <- getBody r2
    assertEqual "foo.html partial suffix" "OO\n" b2

    r3 <- goRangePrefix fsSingle "foo.html" 2
    assertEqual "foo.html 206" 206 $ rspStatus r3
    b3 <- getBody r3
    assertEqual "foo.html partial prefix" "O\n" b3


testMultiRange :: Test
testMultiRange = testCase "fileServe/range/multi" $ do
    r1 <- goMultiRange fsSingle "foo.html" (1,2) (3,3)

    -- we don't support multiple ranges so it's ok for us to return 200 here;
    -- test this behaviour
    assertEqual "foo.html 200" 200 $ rspStatus r1
    b1 <- getBody r1

    assertEqual "foo.html" "FOO\n" b1


testRangeBad :: Test
testRangeBad = testCase "fileServe/range/bad" $ do
    r1 <- goRange fsSingle "foo.html" (1,17)
    assertEqual "bad range" 416 (rspStatus r1)
    assertEqual "bad range content-range"
                (Just "bytes */4")
                (getHeader "Content-Range" r1)
    assertEqual "bad range content-length" (Just 0) (rspContentLength r1)
    b1 <- getBody r1
    assertEqual "bad range empty body" "" b1

    r2 <- goRangeSuffix fsSingle "foo.html" 4893
    assertEqual "bad suffix range" 416 $ rspStatus r2


coverMimeMap :: (Monad m) => m ()
coverMimeMap = Prelude.mapM_ f $ Map.toList defaultMimeTypes
  where
    f (!k,!v) = return $ case k `seq` v `seq` () of () -> ()


testIfRange :: Test
testIfRange = testCase "fileServe/range/if-range" $ do
    r <- goIfRange fs "foo.bin" (1,2) "Wed, 15 Nov 1995 04:58:08 GMT"
    assertEqual "foo.bin 200" 200 $ rspStatus r
    b <- getBody r
    assertEqual "foo.bin" "FOO\n" b

    r2 <- goIfRange fs "foo.bin" (1,2) "Tue, 01 Oct 2030 04:58:08 GMT"
    assertEqual "foo.bin 206" 206 $ rspStatus r2
    b2 <- getBody r2
    assertEqual "foo.bin partial" "OO" b2

    r3 <- goIfRange fs "foo.bin" (1,24324) "Tue, 01 Oct 2030 04:58:08 GMT"
    assertEqual "foo.bin 200" 200 $ rspStatus r3
    b3 <- getBody r3
    assertEqual "foo.bin" "FOO\n" b3