{-# LANGUAGE TypeFamilies #-}

{- | A high-level interface to downloading share data as bytes from storage
 servers.
-}
module Tahoe.Download (
    LookupServer,
    DownloadError (..),
    DirectoryDownloadError (..),
    LookupError (..),
    DiscoverError (..),
    discoverShares,
    download,
    downloadDirectory,
    announcementToImmutableStorageServer,
    announcementToMutableStorageServer,
    getShareNumbers,
) where

import Control.Concurrent.Async (mapConcurrently)
import Control.Exception (Exception (displayException), SomeException, try)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Bifunctor (Bifunctor (first, second))
import Data.Binary (Word16)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import Data.Either (partitionEithers, rights)
import Data.List (foldl')
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Tahoe.Announcement (StorageServerAnnouncement)
import Tahoe.CHK.Server (StorageServer (..), StorageServerID)
import Tahoe.CHK.Types (ShareNum, StorageIndex)
import Tahoe.Directory (Directory, DirectoryCapability (DirectoryCapability))
import qualified Tahoe.Directory as Directory
import Tahoe.Download.Internal.Capability
import Tahoe.Download.Internal.Client
import Tahoe.Download.Internal.Immutable
import Tahoe.Download.Internal.Mutable

-- | Partially describe one share download.
type DownloadTask = (ShareNum, StorageServer)

-- | A downloaded share
type DownloadedShare = (ShareNum, LB.ByteString)

{- | Recover the application data associated with a given capability from the
 given servers, if possible.
-}
download ::
    -- To download, we require a capability for which there is a Readable
    -- instance because are also going to decrypt the ciphertext.  A different
    -- download interface that skips decryption could settle for a capability
    -- with a Verifiable instance.  We also require that the Verifier type for
    -- the read capability has a Verifiable instance because Verifiable is
    -- what gives us the ability to locate the shares.  If we located
    -- separately from decrypting this might be simpler.
    (MonadIO m, Readable readCap, Verifiable v, Verifier readCap ~ v) =>
    -- | Information about the servers from which to consider downloading shares
    -- representing the application data.
    Map.Map StorageServerID StorageServerAnnouncement ->
    -- | The read capability for the application data.
    readCap ->
    -- | Get functions for interacting with a server given its URL.
    LookupServer IO ->
    -- | Either a description of how the recovery failed or the recovered
    -- application data.
    m (Either DownloadError LB.ByteString)
download servers cap lookupServer = do
    print' ("Downloading: " <> show (getStorageIndex $ getVerifiable cap))
    let verifier = getVerifiable cap
    let storageIndex = getStorageIndex verifier
    -- TODO: If getRequiredTotal fails on the first storage server, we may
    -- need to try more.  If it fails for all of them, we need to represent
    -- the failure coherently.
    someParam <- liftIO $ firstRightM lookupServer (getRequiredTotal verifier) (Map.elems servers)
    case someParam of
        Left errs -> pure . Left $ if servers == mempty then NoConfiguredServers else NoReachableServers (StorageServerUnreachable <$> errs)
        Right (required, _) -> do
            locationE <- liftIO $ locateShares servers lookupServer storageIndex (fromIntegral required)
            print' "Finished locating shares"
            case locationE of
                Left err -> do
                    print' "Got an error locating shares"
                    pure $ Left err
                Right discovered -> do
                    print' "Found some shares, fetching them"
                    -- XXX note shares can contain failures
                    shares <- liftIO $ executeDownloadTasks storageIndex (makeDownloadTasks =<< discovered)
                    print' "Fetched the shares, decoding them"
                    s <- liftIO $ decodeShares cap shares required
                    print' "Decoded them"
                    pure s

{- | Apply a monadic operation to each element of a list and another monadic
 operation values in the resulting Rights.  If all of the results are Lefts or
 Nothings, return a list of the values in the Lefts.  Otherwise, return the
 *first* Right.
-}
firstRightM :: Monad m => (a -> m (Either b c)) -> (c -> m (Maybe d)) -> [a] -> m (Either [b] d)
firstRightM _ _ [] = pure $ Left []
firstRightM f op (x : xs) = do
    s <- f x
    case s of
        Left bs -> first (bs :) <$> recurse
        Right ss -> do
            r <- op ss
            case r of
                Nothing -> recurse
                Just d -> pure $ Right d
  where
    recurse = firstRightM f op xs

{- | Execute each download task sequentially and return only the successful
 results.
-}
executeDownloadTasks ::
    -- | The storage index of the shares to download.
    StorageIndex ->
    -- | The downloads to attempt.
    [DownloadTask] ->
    -- | The results of all successful downloads.
    IO [DownloadedShare]
executeDownloadTasks storageIndex tasks = do
    downloadResults <- mapConcurrently (downloadShare storageIndex) tasks
    pure . rights $ inject <$> downloadResults
  where
    inject (a, b) = (a,) <$> b

-- | Find out which servers claim to have shares related to a given storage index.
locateShares ::
    -- | Information about the servers from which to consider downloading shares
    -- representing the application data.
    Map.Map StorageServerID StorageServerAnnouncement ->
    -- | Get functions for interacting with a server given its URL.
    LookupServer IO ->
    -- | The storage index about which to retrieve information.
    B.ByteString ->
    -- | The number of shares we need to locate.  If we cannot find at least
    -- this many shares the result will be an error.
    Word16 ->
    -- | Either an error or a guide to where shares are placed.
    IO (Either DownloadError [(StorageServer, Set.Set ShareNum)])
locateShares servers lookupServer storageIndex required =
    case Map.toList servers of
        [] -> pure . Left $ NoConfiguredServers
        serverList -> do
            print' "Discovering shares"
            -- Ask each server for all shares it has.
            ( problems :: [DiscoverError]
                , discovered :: [(StorageServer, Set.Set ShareNum)]
                ) <-
                partitionEithers <$> mapConcurrently (discoverShares lookupServer storageIndex) serverList
            if null discovered
                then pure . Left . NoReachableServers $ problems
                else
                    if (fromIntegral required >) . countDistinctShares $ discovered
                        then pure $ Left NotEnoughShares{notEnoughSharesNeeded = fromIntegral required, notEnoughSharesFound = countDistinctShares discovered}
                        else pure $ Right discovered

{- | Given the results of downloading shares related to a given capability,
 decode them and decrypt the contents of possible.
-}
decodeShares ::
    (Readable readCap, Verifiable v, v ~ Verifier readCap) =>
    -- | The read capability which allows the contents to be decrypted.
    readCap ->
    -- | The results of downloading the shares.
    [DownloadedShare] ->
    Int ->
    IO (Either DownloadError LB.ByteString)
decodeShares r shares required = do
    -- Filter down to shares we actually got.
    let fewerShares = second (deserializeShare (getVerifiable r)) <$> shares
        onlyDecoded = rights $ (\(a, b) -> (fromIntegral a,) <$> b) <$> fewerShares
    if length onlyDecoded < required
        then pure $ Left NotEnoughDecodedShares{notEnoughDecodedSharesNeeded = fromIntegral required, notEnoughDecodedSharesFound = length onlyDecoded}
        else do
            decodeShare r onlyDecoded

{- | Figure the total number of distinct shares reported by all of the servers
 we asked.
-}
countDistinctShares :: Ord b => [(a, Set.Set b)] -> Int
countDistinctShares = Set.size . foldl' Set.union mempty . map snd

{- | Ask one server which shares it has related to the storage index in
 question.
-}
discoverShares ::
    LookupServer IO ->
    StorageIndex ->
    (StorageServerID, StorageServerAnnouncement) ->
    IO (Either DiscoverError (StorageServer, Set.Set ShareNum))
discoverShares lookupServer storageIndex (_sid, sann) = do
    print' "Looking up server from announcement"
    server <- lookupServer sann
    print' "Looked it up"
    case server of
        Left e -> pure . Left . StorageServerUnreachable $ e
        Right ss@StorageServer{storageServerGetBuckets} -> do
            print' $ "Getting buckets for " <> show storageIndex
            buckets <- try (storageServerGetBuckets storageIndex)
            let massaged = first (StorageServerCommunicationError . (displayException :: SomeException -> String)) buckets
            print' $ "Got them " <> show massaged
            pure $ (ss,) <$> massaged

{- | Expand a one-to-many mapping into a list of pairs with each of the "many"
   values as the first element and the corresponding "one" value as the second
   element.
-}
makeDownloadTasks :: Ord k => (v, Set.Set k) -> [(k, v)]
makeDownloadTasks (v, ks) = zip (Set.toList ks) (repeat v)

-- | Download the bytes of a share from one (or more!) of the given servers.
downloadShare ::
    -- | The storage index of the share to download.
    StorageIndex ->
    -- | Addressing information about the share to download.
    DownloadTask ->
    -- | The bytes of the share or some error that was encountered during
    -- download.
    IO (ShareNum, Either DownloadError LB.ByteString)
downloadShare storageIndex (shareNum, s) = do
    print' $ "Going to download " <> show storageIndex <> " " <> show shareNum
    shareBytes <- try (storageServerRead s storageIndex shareNum)
    let massaged = first (ShareDownloadError . (displayException :: SomeException -> String)) shareBytes
    print' "Downloaded it"
    pure (shareNum, LB.fromStrict <$> massaged)

{- | Download the data associated with a directory capability and interpret it
 as a collection of entries.
-}
downloadDirectory ::
    (MonadIO m, Readable readCap, Verifiable v, Verifier readCap ~ v) =>
    -- | Information about the servers from which to consider downloading shares
    -- representing the application data.
    Map.Map StorageServerID StorageServerAnnouncement ->
    -- | The read capability for the application data.
    DirectoryCapability readCap ->
    -- | Get functions for interacting with a server given its URL.
    LookupServer IO ->
    -- | Either a description of how the recovery failed or the recovered
    -- application data.
    m (Either DirectoryDownloadError Directory)
downloadDirectory anns (DirectoryCapability cap) lookupServer = do
    bs <- download anns cap lookupServer
    pure $ do
        bs' <- first UnderlyingDownloadError bs
        first (const DecodingError) $ Directory.parse (LB.toStrict bs')

data DirectoryDownloadError
    = UnderlyingDownloadError DownloadError
    | DecodingError
    deriving (Ord, Eq, Show)