module Compression ( Compression (Lzma, Lz4, Lzip, BZip)
                   , CompressionLevel (..)
                   , 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
#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 qualified Data.ByteString.Lazy         as BSL
import           Data.List                    (isSuffixOf)
import           System.FilePath              (dropExtension, (-<.>))
import           System.IO                    (IOMode (ReadMode), hFileSize, withFile)

data CompressionLevel = Best
                      | Fastest
                      | Default
                      | Custom !Int

data Compression = Lzma
                 | Lzip
                 | BZip
                 | GZip
                 | Zstd
#ifdef BROTLI
                 | Brotli
#endif
#ifdef SNAPPY
                 | Snappy
#endif
                 | Lz4
                 | Z
                 | None
                 deriving (Enum)

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 None   = ""

uncompressedExt :: FilePath -> FilePath
uncompressedExt fp | ".tlz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".txz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".tlz" `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"
                   | otherwise                 = dropExtension fp

detectCompression :: FilePath -> Compression
detectCompression fp
    | ".xz" `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
    | 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 None   = id

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

levelGuard :: (Int, Int) -> Int -> Int
levelGuard (min', max') lvl | lvl < min' || lvl > max' =
    error ("Invalid compression level. Compression must be between " ++ show min' ++ " and " ++ show max')
                            | otherwise = lvl

toFileCompressor :: Compression -> CompressionLevel -> FilePath -> IO BSL.ByteString
toFileCompressor Lzip lvl = Lzip.compressFileLevel (toEnum $ toInt Lzip lvl)
toFileCompressor x lvl    = fmap (toCompressor x lvl Nothing) . BSL.readFile

gzipCompression :: Int -> GZip.CompressParams
gzipCompression lvl = GZip.defaultCompressParams { GZip.compressLevel = GZip.compressionLevel lvl }

zlibCompression :: Int -> Zlib.CompressParams
zlibCompression lvl = Zlib.defaultCompressParams { Zlib.compressLevel = Zlib.compressionLevel lvl }

lzmaCompression :: Int -> Lzma.CompressParams
lzmaCompression lvl = Lzma.defaultCompressParams { Lzma.compressLevel = toEnum lvl }

#ifdef BROTLI
brotliCompression :: Int -> Br.CompressParams
brotliCompression lvl = Br.defaultCompressParams { Br.compressLevel = toEnum lvl }
#endif

toInt :: Compression -> CompressionLevel -> Int
#ifdef SNAPPY
toInt Snappy _          = undefined
#endif
#ifdef BROTLI
toInt Brotli Fastest    = fromEnum (minBound :: Br.CompressionLevel)
toInt Brotli (Custom i) = levelGuard (0,11) i
toInt Brotli _          = fromEnum (maxBound :: Br.CompressionLevel)
#endif
toInt Zstd Best         = Zstd.maxCLevel
toInt Zstd Fastest      = 1
toInt Zstd (Custom i)   = levelGuard (1, Zstd.maxCLevel) i
toInt Zstd Default      = 3
toInt Lzip Best         = fromEnum (maxBound :: Lzip.CompressionLevel)
toInt Lzip Fastest      = fromEnum (minBound :: Lzip.CompressionLevel)
toInt Lzip (Custom i)   = levelGuard (0, 9) i
toInt Lzip Default      = 6
toInt Lzma Best         = 9
toInt Lzma Fastest      = 0
toInt Lzma (Custom i)   = levelGuard (0, 9) i
toInt Lzma Default      = 6
toInt BZip Best         = 9
toInt BZip Fastest      = 1
toInt BZip (Custom i)   = i
toInt BZip Default      = 7
toInt GZip Best         = 9
toInt GZip Fastest      = 0
toInt GZip (Custom i)   = i
toInt GZip Default      = 6
toInt Z Best            = 9
toInt Z Fastest         = 0
toInt Z (Custom i)      = i
toInt Z Default         = 6
toInt Lz4 Best          = Lz4.lZ4HCClevelMax
toInt Lz4 Fastest       = 0
toInt Lz4 (Custom i)    = levelGuard (0, Lz4.lZ4HCClevelMax) i
toInt Lz4 Default       = 0 -- 1?
toInt None _            = error "Internal error."

toCompressor :: Compression -> CompressionLevel -> Maybe Int -> BSL.ByteString -> BSL.ByteString
toCompressor Lzma lvl _         = Lzma.compressWith (lzmaCompression $ toInt Lzma lvl)
toCompressor Lzip lvl (Just sz) = flip (Lzip.compressWithSz (toEnum $ toInt Lzip lvl)) sz
toCompressor Lzip _ _           = error "Internal error."
toCompressor BZip lvl _         = BZip.compressWith (fromIntegral $ toInt BZip lvl) 30
toCompressor GZip lvl _         = GZip.compressWith (gzipCompression $ toInt GZip lvl)
toCompressor Z    lvl _         = Zlib.compressWith (zlibCompression $ toInt Z lvl)
toCompressor Zstd lvl _         = Zstd.compress (toInt Zstd lvl)
toCompressor Lz4  lvl _         = Lz4.compressSz (toInt Lz4 lvl)
#ifdef BROTLI
toCompressor Brotli lvl _       = Br.compressWith (brotliCompression $ toInt Brotli lvl)
#endif
#ifdef SNAPPY
toCompressor Snappy _ _         = Snappy.compress
#endif
toCompressor None _ _           = id

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