{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Vectors where

import Control.Exception (IOException, try)
import Data.Aeson (
    FromJSON (..),
    ToJSON (..),
    Value (..),
    withText,
    (.:),
 )
import Data.Aeson.Types (
    parseFail,
    withObject,
 )
import qualified Data.Binary as Binary
import qualified Data.ByteString as B
import Data.ByteString.Base32 (decodeBase32Unpadded)
import qualified Data.ByteString.Base64 as Base64
import qualified Data.ByteString.Lazy as BL
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import Data.Word (Word32)
import Data.Yaml (ParseException, decodeEither')
import GHC.Generics (
    Generic,
 )
import Tahoe.CHK.Capability (Reader, pReader)
import Tahoe.CHK.Types (
    Parameters (..),
 )
import qualified Text.Megaparsec as M

newtype JSONByteString = JSONByteString B.ByteString deriving newtype (Ord, Eq)

instance Show JSONByteString where
    show (JSONByteString bs) = T.unpack . T.decodeLatin1 . Base64.encode $ bs

instance FromJSON JSONByteString where
    parseJSON =
        withText
            "base64 encoded bytestring"
            ( \t ->
                case Base64.decode . T.encodeUtf8 $ t of
                    Left err -> parseFail $ "parsing base64-encoded byte string failed" <> show err
                    Right stuff -> pure $ JSONByteString stuff
            )

instance ToJSON JSONByteString where
    toJSON (JSONByteString bs) = String . T.decodeLatin1 . Base64.encode $ bs

data SSKFormat
    = SDMF
        { sskPrivateKey :: T.Text
        }
    | MDMF
        { sskPrivateKey :: T.Text
        }
    deriving (Show, Ord, Eq)

data Format = CHK | SSK SSKFormat deriving (Show, Ord, Eq)

instance FromJSON Format where
    parseJSON = withObject "format" $ \o -> do
        kind <- o .: "kind"
        case kind of
            "chk" -> pure CHK
            "ssk" -> SSK <$> o .: "params"
            invalid -> parseFail $ "Unsupported format: " <> T.unpack invalid

instance FromJSON SSKFormat where
    parseJSON = withObject "ssk-format" $ \o -> do
        format <- o .: "format"
        key <- o .: "key"
        sskFormat <- case format of
            "sdmf" -> pure SDMF
            "mdmf" -> pure MDMF
            invalid -> parseFail $ "Unsupported SSK format: " <> T.unpack invalid
        pure $ sskFormat key

data Sample = Sample
    { sampleTemplate :: JSONByteString
    , sampleLength :: Int
    }
    deriving (Show, Ord, Eq)

instance FromJSON Sample where
    parseJSON = withObject "sample" $ \o ->
        Sample <$> o .: "seed" <*> o .: "length"

data VectorSpec = VectorSpec
    { version :: T.Text
    , vector :: [TestCase]
    }
    deriving (Generic, Show, Ord, Eq, FromJSON)

data TestCase = TestCase
    { convergence :: JSONByteString
    , format :: Format
    , sample :: Sample
    , zfec :: Parameters
    , expected :: T.Text
    }
    deriving (Generic, Show, Ord, Eq, FromJSON)

instance FromJSON Parameters where
    parseJSON = withObject "parameters" $ \o ->
        Parameters <$> o .: "segmentSize" <*> o .: "total" <*> pure 1 <*> o .: "required"

data LoadError = IOError IOException | ParseError ParseException deriving (Show)

-- | Load the test vectors from the yaml file.
loadTestVectorData :: IO (Either LoadError VectorSpec)
loadTestVectorData = go "test_vectors.yaml"
  where
    go :: String -> IO (Either LoadError VectorSpec)
    go path = do
        bs <- read' path
        pure $
            case bs of
                Left le -> Left . IOError $ le
                Right bs' -> parse bs'

    read' :: String -> IO (Either IOError B.ByteString)
    read' = try . B.readFile

    parse :: B.ByteString -> Either LoadError VectorSpec
    parse = either (Left . ParseError) pure . decodeEither'

data WellKnown = WellKnown
    { wellKnownPlaintext :: BL.ByteString
    , wellKnownConvergenceSecret :: B.ByteString
    , wellKnownParameters :: Parameters
    , wellKnownShares :: [BL.ByteString]
    , wellKnownCapability :: Reader
    }

{- | Load one "well-known" (that is: hard-coded and generated by Tahoe-LAFS)
 cases of CHK shares.
-}
loadWellKnownCase :: Parameters -> T.Text -> IO WellKnown
loadWellKnownCase wellKnownParameters@Parameters{paramRequiredShares, paramTotalShares} capText =
    do
        wellKnownShares <- mapM (fmap extractShareData . BL.readFile . sharePath) [0 .. paramTotalShares - 1]
        pure WellKnown{..}
  where
    -- Must be at least 56 bytes or we cannot get shares for comparison out of
    -- Tahoe-LAFS (instead, it emits a LIT cap).
    wellKnownPlaintext =
        "abcdefghijklmnopqrstuvwxyz\
        \ZYXWVUTSRQPONMLKJIJHGRFCBA\
        \1357"

    -- Hard-code the particular convergence secret used to generated the
    -- expected value.
    Right wellKnownConvergenceSecret = decodeBase32Unpadded "lcngfrvgaksfwrelc6ae5kucb3zufssoe6cj74rozcqibnl6uy2a"

    Right wellKnownCapability = M.parse pReader "" capText

    -- Find the share data in the source tree.
    sharePath n = concat ["test/", show paramRequiredShares, "of", show paramTotalShares, "." <> show n]

    -- Strip the share "container" bytes that is wrapped around the share data
    -- "proper" in our test data files.  These are the server-side bookkeeping
    -- bytes written by the Tahoe-LAFS storage server which generated our test
    -- data.
    extractShareData :: BL.ByteString -> BL.ByteString
    extractShareData container = shareData
      where
        shareData = BL.take (fromIntegral shareDataLength) . BL.drop 0x0c $ container
        shareDataLength = Binary.decode . BL.take 4 . BL.drop 4 $ container :: Word32