module Stackage.Install
    ( install
    , download
    , Settings
    , defaultSettings
    , setGetManager
    , setPackageLocation
    , defaultPackageLocation
    , setIndexLocation
    , defaultIndexLocation
    ) where
import qualified Codec.Archive.Tar        as Tar
import           Control.Applicative      ((*>), (<$>), (<*>))
import           Control.Concurrent.Async (wait, withAsync)
import           Control.Concurrent.Async (Concurrently (..))
import           Control.Concurrent.STM   (atomically, newTVarIO, readTVar,
                                           writeTVar)
import           Control.Exception        (Exception, throwIO)
import           Control.Monad            (join, unless, when)
import           Crypto.Hash              (Context, Digest, SHA512,
                                           digestToHexByteString, hashFinalize,
                                           hashInit, hashUpdate)
import           Data.Aeson               (FromJSON (..), decode, withObject,
                                           (.!=), (.:?))
import           Data.ByteString          (ByteString)
import qualified Data.ByteString          as S
import qualified Data.ByteString.Char8    as S8
import qualified Data.ByteString.Lazy     as L
import qualified Data.Foldable            as F
import           Data.Function            (fix)
import           Data.List                (isPrefixOf)
import           Data.Map                 (Map)
import qualified Data.Map                 as Map
import           Data.Set                 (Set)
import qualified Data.Set                 as Set
import           Data.Text                (Text)
import qualified Data.Text                as T
import           Data.Text.Encoding       (encodeUtf8)
import           Data.Typeable            (Typeable)
import           Data.Word                (Word64)
import           Network.HTTP.Client      (Manager, brRead,
                                           managerResponseTimeout, newManager,
                                           responseBody,
                                           responseStatus, withResponse
#if MIN_VERSION_http_client(0,5,0)
                                           , parseRequest
                                           , responseTimeoutMicro
#else
                                           , parseUrl
                                           , checkStatus
#endif
                                           )
import           Network.HTTP.Client.TLS  (tlsManagerSettings)
import           Network.HTTP.Types       (statusCode)
import           System.Directory         (createDirectoryIfMissing,
                                           doesFileExist,
                                           getAppUserDataDirectory, renameFile)
import           System.Exit              (ExitCode)
import           System.FilePath          (takeDirectory, (<.>), (</>), takeExtension)
import           System.IO                (IOMode (ReadMode, WriteMode), stdout,
                                           withBinaryFile)
import           System.Process           (rawSystem, readProcess)
install :: Settings -> [String] -> IO ExitCode
install s args = do
    out <- readProcess (_cabalCommand s)
        ("install":"--dry-run":if null args then ["."] else args)
        ""
    let pkgs = map toPair $ filter (not . toIgnore) $ lines out
    download s pkgs
    rawSystem (_cabalCommand s) ("install":args)
  where
    toIgnore str = ' ' `elem` str || '-' `notElem` str
    toPair :: String -> (String, String)
    toPair orig =
        (pkg, ver)
      where
        (ver', pkg') = break (== '-') $ reverse orig
        ver = reverse ver'
        pkg = reverse $ drop 1 pkg'
data Settings = Settings
    { _getManager     :: !(IO Manager)
    , _cabalCommand   :: !FilePath
    , _downloadPrefix :: !String
    , _onDownload     :: !(String -> IO ())
    , _onDownloadErr  :: !(String -> IO ())
    , _connections    :: !Int
    , _packageLocation :: !(IO (String -> String -> FilePath))
    , _indexLocation :: !(IO FilePath)
    }
defaultSettings :: Settings
defaultSettings = Settings
    { _getManager = newManager tlsManagerSettings
        { managerResponseTimeout =
#if MIN_VERSION_http_client(0,5,0)
            responseTimeoutMicro
#else
            Just
#endif
            90000000
        }
    , _cabalCommand = "cabal"
    , _downloadPrefix = "https://s3.amazonaws.com/hackage.fpcomplete.com/package/"
    , _onDownload = \s -> S8.hPut stdout $ S8.pack $ concat
        [ "Downloading "
        , s
        , "\n"
        ]
    , _onDownloadErr = \s -> S8.hPut stdout $ S8.pack $ concat
        [ "Error downloading "
        , s
        , ", if this is a local package, this message can be ignored\n"
        ]
    , _connections = 8
    , _packageLocation = defaultPackageLocation
    , _indexLocation = defaultIndexLocation
    }
setGetManager :: IO Manager -> Settings -> Settings
setGetManager x s = s { _getManager = x }
data Package = Package
    { packageHashes    :: Map Text Text
    , packageLocations :: [Text]
    , packageSize      :: Maybe Word64
    }
    deriving Show
instance FromJSON Package where
    parseJSON = withObject "Package" $ \o -> Package
        <$> o .:? "package-hashes" .!= Map.empty
        <*> o .:? "package-locations" .!= []
        <*> o .:? "package-size"
getPackageInfo :: FilePath -> Set (String, String) -> IO (Map (String, String) Package)
getPackageInfo indexTar pkgs0 = withBinaryFile indexTar ReadMode $ \h -> do
    lbs <- L.hGetContents h
    loop pkgs0 Map.empty False $ Tar.read lbs
  where
    loop pkgs m sawJSON Tar.Done = do
        when (not (Set.null pkgs) && sawJSON) $
            putStrLn $ "Warning: packages not found in index: " ++ show (Set.toList pkgs)
        return m
    loop _ m _ (Tar.Fail e) = throwIO $ Couldn'tReadIndexTarball indexTar e
    loop pkgs m sawJSON (Tar.Next e es) =
        case (getName $ Tar.entryPath e, Tar.entryContent e) of
            (Just pair, Tar.NormalFile lbs _)
                    | pair `Set.member` pkgs
                    , Just p <- decode lbs ->
                loop (Set.delete pair pkgs) (Map.insert pair p m) sawJSON' es
            _ -> loop pkgs m sawJSON' es
      where
        sawJSON' = sawJSON || takeExtension (Tar.entryPath e) == ".json"
    getName name =
        case T.splitOn "/" $ T.pack name of
            [pkg, ver, fp] | T.stripSuffix ".json" fp == Just pkg
                -> Just (T.unpack pkg, T.unpack ver)
            _ -> Nothing
