{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# 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.ByteString as B
import qualified Data.ByteString.Base64 as Base64
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import Data.Yaml (ParseException, decodeEither')
import GHC.Generics (
    Generic,
 )
import Tahoe.CHK.Types (
    Parameters (..),
 )

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'