module Compression ( Compression (..)
                   , detectCompression
                   , toCompressor
                   , toDecompressor
                   , toFileCompressor
                   , toFileDecompressor
                   , uncompressedExt
                   , check
                   , ext
                   -- * Utilities
                   , fileSize
                   ) where

#ifdef BROTLI
import qualified Codec.Compression.Brotli     as Br
#endif
import qualified Codec.Compression.BZip       as BZip
import qualified Codec.Compression.GZip       as GZip
import qualified Codec.Compression.Lzma       as Lzma
import qualified Codec.Compression.Lzo        as Lzo
#ifdef SNAPPY
import qualified Codec.Compression.Snappy.BSL as Snappy
#endif
import qualified Codec.Compression.Zlib       as Zlib
import qualified Codec.Compression.Zstd.Lazy  as Zstd
import qualified Codec.Lz4                    as Lz4
import qualified Codec.Lzip                   as Lzip
import           Compression.Level
import           Compression.Type
import qualified Data.ByteString.Lazy         as BSL
import           Data.List                    (isSuffixOf)
import           System.FilePath              (dropExtension, (-<.>))
import           System.IO                    (IOMode (ReadMode), hFileSize, withFile)

ext :: Compression -> String
ext Lzma   = ".xz"
ext Lzip   = ".lz"
ext BZip   = ".bz2"
ext GZip   = ".gz"
ext Zstd   = ".zst"
ext Lz4    = ".lz4"
ext Z      = ".Z"
#ifdef BROTLI
ext Brotli = ".br"
#endif
#ifdef SNAPPY
ext Snappy = ".sz"
#endif
ext Lzo    = ".lzo"
ext None   = ""

uncompressedExt :: FilePath -> FilePath
uncompressedExt fp | ".tlz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".txz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".tlz4" `isSuffixOf` fp   = fp -<.> ".tar"
                   | ".tgz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".cpgz" `isSuffixOf` fp   = fp -<.> ".cpio"
                   | ".cpbz2" `isSuffixOf` fp  = fp -<.> ".cpio"
                   | ".xcfgz" `isSuffixOf` fp  = fp -<.> ".xcf"
                   | ".xcfbz2" `isSuffixOf` fp = fp -<.> ".xcf"
                   | ".tbz2" `isSuffixOf` fp   = fp -<.> "tar"
                   | ".tbz" `isSuffixOf` fp    = fp -<.> "tar"
                   | ".tbr" `isSuffixOf` fp    = fp -<.> "tar"
                   | ".tsz" `isSuffixOf` fp    = fp -<.> "tar"
                   | ".tzo" `isSuffixOf` fp    = fp -<.> "tar"
                   | otherwise                 = dropExtension fp

detectCompression :: FilePath -> Compression
detectCompression fp
    | ".xz" `isSuffixOf` fp     = Lzma
    | ".lzma" `isSuffixOf` fp = Lzma
    | ".txz" `isSuffixOf` fp    = Lzma
    | ".lz" `isSuffixOf` fp     = Lzip
    | ".tlz" `isSuffixOf` fp    = Lzip
    | ".tbz2" `isSuffixOf` fp   = BZip
    | ".tbz" `isSuffixOf` fp    = BZip
    | ".bz2" `isSuffixOf` fp    = BZip
    | ".gz" `isSuffixOf` fp     = GZip
    | ".tgz" `isSuffixOf` fp    = GZip
    | ".Z" `isSuffixOf` fp      = Z
    | ".zst" `isSuffixOf` fp    = Zstd
    | ".lz4" `isSuffixOf` fp    = Lz4
    | ".xcfgz" `isSuffixOf` fp  = GZip
    | ".xcfbz2" `isSuffixOf` fp = BZip
    | ".cpgz" `isSuffixOf` fp   = GZip
    | ".cpbz2" `isSuffixOf` fp  = BZip
#ifdef BROTLI
    | ".br" `isSuffixOf` fp     = Brotli
#endif
#ifdef SNAPPY
    | ".sz" `isSuffixOf` fp     = Snappy
#endif
    | ".lzo" `isSuffixOf` fp    = Lzo
    | ".tzo" `isSuffixOf` fp    = Lzo
    | otherwise                 = None -- error "Suffix not supported or invalid"

toFileDecompressor :: Compression -> FilePath -> IO BSL.ByteString
toFileDecompressor Lz4 = lz4DecompressFile
toFileDecompressor x   = fmap (toDecompressor x) . BSL.readFile

toDecompressor :: Compression -> BSL.ByteString -> BSL.ByteString
toDecompressor Lzma   = Lzma.decompress
toDecompressor Lzip   = Lzip.decompress
toDecompressor BZip   = BZip.decompress
toDecompressor GZip   = GZip.decompress
toDecompressor Z      = Zlib.decompress
toDecompressor Zstd   = Zstd.decompress
toDecompressor Lz4    = Lz4.decompress
#ifdef BROTLI
toDecompressor Brotli = Br.decompress
#endif
#ifdef SNAPPY
toDecompressor Snappy = Snappy.decompress
#endif
toDecompressor Lzo    = Lzo.decompressFile
toDecompressor None   = id

check :: Compression -> BSL.ByteString -> IO ()
check = (forceBSL .) . toDecompressor

fileSize :: FilePath -> IO Integer
fileSize fp = withFile fp ReadMode hFileSize

lz4DecompressFile :: FilePath -> IO BSL.ByteString
lz4DecompressFile fp = do
    fSz <- fileSize fp
    let f = if fSz <= 32 * 1024
        then Lz4.decompressBufSz (32 * 1024)
        else Lz4.decompressBufSz (128 * 1024)
    f <$> BSL.readFile fp

forceBSL :: BSL.ByteString -> IO ()
forceBSL bsl = last (BSL.toChunks bsl) `seq` mempty