{-# LANGUAGE RecordWildCards #-}

module Generators where

import Crypto.Hash (
    hashDigestSize,
 )
import Crypto.Hash.Algorithms (
    SHA256 (SHA256),
 )
import qualified Data.ByteString as BS
import Data.ByteString.Base32 (encodeBase32Unpadded)
import qualified Data.ByteString.Lazy as LBS
import Data.Int (Int64)
import qualified Data.Text as T
import Hedgehog (MonadGen)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Tahoe.CHK.Crypto (storageIndexLength)
import Tahoe.CHK.Merkle (MerkleTree, makeTreePartial)
import Tahoe.CHK.Server (StorageServerAnnouncement (StorageServerAnnouncement))
import Tahoe.CHK.Share (Share (..))
import Tahoe.CHK.Types (Parameters (..), ShareNum, StorageIndex)
import Tahoe.CHK.URIExtension (URIExtension (URIExtension))

-- | The maximum value an Int64 can represent.
maxInt64 :: Integer
maxInt64 = fromIntegral (maxBound :: Int64)

-- | Generate Parameters values for which all field invariants hold.
genParameters :: MonadGen m => m Parameters
genParameters = do
    paramSegmentSize <- Gen.integral (Range.exponential 1 maxInt64)
    paramTotalShares <- Gen.integral (Range.linear 1 256)
    paramRequiredShares <- Gen.integral (Range.linear 1 paramTotalShares)
    -- XXX We're going to get rid of "Happy" from this type.  For now it's
    -- easier not to let this value vary and it doesn't hurt anything.
    let paramHappyShares = 1
    pure $ Parameters{paramSegmentSize, paramTotalShares, paramHappyShares, paramRequiredShares}

-- | Generate URIExtension values which are not necessarily well-formed.
genURIExtension :: MonadGen m => m URIExtension
genURIExtension =
    URIExtension
        <$> Gen.bytes (Range.linear 1 32)
        <*> genParameters
        <*> genParameters
        <*> Gen.integral (Range.exponential 1 maxInt64)
        <*> Gen.integral (Range.exponential 1 maxInt64)
        <*> Gen.integral (Range.exponential 1 (maxBound :: Int))
        <*> Gen.integral (Range.linear 1 256)
        <*> Gen.integral (Range.linear 1 256)
        <*> genHash
        <*> genHash
        <*> genHash

-- | Generate ByteStrings which could be sha256d digests.
genHash :: MonadGen m => m BS.ByteString
genHash = Gen.bytes . Range.singleton . hashDigestSize $ SHA256

shares :: MonadGen m => m Share
shares = do
    -- XXX It would be nice to explore the full space but the tests operate in
    -- memory (and even if they didn't, they would be constrained by disk
    -- space and speed) and maxBound :: Int64 is a lot of bytes...
    let maxSize = 65536
    shareBlockSize <- Gen.integral (Range.exponential 1 maxSize)
    numBlocks <- Gen.integral (Range.exponential 1 32)

    -- We don't make shareDataSize agree with the rest of the share data
    -- because the field is supposedly unused so everyone should just ignore
    -- it and not mind if we put garbage there.
    --
    -- We can go all the way up to an unreasonable maximum here because this
    -- doesn't influence how many bytes are actually in the share.
    shareDataSize <- fromIntegral <$> Gen.integral (Range.linear 1 maxInt64)

    shareBlocks <- Gen.list (Range.singleton numBlocks) (LBS.fromStrict <$> Gen.bytes (Range.singleton $ fromIntegral shareBlockSize))

    -- XXX These merkle trees and the "needed hashes" list all have a size
    -- that really needs to be dictated by the encoding parameters (k and n).
    sharePlaintextHashTree <- merkleTrees (Range.exponential 1 256)
    shareCrypttextHashTree <- merkleTrees (Range.exponential 1 256)
    shareBlockHashTree <- merkleTrees (Range.exponential 1 256)
    shareNeededHashes <- Gen.list (Range.exponential 1 100) ((,) <$> Gen.integral (Range.exponential 1 255) <*> Gen.bytes (Range.singleton 32))

    -- XXX A valid share will have a URI extension that agrees with some of
    -- the other fields we've just generated, which we're not even trying to
    -- do here.
    shareURIExtension <- genURIExtension

    pure $ Share{..}

merkleTrees :: MonadGen m => Range.Range Int -> m MerkleTree
merkleTrees r = makeTreePartial <$> Gen.list r genHash

storageIndexes :: MonadGen m => m StorageIndex
storageIndexes = Gen.bytes (Range.singleton storageIndexLength)

shareNumbers :: MonadGen m => m ShareNum
shareNumbers = Gen.integral Range.linearBounded

storageServerIdentifiers :: MonadGen m => m T.Text
storageServerIdentifiers =
    Gen.choice
        -- XXX Maybe more than alpha?
        [ Gen.text (Range.linear 1 64) Gen.alpha
        , encodeBase32Unpadded <$> Gen.bytes (Range.linear 1 64)
        ]

-- | Generate storage server anonymous storage service announcements.
storageServerAnnouncements :: MonadGen m => m StorageServerAnnouncement
storageServerAnnouncements =
    StorageServerAnnouncement
        <$> Gen.maybe storageServiceFURLs
        -- XXX Maybe more than alpha?
        <*> Gen.maybe (Gen.text (Range.linear 1 32) Gen.alpha)
        -- XXX 32 bytes?
        <*> Gen.maybe (Gen.bytes (Range.singleton 32))

{- | Generate text that could be a storage server fURL.  TODO: Represent fURLs
 _and NURLs_ in a structured way instead of with Text.
-}
storageServiceFURLs :: MonadGen m => m T.Text
storageServiceFURLs = do
    -- XXX 32 bytes?
    tubid <- encodeBase32Unpadded <$> Gen.bytes (Range.singleton 32)
    -- XXX 32 bytes?
    swissnum <- encodeBase32Unpadded <$> Gen.bytes (Range.singleton 32)
    let location = "@tcp:"
    pure $ "pb://" <> tubid <> location <> "/" <> swissnum