{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module Tahoe.CHK.URIExtension (
    URIExtension (..),
    uriExtensionToBytes,
    showBytes,
    pURIExtension,
) where

import Control.Applicative.Combinators (count)
import Control.Applicative.Permutations (runPermutation, toPermutation)
import Control.Monad (join, void)
import Data.TreeDiff.Class (ToExpr)
import Data.Void (Void)
import GHC.Generics (Generic)

import Text.Megaparsec (
    MonadParsec (takeP),
    Parsec,
    anySingle,
 )
import Text.Megaparsec.Byte (string)
import Text.Megaparsec.Byte.Lexer (decimal)

import Data.ByteString.Base32 (
    encodeBase32Unpadded,
 )

import qualified Data.ByteString as B
import qualified Data.Text as T
import Data.Text.Encoding (
    decodeLatin1,
    encodeUtf8,
 )

import Data.List (
    sort,
 )

import qualified Tahoe.CHK.Parsing
import Tahoe.CHK.Types (
    CrypttextHash,
    Parameters (..),
    Required,
    SegmentNum,
    Size,
    Total,
 )

import Tahoe.Netstring (
    netstring,
 )

-- | Represent additional metadata that appears at the end of each share.
data URIExtension = URIExtension
    { -- | The name of the encoding function.  Only "zfec" is implemented.
      uriExtCodecName :: B.ByteString
    , -- | The parameters for the encoding function for all except the final
      -- segment.
      uriExtCodecParams :: Parameters
    , -- | The parameters for the encoding function for the final segment.
      uriExtTailCodecParams :: Parameters
    , -- | The application data size in bytes.
      uriExtSize :: Size
    , -- | The individual segment size in bytes.
      uriExtSegmentSize :: Size
    , -- | The number of segments of application data.  Note the last segment
      -- may be short so it is not necessarily the case that uriExtSize ==
      -- uriExtSegmentSize * uriExtNumSegments.
      uriExtNumSegments :: SegmentNum
    , -- | The required (K) parameter to the encoding function.  This is a
      -- duplicate of the values in uriExtCodecParams and uriExtTailCodecParams.
      uriExtNeededShares :: Required
    , -- | The total (N) parameter to the encoding function.  This too is a
      -- duplicate.
      uriExtTotalShares :: Total
    , -- | A tagged sha256d hash of the complete ciphertext.
      uriExtCrypttextHash :: CrypttextHash
    , -- | The root hash of a merkle tree where the leaf hashes are of segments of ciphertext.
      uriExtCrypttextRootHash :: CrypttextHash
    , -- | The root hash of a merkle tree where leaf hashes are the root hashes of all of the block hash trees.
      uriExtShareRootHash :: CrypttextHash
    }
    deriving (Eq, Ord, Generic, ToExpr)

instance Show URIExtension where
    show (URIExtension name params tailParams size segSize numSegs needed total hash1 hash2 hash3) =
        T.unpack . T.concat $
            [ "URIExtension { "
            , "codec = "
            , decodeLatin1 name
            , "; codec-params = "
            , showText params
            , "; tail-codec-params = "
            , showText tailParams
            , "; size = "
            , showText size
            , "; segment-size = "
            , showText segSize
            , "; num-segments = "
            , showText numSegs
            , "; needed-shares = "
            , showText needed
            , "; total-shares = "
            , showText total
            , "; crypttext-hash = "
            , showText $ b32 hash1
            , "; crypttext-root-hash = "
            , showText $ b32 hash2
            , "; share-root-hash = "
            , showText $ b32 hash3
            , " }"
            ]
      where
        showText :: Show s => s -> T.Text
        showText = T.pack . show
        b32 = encodeBase32Unpadded

-- Serialize a URIExtension to bytes in the format it appears in a CHK share.
uriExtensionToBytes :: URIExtension -> B.ByteString
uriExtensionToBytes =
    toWeirdString
        -- all of the below values are authenticated by the capability you get when you store data in Tahoe
        [ ("codec_name", uriExtCodecName)
        , ("codec_params", paramsToBytes . uriExtCodecParams)
        , ("tail_codec_params", paramsToBytes . uriExtTailCodecParams)
        , ("size", showBytes . uriExtSize)
        , ("segment_size", showBytes . uriExtSegmentSize)
        , ("num_segments", showBytes . uriExtNumSegments)
        , ("needed_shares", showBytes . uriExtNeededShares)
        , ("total_shares", showBytes . uriExtTotalShares)
        , ("crypttext_hash", uriExtCrypttextHash) -- hash of the *entire* cipher text
        , ("crypttext_root_hash", uriExtCrypttextRootHash) -- root hash of the *cipher text* merkle tree
        , ("share_root_hash", uriExtShareRootHash) -- root hash of the *share* merkle tree
        ]

type Parser = Parsec Void B.ByteString

-- | A version of bounded specialized to parsing bytestrings.
bounded :: (Ord n, Integral n) => n -> n -> Parser n
bounded = Tahoe.CHK.Parsing.bounded decimal

{- | Parse the representation of a URIExtension which appears in CHK shares
 back into a URIExtension.
-}
pURIExtension :: Parser URIExtension
pURIExtension =
    runPermutation $
        URIExtension
            <$> toPermutation (B.pack <$> pField "codec_name" (`count` anySingle))
            <*> toPermutation (pField "codec_params" $ const pParameters)
            <*> toPermutation (pField "tail_codec_params" $ const pParameters)
            <*> toPermutation (pField "size" $ const decimal)
            <*> toPermutation (pField "segment_size" $ const decimal)
            <*> toPermutation (pField "num_segments" $ const (bounded 1 maxBound))
            <*> toPermutation (pField "needed_shares" $ const (bounded 1 256))
            <*> toPermutation (pField "total_shares" $ const (bounded 1 256))
            <*> toPermutation (pField "crypttext_hash" $ takeP Nothing)
            <*> toPermutation (pField "crypttext_root_hash" $ takeP Nothing)
            <*> toPermutation (pField "share_root_hash" $ takeP Nothing)

-- | Parse one field of a serialized URIExtension.
pField ::
    -- | The serialized label for the field.
    B.ByteString ->
    -- | A function that takes the length of the field value and returns a parser for the field value.
    (Int -> Parser a) ->
    -- | A parser for the field.
    Parser a
pField label pInner = do
    void $ string (label <> ":")
    len <- decimal -- XXX Could overflow
    void $ string ":"
    result <- pInner len
    void $ string ","
    pure result

-- | Serialize some named URIExtension fields to bytes.
toWeirdString ::
    -- | A list of pairs of field names and functions to get serialized
    -- field values.
    [(B.ByteString, URIExtension -> B.ByteString)] ->
    -- | The URIExtension to get the field values from.
    URIExtension ->
    -- | The concatenation of all of the serialized fields.
    B.ByteString
toWeirdString fields ext =
    B.concat . join . sort $ map (encodedField ext) fields
  where
    encodedField ext' (name, extract) =
        [name, ":", netstring (extract ext')]

-- | Show a value as a UTF-8-encoded byte string.
showBytes :: (Show s) => s -> B.ByteString
showBytes = encodeUtf8 . T.pack . show

{- | Serialize Parameters to a byte string in the format it appears within the
 URI extension block in a CHK share.
-}
paramsToBytes :: Parameters -> B.ByteString
paramsToBytes Parameters{paramSegmentSize, paramTotalShares, paramRequiredShares} =
    B.concat [showBytes paramSegmentSize, "-", showBytes paramRequiredShares, "-", showBytes paramTotalShares]

{- | Parse a serialized Parameters value in the format produced by
 paramsToBytes.
-}
pParameters :: Parser Parameters
pParameters =
    (\segSize required total -> Parameters{paramSegmentSize = segSize, paramRequiredShares = required, paramHappyShares = 1, paramTotalShares = total})
        <$> decimal
        <* string "-"
        <*> bounded 1 maxShares
        <* string "-"
        <*> bounded 1 maxShares
  where
    maxShares = 256