{- Copyright 2016 Joey Hess <id@joeyh.name>
 -
 - Licensed under the GNU AGPL version 3 or higher.
 -}

module Storage (module Storage, module Types.Storage) where

import Types
import Types.Storage
import Share
import Storage.Local
import Storage.Network
import Data.Monoid
import Data.Maybe
import System.FilePath
import Control.Monad
import qualified Data.Set as S

allStorageLocations :: Maybe LocalStorageDirectory -> IO StorageLocations
allStorageLocations d = do
	servers <- networkServers
	return $ StorageLocations $
		map networkStorage servers <> map (uploadQueue d) servers

localStorageLocations :: Maybe LocalStorageDirectory -> StorageLocations
localStorageLocations d = StorageLocations $
	map (localStorage (storageDir d) . ("local" </>) . show)
		[1..100 :: Int]

type UpdateProgress = IO ()

-- | Stores the shares amoung the storage locations. Each location
-- gets at most one share from each set.
--
-- TODO: Add shuffling and queueing/chaffing to prevent
-- correlation of related shares.
storeShares :: StorageLocations -> ShareIdents -> [S.Set Share] -> UpdateProgress -> IO StoreResult
storeShares (StorageLocations locs) allsis shares updateprogress = do
	(r, usedlocs) <- go allsis shares []
	_ <- mapM_ obscureShares usedlocs
	return r
  where
	go sis (s:rest) usedlocs = do
		let (is, sis') = nextShareIdents sis
		(r, usedlocs') <- storeset locs [] Nothing (zip (S.toList is) (S.toList s))
		case r of
			StoreSuccess -> go sis' rest (usedlocs ++ usedlocs')
			_ -> return (r, usedlocs ++ usedlocs')
	go _ [] usedlocs = return (StoreSuccess, usedlocs)

	storeset _ usedlocs _ [] = return (StoreSuccess, usedlocs)
	storeset [] usedlocs lasterr _ =
		return (fromMaybe (StoreFailure "no storage locations") lasterr, usedlocs)
	storeset (loc:otherlocs) usedlocs _ ((i, s):rest) = do
		r <- storeShare loc i s
		case r of
			StoreSuccess -> do
				_ <- updateprogress
				storeset otherlocs (loc:usedlocs) Nothing rest
			-- Give up if any location complains a share
			-- already exists, because we have a name conflict.
			StoreAlreadyExists -> return (StoreAlreadyExists, usedlocs)
			-- Try storing it somewhere else on failure.
			StoreFailure _ ->
				storeset otherlocs usedlocs (Just r) ((i, s):rest)

-- | Retrieves one set of shares from the storage locations.
-- Returns all the shares it can find, which may not be enough,
-- and the remaining Shareidents, to use to get subsequent sets.
--
-- Assumes that each location only contains one share. So, once a
-- share has been found on a location, can avoid asking that location
-- for any other shares.
retrieveShares :: StorageLocations -> ShareIdents -> UpdateProgress -> IO (S.Set Share, ShareIdents)
retrieveShares (StorageLocations locs) sis updateprogress = do
	let (is, sis') = nextShareIdents sis
	let want = zip [1..] (S.toList is)
	(shares, usedlocs, _unusedlocs) <- go locs [] want []
	_ <- mapM_ obscureShares usedlocs
	return (S.fromList shares, sis')
  where
	go unusedlocs usedlocs [] shares = return (shares, usedlocs, unusedlocs)
	go [] usedlocs _ shares = return (shares, usedlocs, [])
	go (loc:otherlocs) usedlocs ((n, i):rest) shares = do
		r <- retrieveShare loc n i
		case r of
			RetrieveSuccess s -> do
				_ <- updateprogress
				go otherlocs (loc:usedlocs) rest (s:shares)
			RetrieveFailure _ -> do
				-- Try to get the share from other locations.
				(shares', usedlocs', unusedlocs) <-
					go otherlocs usedlocs [(n, i)] shares
				-- May need to ask the location that didn't
				-- have the share for a later share, but
				-- ask it last. This way, the first
				-- location on the list can't deny having
				-- all shares and so learn the idents of
				-- all of them.
				go (unusedlocs++[loc]) usedlocs' rest shares'

uploadQueued :: Maybe LocalStorageDirectory -> IO ()
uploadQueued d = do
	servers <- networkServers
	forM_ servers $ \s -> moveShares (uploadQueue d s) (networkStorage s)