module Generators where

import qualified Crypto.PubKey.RSA as RSA
import Data.ASN1.BinaryEncoding (DER (DER))
import Data.ASN1.Encoding (ASN1Decoding (decodeASN1))
import Data.ASN1.Types (ASN1Object (fromASN1))
import Data.Bifunctor (Bifunctor (first))
import Data.ByteString.Base32 (encodeBase32Unpadded)
import qualified Data.ByteString.Lazy as LB
import Data.Int (Int64)
import qualified Data.Text as T
import Data.X509 (PrivKey (PrivKeyRSA))
import Hedgehog (MonadGen)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import System.IO.Unsafe (unsafePerformIO)
import Tahoe.Announcement (Announcements (..), StorageServerAnnouncement (..))
import Tahoe.CHK.Types (Parameters (..))

-- | 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 2 256)
    paramRequiredShares <- Gen.integral (Range.linear 1 (paramTotalShares - 1))
    -- 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}

genAnnouncements :: MonadGen m => Range.Range Int -> m Announcements
genAnnouncements size =
    Announcements <$> Gen.map size ((,) <$> genServerIDs <*> genStorageServerAnnouncements)

genServerIDs :: MonadGen m => m T.Text
genServerIDs = T.toLower . encodeBase32Unpadded <$> Gen.bytes (Range.singleton 32)

genStorageServerAnnouncements :: MonadGen m => m StorageServerAnnouncement
genStorageServerAnnouncements =
    StorageServerAnnouncement
        <$> Gen.maybe (Gen.text (Range.linear 16 32) Gen.ascii)
        <*> Gen.maybe (Gen.text (Range.linear 16 32) Gen.ascii)
        <*> Gen.maybe (Gen.bytes $ Range.singleton 32)

{- | Build RSA key pairs.

 Because the specific bits of the key pair shouldn't make any difference to
 any application logic, generating new RSA key pairs is expensive, and
 generating new RSA key pairs in a way that makes sense in Hedgehog is
 challenging, this implementation just knows a few RSA key pairs already and
 will give back one of them.
-}
genRSAKeys :: MonadGen m => m RSA.PrivateKey
genRSAKeys = Gen.element (map rsaKeyPair rsaKeyPairBytes)

-- I'm not sure how to do IO in MonadGen so do the IO up front unsafely (but
-- hopefully not really unsafely).
rsaKeyPairBytes :: [LB.ByteString]
{-# NOINLINE rsaKeyPairBytes #-}
rsaKeyPairBytes = unsafePerformIO $ mapM (\n -> LB.readFile ("test/data/rsa-privkey-" <> show n <> ".der")) [0 .. 4 :: Int]

rsaKeyPair :: LB.ByteString -> RSA.PrivateKey
rsaKeyPair bs = do
    let (Right kp) = do
            asn1s <- first show (decodeASN1 DER bs)
            (r, _) <- fromASN1 asn1s
            case r of
                PrivKeyRSA pk -> pure pk
                _ -> error "Expected RSA Private Key"
    kp