{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}
module Network.Wai.StaticCache
(cacheDir, staticfiles)
where

import Data.Conduit 
    -- ( ($$), (=$), runResourceT, ResourceT, ConduitM, awaitForever, yield )
import Data.Conduit.Combinators (sourceDirectoryDeep)
import Control.Monad.IO.Class (liftIO)
import qualified Data.Conduit.List as CL
import Prelude hiding (FilePath)
import Filesystem.Path ( FilePath, extension, stripPrefix )
import Filesystem.Path.CurrentOS ( encodeString )
import Data.Word (Word64)
import Data.Digest.CityHash (cityHash64)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC
import Data.Map.Lazy (Map)
import qualified Data.Map.Lazy as M
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Vector as V
import Data.Vector.Unboxed (Vector, findIndex)
import qualified Data.Vector.Unboxed as VU
import Network.Wai (Response, Middleware, responseFile, rawPathInfo)
import Network.HTTP.Types (status200)
import Control.Applicative ((<$>))

data MetaFile = MF { fileId  :: !Word64
                   , path    :: !String
                   , rawPath :: !ByteString
                   , mime    :: !ByteString
                   , etag    :: !ByteString
                   } deriving (Show)

data FileCache = FC !(Vector Word64) !(V.Vector MetaFile)
                    deriving (Show)

-- responseFile :: Status -> ResponseHeaders -> FilePath -> Maybe FilePart -> Response

cacheDir :: FilePath -> IO FileCache
cacheDir fp = readDir fp >>= return . buildCache

staticfiles :: FileCache -> Middleware 
staticfiles fc app req = fromMaybe (app req) response
    where 
        response :: Maybe (IO Response)
        response = do 
            let bs = rawPathInfo req
                isFolder = if BSC.null bs then True else (BSC.last bs == '/')
                reqBS = if isFolder then BSC.append bs "index.html" else bs
            mf <- getMetaFile fc reqBS
            let headers = [ ("Content-Type", mime mf), ("ETag", etag mf) ]
            return . return $ responseFile status200 headers (path mf) Nothing

getMetaFile :: FileCache -> ByteString -> Maybe MetaFile
getMetaFile (FC ix m) bs = (V.!) m <$> maybePos
    where
        maybePos :: Maybe Int
        maybePos = findIndex (== (cityHash64 bs)) ix

buildCache :: [MetaFile] -> FileCache
buildCache lst = FC ix metaVec
    where
        metaVec = V.fromList lst
        ix = VU.generate (V.length metaVec) (\i -> fileId (metaVec V.! i))

readDir :: FilePath -> IO [MetaFile]
-- readDir :: String -> IO [MetaFile]
readDir f = do 
    l <- runResourceT $ sourceDirectoryDeep False f $$ readIt =$ CL.consume
    runResourceT $ sequence l

-- Reads a filepath from upstream and returns the corresponding MetaFile for it
readIt :: ConduitM FilePath (ResourceT IO MetaFile) (ResourceT IO) ()
readIt = CL.map $ \i -> (liftIO $ toMetaFile i)
    where
        toMetaFile f = do 
            let fpath = encodeString (fromMaybe f $ stripPrefix "." f)
                rpath = BSC.dropWhile (/= '/') $ BSC.pack fpath
                fmime = getMimeType f
            file <- BS.readFile fpath
            let ftag = BSC.pack $ show $ cityHash64 file
            return $ MF (cityHash64 rpath) 
                        fpath
                        rpath
                        fmime
                        ftag

getMimeType :: FilePath -> ByteString
getMimeType fp = fromMaybe defaultMimeType $ do 
    e <- extension fp
    M.lookup e defaultMimeTypes

defaultMimeType :: ByteString
defaultMimeType = "application/octet-stream"

-- This list taken from snap-core's Snap.Util.FileServe
defaultMimeTypes :: Map Text ByteString
defaultMimeTypes = M.fromList [
  ( "asc"     , "text/plain"                        ),
  ( "asf"     , "video/x-ms-asf"                    ),
  ( "asx"     , "video/x-ms-asf"                    ),
  ( "avi"     , "video/x-msvideo"                   ),
  ( "bz2"     , "application/x-bzip"                ),
  ( "c"       , "text/plain"                        ),
  ( "class"   , "application/octet-stream"          ),
  ( "conf"    , "text/plain"                        ),
  ( "cpp"     , "text/plain"                        ),
  ( "css"     , "text/css"                          ),
  ( "cxx"     , "text/plain"                        ),
  ( "dtd"     , "text/xml"                          ),
  ( "dvi"     , "application/x-dvi"                 ),
  ( "gif"     , "image/gif"                         ),
  ( "gz"      , "application/x-gzip"                ),
  ( "hs"      , "text/plain"                        ),
  ( "htm"     , "text/html"                         ),
  ( "html"    , "text/html"                         ),
  ( "jar"     , "application/x-java-archive"        ),
  ( "jpeg"    , "image/jpeg"                        ),
  ( "jpg"     , "image/jpeg"                        ),
  ( "js"      , "text/javascript"                   ),
  ( "json"    , "application/json"                  ),
  ( "log"     , "text/plain"                        ),
  ( "m3u"     , "audio/x-mpegurl"                   ),
  ( "mov"     , "video/quicktime"                   ),
  ( "mp3"     , "audio/mpeg"                        ),
  ( "mpeg"    , "video/mpeg"                        ),
  ( "mpg"     , "video/mpeg"                        ),
  ( "ogg"     , "application/ogg"                   ),
  ( "pac"     , "application/x-ns-proxy-autoconfig" ),
  ( "pdf"     , "application/pdf"                   ),
  ( "png"     , "image/png"                         ),
  ( "ps"      , "application/postscript"            ),
  ( "qt"      , "video/quicktime"                   ),
  ( "sig"     , "application/pgp-signature"         ),
  ( "spl"     , "application/futuresplash"          ),
  ( "svg"     , "image/svg+xml"                     ),
  ( "swf"     , "application/x-shockwave-flash"     ),
  ( "tar"     , "application/x-tar"                 ),
  ( "tar.bz2" , "application/x-bzip-compressed-tar" ),
  ( "tar.gz"  , "application/x-tgz"                 ),
  ( "tbz"     , "application/x-bzip-compressed-tar" ),
  ( "text"    , "text/plain"                        ),
  ( "tgz"     , "application/x-tgz"                 ),
  ( "torrent" , "application/x-bittorrent"          ),
  ( "ttf"     , "application/x-font-truetype"       ),
  ( "txt"     , "text/plain"                        ),
  ( "wav"     , "audio/x-wav"                       ),
  ( "wax"     , "audio/x-ms-wax"                    ),
  ( "wma"     , "audio/x-ms-wma"                    ),
  ( "wmv"     , "video/x-ms-wmv"                    ),
  ( "xbm"     , "image/x-xbitmap"                   ),
  ( "xml"     , "text/xml"                          ),
  ( "xpm"     , "image/x-xpixmap"                   ),
  ( "xwd"     , "image/x-xwindowdump"               ),
  ( "zip"     , "application/zip"                   ) ]