{-# LANGUAGE ScopedTypeVariables, FlexibleContexts, TupleSections #-}
module Distribution.Skete.TarUtils (
    explodeTar
  , implodeTar
  , TarExplodeError(..)
  ) where

import qualified Codec.Archive.Tar as Tar
import qualified Codec.Archive.Tar.Entry as Tar
import qualified Codec.Archive.Tar.Check as Tar
import qualified Codec.Compression.GZip as GZip
import           Control.Monad.Except
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import           Data.Bits ((.|.))
import qualified Data.ByteString.Lazy as BSL
import           Data.Git.FileTree (FileData)
import qualified Data.ListTrie.Map.Ord as LT
import qualified Control.Exception as EB
import           System.FilePath
import           System.Filesystem.FileTree
import           System.Filesystem.PathComponent
import           System.Posix.Files

data TarExplodeError =
    TarSymlinkError (LT.TrieMap FilePath FilePath)
  | TarFormatError Tar.FormatError
  | TarFileNameError Tar.FileNameError
  | TarBombError Tar.TarBombError
  | BadPathComponent
  deriving (Show)

instance EB.Exception TarExplodeError

explodeTar :: FilePath -> BSL.ByteString -> Either TarExplodeError (FileTree FileData)
explodeTar expectedSubDir bs = runExcept $ do
    td <- transError . fmap (uncurry resolveLinks . filesAndLinks) . entryTree . Tar.checkTarbomb expectedSubDir . Tar.checkSecurity . Tar.read . GZip.decompress $ bs
    mapKeysFromTar td
  where
    transError :: Either (Either (Either Tar.FormatError Tar.FileNameError) Tar.TarBombError) (Either (LT.TrieMap FilePath FilePath) (LT.TrieMap FilePath FileData))
               -> Except TarExplodeError (LT.TrieMap FilePath FileData)
    transError (Right (Right good))      = pure good
    transError (Right (Left symerr))     = throwError . TarSymlinkError $ symerr
    transError (Left (Left (Left tfe)))  = throwError . TarFormatError $ tfe
    transError (Left (Left (Right fne))) = throwError . TarFileNameError $ fne
    transError (Left (Right tbe))        = throwError . TarBombError $ tbe


mapKeysFromTar :: LT.TrieMap FilePath a -> Except TarExplodeError (FileTree a)
mapKeysFromTar t = fmap LT.fromList . forM (LT.toList t) $ \(p, d) -> do
  case (,d) <$> (mapM (pathComponent . TE.encodeUtf8 . T.pack) p) of
    Nothing -> throwError BadPathComponent
    Just r  -> pure r

mapKeysToTar :: FileTree a -> LT.TrieMap FilePath a
mapKeysToTar =   LT.mapKeys (map $ T.unpack . TE.decodeUtf8 . getPC)

implodeTar :: FileTree FileData -> BSL.ByteString
implodeTar dt = Tar.write . (flip map) (LT.toList . mapKeysToTar $ dt) $ \(fp, (bs, exec)) ->
  Tar.Entry
  (either (const (error "No toTarPath")) id . Tar.toTarPath False . joinPath $ fp)
  (Tar.NormalFile bs (BSL.length bs))
  (if exec then Tar.executableFilePermissions else Tar.ordinaryFilePermissions)
  (Tar.Ownership "" "" 0 0)
  0
  Tar.V7Format

splitEntryPath :: Tar.Entry -> [FilePath]
splitEntryPath = splitDirectories . Tar.entryPath

entryTree :: forall a. Show a => Tar.Entries a -> Either a (LT.TrieMap FilePath Tar.Entry)
entryTree = Tar.foldEntries entry (Right LT.empty) (Left)
  where
    entry :: Tar.Entry -> Either a (LT.TrieMap FilePath Tar.Entry) -> Either a (LT.TrieMap FilePath Tar.Entry)
    entry e = fmap (LT.insert (splitEntryPath e) e)

linkTarget :: Tar.Entry -> Maybe FilePath
linkTarget e = case Tar.entryContent e of
                 Tar.SymbolicLink t -> target t
                 Tar.HardLink     t -> target t
                 _                  -> Nothing
    where target = Just . (dir </>) . Tar.fromLinkTarget
          dir    = takeDirectory . Tar.fromTarPath . Tar.entryTarPath $ e

filesAndLinks :: LT.TrieMap FilePath Tar.Entry -> (LT.TrieMap FilePath FileData, LT.TrieMap FilePath FilePath)
filesAndLinks t = LT.mapMaybe linkTarget <$> LT.mapEither part t
    where part e | Tar.NormalFile fd _ <- Tar.entryContent e = Left (fd, isExecutable e)
                 | otherwise                                = Right e

resolveLink :: LT.TrieMap FilePath a -> FilePath -> Either a FilePath
resolveLink fd tgt = maybe (Right tgt) Left (splitTarget tgt `LT.lookup` fd)
    where splitTarget = dedot . splitDirectories
          dedot []          = []
          dedot (".":ps)    = dedot ps
          dedot (_:"..":ps) = dedot ps -- fine unless a path starts with ".."?
          dedot (p:ps)      = p : dedot ps


resolveLinks :: LT.TrieMap FilePath FileData -> LT.TrieMap FilePath FilePath -> Either (LT.TrieMap FilePath FilePath) (LT.TrieMap FilePath FileData)
resolveLinks fd fp | LT.null fp = Right fd
                   | unresolved = Left u
                   | otherwise  = resolveLinks (fd `LT.union` r) u
    where (r,u) = LT.mapEither (resolveLink fd) fp
          unresolved = LT.size fp == (LT.size u :: Int)

isExecutable :: Tar.Entry -> Bool
isExecutable e = 0 /= Tar.entryPermissions e
                     `intersectFileModes`
                     (ownerExecuteMode .|. groupExecuteMode .|. otherExecuteMode)