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

{-# LANGUAGE OverloadedStrings #-}

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

import Types
import Types.Storage
import Types.Server
import Share
import Storage.Local
import Storage.Network
import Servers
import Tunables
import Data.Maybe
import Data.List
import Data.Monoid
import System.IO
import System.FilePath
import Control.Monad
import Crypto.Random
import Control.Concurrent.Async
import qualified Data.Set as S
import Network.Wai.Handler.Warp (Port)

networkStorageLocations :: Maybe LocalStorageDirectory -> StorageLocations
networkStorageLocations d = StorageLocations $ 
	map (networkStorage d) networkServers

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.
--
-- If a server is not currently accessible, it will be queued locally.
-- If this happens at all, returns True.
--
-- TODO: Add shuffling and queueing/chaffing to prevent
-- correlation of related shares.
storeShares :: StorageLocations -> ShareIdents -> [S.Set Share] -> UpdateProgress -> IO (StoreResult, Bool, [Storage])
storeShares (StorageLocations locs) allsis shares updateprogress = do
	((r, anyqueued), usedlocs) <- go allsis shares [] False
	_ <- mapM_ obscureShares usedlocs
	return (r, anyqueued, usedlocs)
  where
	go sis (s:rest) usedlocs anyqueued = do
		let (is, sis') = nextShareIdents sis
		(r, usedlocs', queued) <- storeset locs [] Nothing (zip (S.toList is) (S.toList s)) False
		case r of
			StoreSuccess -> go sis' rest (usedlocs ++ usedlocs') (anyqueued || queued)
			_ -> return ((r, anyqueued || queued), usedlocs ++ usedlocs')
	go _ [] usedlocs anyqueued = return ((StoreSuccess, anyqueued), usedlocs)

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

-- | 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, [Server])
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
	let usedservers = mapMaybe getServer usedlocs
	return (S.fromList shares, sis', usedservers)
  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'

-- | Returns descriptions of any failures.
tryUploadQueued :: Maybe LocalStorageDirectory -> IO [String]
tryUploadQueued d = do
	results <- forM locs $ \loc -> case uploadQueue loc of
		Nothing -> return []
		Just q -> moveShares q loc
	return $ processresults (concat results) []
  where
	StorageLocations locs = networkStorageLocations d
	processresults [] c = nub c
	processresults (StoreSuccess:rs) c = processresults rs c
	processresults (StoreFailure e:rs) c = processresults rs (e:c)
	processresults (StoreAlreadyExists:rs) c =
		processresults rs ("Unable to upload a share to a server due to a name conflict.":c)

storeChaff :: HostName -> Port -> IO ()
storeChaff hn port = forever $ do
	putStrLn $ "Sending chaff to " ++ hn ++ " (press ctrl-c to stop)"
	putStrLn "Legend: + = successful upload, ! = upload failure"
	rng <- (cprgCreate <$> createEntropyPool) :: IO SystemRNG
	let (randomname, rng') = cprgGenerate 128 rng
	-- It's ok the use the testModeTunables here because
	-- the randomname is not something that can be feasibly guessed.
	-- Prefix "random chaff" to the name to avoid ever using a name
	-- that a real user might want to use.
	let sis = shareIdents testModeTunables (Name $ "random chaff:" <> randomname) (KeyFile "random")
	mapConcurrently (go sis rng')
		[1..totalObjects (shareParams testModeTunables)]
  where
	server = networkStorage Nothing $ Server (ServerName hn) 
		[ServerAddress hn port]
	objsize = objectSize defaultTunables * shareOverhead defaultTunables
	go sis rng n = do
		let (b, rng') = cprgGenerate objsize rng
		let share = Share 0 (StorableObject b)
		let (is, sis') = nextShareIdents sis
		let i = S.toList is !! (n - 1)
		r <- storeShare server i share
		case r of
			StoreSuccess -> putStr "+"
			_ -> putStr "!"
		hFlush stdout
		go sis' rng' n