data StackageInstallException
    = Couldn'tReadIndexTarball FilePath Tar.FormatError
    | InvalidDownloadSize
        { _idsUrl             :: String
        , _idsExpected        :: Word64
        , _idsTotalDownloaded :: Word64
        }
    | InvalidHash
        { _ihUrl      :: String
        , _ihExpected :: Text
        , _ihActual   :: Digest SHA512
        }
    deriving (Show, Typeable)
instance Exception StackageInstallException
defaultPackageLocation :: IO (String -> String -> FilePath)
defaultPackageLocation = do
    cabalDir <- getAppUserDataDirectory "cabal"
    let packageDir = cabalDir </> "packages" </> "hackage.haskell.org"
    return $ \name version ->
             packageDir </>
             name </>
             version </>
             concat [name, "-", version, ".tar.gz"]
setPackageLocation :: IO (String -> String -> FilePath) -> Settings -> Settings
setPackageLocation x s = s { _packageLocation = x }
setIndexLocation :: IO FilePath -> Settings -> Settings
setIndexLocation x s = s { _indexLocation = x }
defaultIndexLocation :: IO FilePath
defaultIndexLocation = do
    cabalDir <- getAppUserDataDirectory "cabal"
    return $ cabalDir </> "packages" </> "hackage.haskell.org" </> "00-index.tar"
download :: F.Foldable f => Settings -> f (String, String) -> IO ()
download s pkgs = do
    indexFP <- _indexLocation s
    packageLocation <- _packageLocation s
    withAsync (getPackageInfo indexFP $ Set.fromList $ F.toList pkgs) $ \a -> do
        man <- _getManager s
        parMapM_ (_connections s) (go packageLocation man (wait a)) pkgs
  where
    unlessM p f = do
        p' <- p
        unless p' f
    go packageLocation man getPackageInfo pair@(name, version) = do
        unlessM (doesFileExist fp) $ do
            _onDownload s pkg
            packageInfo <- getPackageInfo
            let (msha512, url, msize) =
                    case Map.lookup pair packageInfo of
                        Nothing -> (Nothing, defUrl, Nothing)
                        Just p ->
                            ( Map.lookup "SHA512" $ packageHashes p
                            , case reverse $ packageLocations p of
                                [] -> defUrl
                                x:_ -> T.unpack x
                            , packageSize p
                            )
            createDirectoryIfMissing True $ takeDirectory fp
#if MIN_VERSION_http_client(0,5,0)
            req' <- parseRequest url
#else
            req <- parseUrl url
            let req' = req
                    { checkStatus = \s x y ->
                        if statusCode s `elem` [401, 403]
                            
                            then Nothing
                            else checkStatus req s x y
                    }
#endif
            withResponse req' man $ \res -> if statusCode (responseStatus res) == 200
                then do
                    let tmp = fp <.> "tmp"
                    withBinaryFile tmp WriteMode $ \h -> do
                        let loop total ctx = do
                                bs <- brRead $ responseBody res
                                if S.null bs
                                    then
                                        case msize of
                                            Nothing -> return ()
                                            Just expected
                                                | expected /= total ->
                                                    throwIO InvalidDownloadSize
                                                        { _idsUrl = url
                                                        , _idsExpected = expected
                                                        , _idsTotalDownloaded = total
                                                        }
                                                | otherwise -> validHash url msha512 ctx
                                    else do
                                        S.hPut h bs
                                        let total' = total + fromIntegral (S.length bs)
                                        case msize of
                                            Just expected | expected < total' ->
                                                throwIO InvalidDownloadSize
                                                    { _idsUrl = url
                                                    , _idsExpected = expected
                                                    , _idsTotalDownloaded = total'
                                                    }
                                            _ -> loop total' $! hashUpdate ctx bs
                        loop 0 hashInit
                    renameFile tmp fp
                else _onDownloadErr s pkg
      where
        pkg = concat [name, "-", version]
        targz = pkg ++ ".tar.gz"
        defUrl = _downloadPrefix s ++ targz
        fp = packageLocation name version
validHash :: String -> Maybe Text -> Context SHA512 -> IO ()
validHash _ Nothing _ = return ()
validHash url (Just sha512) ctx
    | encodeUtf8 sha512 == digestToHexByteString dig = return ()
    | otherwise = throwIO InvalidHash
        { _ihUrl = url
        , _ihExpected = sha512
        , _ihActual = dig
        }
  where
    dig = hashFinalize ctx
parMapM_ :: F.Foldable f
         => Int
         -> (a -> IO ())
         -> f a
         -> IO ()
parMapM_ (max 1 -> 1) f xs = F.mapM_ f xs
parMapM_ cnt f xs0 = do
    var <- newTVarIO $ F.toList xs0
    let worker :: IO ()
        worker = fix $ \loop -> join $ atomically $ do
            xs <- readTVar var
            case xs of
                [] -> return $ return ()
                x:xs' -> do
                    writeTVar var xs'
                    return $ do
                        f x
                        loop
        workers 1 = Concurrently worker
        workers i = Concurrently worker *> workers (i  1)
    runConcurrently $ workers cnt