-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Wire.API.MLS.Keys where

import Control.Lens ((?~))
import Crypto.ECC (Curve_P256R1, Curve_P384R1, Curve_P521R1)
import Crypto.PubKey.ECDSA qualified as ECDSA
import Data.Aeson (FromJSON (..), ToJSON (..))
import Data.Aeson qualified as A
import Data.Bifunctor
import Data.ByteArray (ByteArray)
import Data.ByteArray qualified as BA
import Data.Default
import Data.Json.Util
import Data.OpenApi qualified as OpenApi
import Data.OpenApi qualified as S
import Data.Proxy
import Data.Schema hiding (HasField)
import Imports hiding (First, getFirst)
import Web.HttpApiData
import Wire.API.MLS.CipherSuite

data MLSKeysByPurpose a = MLSKeysByPurpose
  { removal :: a
  }
  deriving (Eq, Show, Functor, Foldable, Traversable)
  deriving (FromJSON, ToJSON, S.ToSchema) via Schema (MLSKeysByPurpose a)

instance (Typeable a, ToSchema a) => ToSchema (MLSKeysByPurpose a) where
  schema =
    object $
      MLSKeysByPurpose
        <$> (.removal) .= field "removal" schema

data MLSKeys a = MLSKeys
  { ed25519 :: a,
    ecdsa_secp256r1_sha256 :: a,
    ecdsa_secp384r1_sha384 :: a,
    ecdsa_secp521r1_sha512 :: a
  }
  deriving (Eq, Show, Functor, Foldable, Traversable)
  deriving (FromJSON, ToJSON, S.ToSchema) via Schema (MLSKeys a)

instance (Typeable a, ToSchema a) => ToSchema (MLSKeys a) where
  schema =
    object $
      MLSKeys
        <$> ed25519 .= field "ed25519" schema
        <*> ecdsa_secp256r1_sha256 .= field "ecdsa_secp256r1_sha256" schema
        <*> ecdsa_secp384r1_sha384 .= field "ecdsa_secp384r1_sha384" schema
        <*> ecdsa_secp521r1_sha512 .= field "ecdsa_secp521r1_sha512" schema

data MLSPrivateKeys = MLSPrivateKeys
  { mlsKeyPair_ed25519 :: KeyPair Ed25519,
    mlsKeyPair_ecdsa_secp256r1_sha256 :: KeyPair Ecdsa_secp256r1_sha256,
    mlsKeyPair_ecdsa_secp384r1_sha384 :: KeyPair Ecdsa_secp384r1_sha384,
    mlsKeyPair_ecdsa_secp521r1_sha512 :: KeyPair Ecdsa_secp521r1_sha512
  }

instance ToSchema MLSPrivateKeys where
  schema =
    object $
      MLSPrivateKeys
        <$> (.mlsKeyPair_ed25519) .= field @NamedSwaggerDoc "ed25519" (opaqueSchema "KeyPair Ed25519")
        <*> (.mlsKeyPair_ecdsa_secp256r1_sha256) .= field @NamedSwaggerDoc "ecdsa_secp256r1_sha256" (opaqueSchema "KeyPair Ecdsa_secp256r1_sha256")
        <*> (.mlsKeyPair_ecdsa_secp384r1_sha384) .= field @NamedSwaggerDoc "ecdsa_secp384r1_sha384" (opaqueSchema "KeyPair Ecdsa_secp384r1_sha384")
        <*> (.mlsKeyPair_ecdsa_secp521r1_sha512) .= field @NamedSwaggerDoc "ecdsa_secp521r1_sha512" (opaqueSchema "KeyPair Ecdsa_secp521r1_sha512")

opaqueSchema :: Text -> ValueSchema NamedSwaggerDoc a
opaqueSchema name = mkSchema (pure $ OpenApi.NamedSchema (Just name) (OpenApi.binarySchema & OpenApi.format ?~ "password")) (error "not implemented") (const Nothing)

type MLSPublicKeys = MLSKeys MLSPublicKey

newtype MLSPublicKey = MLSPublicKey {unwrapMLSPublicKey :: ByteString}
  deriving (Eq, Show)
  deriving (ToJSON) via Schema MLSPublicKey

instance ToSchema MLSPublicKey where
  schema = named "MLSPublicKey" $ MLSPublicKey <$> unwrapMLSPublicKey .= base64Schema

