{-# LANGUAGE TypeFamilies #-}

module Tahoe.Download.Internal.Capability where

import Control.Exception (SomeException, throwIO, try)
import Control.Monad.IO.Class
import Data.Bifunctor (Bifunctor (..))
import Data.Binary (Word8, decodeOrFail)
import Data.Binary.Get (ByteOffset)
import qualified Data.ByteString.Lazy as LB
import Data.Foldable (foldlM)
import qualified Data.Set as Set
import Network.HTTP.Types (Status (statusCode))
import Servant.Client (ClientError (FailureResponse), ResponseF (..))
import qualified Tahoe.CHK
import qualified Tahoe.CHK.Capability as CHK
import qualified Tahoe.CHK.Encrypt
import Tahoe.CHK.Server
import qualified Tahoe.CHK.Share
import Tahoe.CHK.Types
import Tahoe.Download.Internal.Client
import qualified Tahoe.SDMF as SDMF
import qualified Tahoe.SDMF.Internal.Keys as SDMF.Keys

-- | A capability which confers the ability to locate and verify some stored data.
class Verifiable v where
    -- | Represent the type of share to operate on.
    type ShareT v

    -- | Ask a storage server which share numbers related to this capability it
    -- is holding.  This is an unverified result and the storage server could
    -- present incorrect information.  Even if it correctly reports that it
    -- holds a share, it could decline to give it out when asked.
    getShareNumbers :: MonadIO m => v -> StorageServer -> m (Set.Set ShareNum)

    -- | Get the encoding parameters used for the shares of this capability.
    -- The information is presented as a tuple of (required, total).

    -- SDMF can fail to figure this out in lots of ways so consider switching
    -- to Either or something?
    getRequiredTotal :: MonadIO m => v -> StorageServer -> m (Maybe (Int, Int))

    -- | Get the location information for shares of this capability.
    getStorageIndex :: v -> StorageIndex

    -- | Deserialize some bytes representing some kind of share to the kind of
    -- share associated with this capability type, if possible.
    deserializeShare ::
        -- | A type witness revealing what type of share to decode to.
        v ->
        -- | The bytes of the serialized share.
        LB.ByteString ->
        Either (LB.ByteString, ByteOffset, String) (ShareT v)

{- | A capability which confers the ability to recover plaintext from
 ciphertext.
-}
class Readable r where
    -- | Represent the type of a Verifiable associated with the Readable.
    type Verifier r

    -- | Attentuate the capability.
    getVerifiable :: r -> Verifier r

    -- | Interpret the required number of shares to recover the plaintext.
    --
    -- Note: might want to split the two functions below out of decodeShare
    --
    -- shareToCipherText :: r -> [(Int, ShareT r)] -> LB.ByteString
    --
    -- cipherTextToPlainText :: r -> LB.ByteString -> LB.ByteString
    decodeShare :: MonadIO m => r -> [(Int, ShareT (Verifier r))] -> m (Either DownloadError LB.ByteString)

instance Verifiable CHK.Verifier where
    type ShareT CHK.Verifier = Tahoe.CHK.Share.Share

    getShareNumbers v s = liftIO $ storageServerGetBuckets s (CHK.storageIndex v)
    getStorageIndex CHK.Verifier{storageIndex} = storageIndex

    -- CHK is pure, we don't have to ask the StorageServer
    getRequiredTotal CHK.Verifier{required, total} _ = pure $ pure (fromIntegral required, fromIntegral total)

    deserializeShare _ = fmap (\(_, _, c) -> c) . decodeOrFail

{- | A capability which confers the ability to interpret some stored data to
 recover the original plaintext.  Additionally, it can be attentuated to a
 Verifiable.
-}
instance Readable CHK.Reader where
    type Verifier CHK.Reader = CHK.Verifier

    getVerifiable = CHK.verifier
    decodeShare r shareList = do
        cipherText <- liftIO $ Tahoe.CHK.decode r shareList
        case cipherText of
            Nothing -> pure $ Left ShareDecodingFailed
            Just ct ->
                pure . Right $ Tahoe.CHK.Encrypt.decrypt (CHK.readKey r) ct

firstJustsM :: (Monad m, Foldable f) => f (m (Maybe a)) -> m (Maybe a)
firstJustsM = foldlM go Nothing
  where
    go :: Monad m => Maybe a -> m (Maybe a) -> m (Maybe a)
    go Nothing action = action
    go result@(Just _) _action = return result

instance Verifiable SDMF.Verifier where
    type ShareT SDMF.Verifier = SDMF.Share

    getShareNumbers v s = liftIO $ storageServerGetBuckets s (SDMF.Keys.unStorageIndex $ SDMF.verifierStorageIndex v)
    getStorageIndex = SDMF.Keys.unStorageIndex . SDMF.verifierStorageIndex
    getRequiredTotal SDMF.Verifier{..} ss = liftIO $ do
        -- Find out what shares it has.  Any share will do but we need to tell
        -- it which we want.
        errorOrShareNums <- try $ storageServerGetBuckets ss storageIndex
        case Set.toList <$> errorOrShareNums of
            -- Literally anything could go wrong with that...
            Left (e :: SomeException) -> throwIO e
            -- Or the server may have no shares for this storage index.
            Right [] -> pure Nothing
            -- Or it might have at least one.  Check each in turn, stopping as
            -- soon as we get a result.
            Right shareNums -> firstJustsM (getParams <$> shareNums)
      where
        -- Get the Required, Total parameters for one share number, if
        -- possible.
        getParams :: MonadIO m => Word8 -> m (Maybe (Int, Int))
        getParams shareNum = liftIO $ do
            errorOrShareBytes <- try $ storageServerRead ss storageIndex shareNum
            case errorOrShareBytes of
                Left e@(FailureResponse _ response) ->
                    -- It should not be very surprising for the requested share to
                    -- be missing from the server (you can never be sure what a
                    -- server will have).  Other issues should probably be kept
                    -- visible.
                    if isStatusCode 404 response
                        then pure Nothing
                        else throwIO e
                Left e -> throwIO e
                Right shareBytes ->
                    case decodeOrFail (LB.fromStrict shareBytes) of
                        Left _ -> pure Nothing
                        Right (_, _, sh) -> pure $ pure (fromIntegral $ SDMF.shareRequiredShares sh, fromIntegral $ SDMF.shareTotalShares sh)

        storageIndex = SDMF.Keys.unStorageIndex verifierStorageIndex

    deserializeShare _ = fmap (\(_, _, c) -> c) . decodeOrFail

-- | Test the status code of a response for equality against a given value.
isStatusCode :: Int -> ResponseF a -> Bool
isStatusCode expected = (expected ==) . statusCode . responseStatusCode

instance Readable SDMF.Reader where
    type Verifier SDMF.Reader = SDMF.Verifier
    getVerifiable = SDMF.readerVerifier
    decodeShare r shareList = do
        cipherText <- Right <$> liftIO (SDMF.decode r (first fromIntegral <$> shareList))
        case cipherText of
            Left _ -> pure $ Left ShareDecodingFailed
            Right ct -> do
                print' ("Got some ciphertext: " <> show ct)
                print' ("Decrypting with iv: " <> show iv)
                pure . Right $ SDMF.decrypt readKey iv ct
              where
                readKey = SDMF.readerReadKey r
                iv = SDMF.shareIV (snd . head $ shareList)

print' :: MonadIO m => String -> m ()
-- print' = liftIO . putStrLn
print' = const $ pure ()