mlsKeysToPublic :: MLSPrivateKeys -> MLSPublicKeys
mlsKeysToPublic (MLSPrivateKeys (_, ed) (_, ec256) (_, ec384) (_, ec521)) =
  MLSKeys
    { ed25519 = MLSPublicKey $ BA.convert ed,
      ecdsa_secp256r1_sha256 = MLSPublicKey $ ECDSA.encodePublic (Proxy @Curve_P256R1) ec256,
      ecdsa_secp384r1_sha384 = MLSPublicKey $ ECDSA.encodePublic (Proxy @Curve_P384R1) ec384,
      ecdsa_secp521r1_sha512 = MLSPublicKey $ ECDSA.encodePublic (Proxy @Curve_P521R1) ec521
    }

data MLSPublicKeyFormat = MLSPublicKeyFormatRaw | MLSPublicKeyFormatJWK
  deriving (Eq, Ord, Show)

instance Default MLSPublicKeyFormat where
  def = MLSPublicKeyFormatRaw

instance FromHttpApiData MLSPublicKeyFormat where
  parseQueryParam "raw" = pure MLSPublicKeyFormatRaw
  parseQueryParam "jwk" = pure MLSPublicKeyFormatJWK
  parseQueryParam _ = Left "invalid MLSPublicKeyFormat"

instance ToHttpApiData MLSPublicKeyFormat where
  toQueryParam MLSPublicKeyFormatRaw = "raw"
  toQueryParam MLSPublicKeyFormatJWK = "jwk"

instance S.ToParamSchema MLSPublicKeyFormat where
  toParamSchema _ =
    mempty
      & S.type_ ?~ S.OpenApiString
      & S.enum_
        ?~ map
          (toJSON . toQueryParam)
          [MLSPublicKeyFormatRaw, MLSPublicKeyFormatJWK]

data JWK = JWK
  { keyType :: String,
    curve :: String,
    pubX :: ByteString,
    pubY :: Maybe ByteString
  }
  deriving (Show, Ord, Eq)
  deriving (ToJSON) via Schema JWK

instance ToSchema JWK where
  schema =
    object $
      JWK
        <$> (.keyType) .= field "kty" schema
        <*> (.curve) .= field "crv" schema
        <*> (.pubX) .= field "x" base64URLSchema
        <*> (.pubY) .= maybe_ (optField "y" base64URLSchema)

type MLSPublicKeysJWK = MLSKeys JWK

mlsPublicKeysToJWK :: MLSPublicKeys -> Maybe MLSPublicKeysJWK
mlsPublicKeysToJWK (MLSKeys (MLSPublicKey ed) (MLSPublicKey ec256) (MLSPublicKey ec384) (MLSPublicKey ec521)) =
  -- The kty parameter for ECDSA is "EC", for Ed25519 it's "OKP" (octet key
  -- pair).
  -- https://www.rfc-editor.org/rfc/rfc7518.html#section-6.1

  -- The crv for P-256, P-384, and P-521 are their names.
  -- https://www.rfc-editor.org/rfc/rfc7518.html#section-6.2.1.1

  -- The x parameter is mandatory for all keys, the y parameter is mandatory for
  -- all ECDSA keys.
  -- https://www.rfc-editor.org/rfc/rfc7518.html#section-6.2.1
  MLSKeys
    (JWK "OKP" "Ed25519" (BA.convert ed) Nothing)
    <$> (uncurry (JWK "EC" "P-256") . second Just <$> splitXY ec256)
    <*> (uncurry (JWK "EC" "P-384") . second Just <$> splitXY ec384)
    <*> (uncurry (JWK "EC" "P-521") . second Just <$> splitXY ec521)
  where
    -- Obtaining X and Y from an encoded curve point follows the logic of
    -- Crypto.ECC's encodeECPoint and decodeECPoint (the module is not
    -- exported). Points need to be encoded in uncompressed representation. This
    -- is true for ECDSA.encodePublic.
    -- https://www.rfc-editor.org/rfc/rfc8422#section-5.4.1
    splitXY :: forall {bs}. (ByteArray bs) => bs -> Maybe (bs, bs)
    splitXY mxy = do
      (m, xy) <- BA.uncons mxy
      -- The first Byte m is 4 for the uncompressed representation of curve points.
      guard (m == 4)
      -- The first half of the following Bytes belong to X and the second half
      -- to Y.
      let size = BA.length xy `div` 2
      pure $ BA.splitAt size xy

data SomeKey = SomeKey A.Value

instance ToSchema SomeKey where
  schema = mkSchema d r w
    where
      d = pure $ S.NamedSchema (Just "SomeKey") mempty
      r = fmap SomeKey . parseJSON
      w (SomeKey x) = Just (toJSON x)

mkSomeKey :: (ToJSON a) => a -> SomeKey
mkSomeKey = SomeKey . toJSON
