{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

module Simplex.Messaging.Notifications.Transport
  ( NTFVersion,
    VersionRangeNTF,
    pattern VersionNTF,
    THandleNTF,
    invalidReasonNTFVersion,
    supportedClientNTFVRange,
    supportedServerNTFVRange,
    alpnSupportedNTFHandshakes,
    ntfServerHandshake,
    ntfClientHandshake,
  ) where

import Control.Monad (forM)
import Control.Monad.Except
import Control.Monad.Trans.Except
import Data.Attoparsec.ByteString.Char8 (Parser)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import Data.Word (Word16)
import qualified Data.X509 as X
import qualified Simplex.Messaging.Crypto as C
import Simplex.Messaging.Encoding
import Simplex.Messaging.Transport
import Simplex.Messaging.Util (liftEitherWith)
import Simplex.Messaging.Version
import Simplex.Messaging.Version.Internal

ntfBlockSize :: Int
ntfBlockSize = 512

data NTFVersion

instance VersionScope NTFVersion

type VersionNTF = Version NTFVersion

type VersionRangeNTF = VersionRange NTFVersion

pattern VersionNTF :: Word16 -> VersionNTF
pattern VersionNTF v = Version v

initialNTFVersion :: VersionNTF
initialNTFVersion = VersionNTF 1

authBatchCmdsNTFVersion :: VersionNTF
authBatchCmdsNTFVersion = VersionNTF 2

invalidReasonNTFVersion :: VersionNTF
invalidReasonNTFVersion = VersionNTF 3

currentClientNTFVersion :: VersionNTF
currentClientNTFVersion = VersionNTF 3

currentServerNTFVersion :: VersionNTF
currentServerNTFVersion = VersionNTF 3

supportedClientNTFVRange :: VersionRangeNTF
supportedClientNTFVRange = mkVersionRange initialNTFVersion currentClientNTFVersion

legacyServerNTFVRange :: VersionRangeNTF
legacyServerNTFVRange = mkVersionRange initialNTFVersion initialNTFVersion

supportedServerNTFVRange :: VersionRangeNTF
supportedServerNTFVRange = mkVersionRange initialNTFVersion currentServerNTFVersion

alpnSupportedNTFHandshakes :: [ALPN]
alpnSupportedNTFHandshakes = ["ntf/1"]

type THandleNTF c p = THandle NTFVersion c p

data NtfServerHandshake = NtfServerHandshake
  { ntfVersionRange :: VersionRangeNTF,
    sessionId :: SessionId,
    -- pub key to agree shared secrets for command authorization and entity ID encryption.
    authPubKey :: Maybe (X.SignedExact X.PubKey)
  }

data NtfClientHandshake = NtfClientHandshake
  { -- | agreed SMP notifications server protocol version
    ntfVersion :: VersionNTF,
    -- | server identity - CA certificate fingerprint
    keyHash :: C.KeyHash
  }

instance Encoding NtfServerHandshake where
  smpEncode NtfServerHandshake {ntfVersionRange, sessionId, authPubKey} =
    B.concat
      [ smpEncode (ntfVersionRange, sessionId),
        encodeAuthEncryptCmds (maxVersion ntfVersionRange) $ C.SignedObject <$> authPubKey
      ]

  smpP = do
    (ntfVersionRange, sessionId) <- smpP
    -- TODO drop SMP v6: remove special parser and make key non-optional
    authPubKey <- authEncryptCmdsP (maxVersion ntfVersionRange) $ C.getSignedExact <$> smpP
    pure NtfServerHandshake {ntfVersionRange, sessionId, authPubKey}

encodeAuthEncryptCmds :: Encoding a => VersionNTF -> Maybe a -> ByteString
encodeAuthEncryptCmds v k
  | v >= authBatchCmdsNTFVersion = maybe "" smpEncode k
  | otherwise = ""

authEncryptCmdsP :: VersionNTF -> Parser a -> Parser (Maybe a)
authEncryptCmdsP v p = if v >= authBatchCmdsNTFVersion then Just <$> p else pure Nothing

instance Encoding NtfClientHandshake where
  smpEncode NtfClientHandshake {ntfVersion, keyHash} =
    smpEncode (ntfVersion, keyHash)
  smpP = do
    (ntfVersion, keyHash) <- smpP
    pure NtfClientHandshake {ntfVersion, keyHash}

-- | Notifcations server transport handshake.
ntfServerHandshake :: forall c. Transport c => C.APrivateSignKey -> c 'TServer -> C.KeyPairX25519 -> C.KeyHash -> VersionRangeNTF -> ExceptT TransportError IO (THandleNTF c 'TServer)
ntfServerHandshake serverSignKey c (k, pk) kh ntfVRange = do
  let th@THandle {params = THandleParams {sessionId}} = ntfTHandle c
  let sk = C.signX509 serverSignKey $ C.publicToX509 k
  let ntfVersionRange = maybe legacyServerNTFVRange (const ntfVRange) $ getSessionALPN c
  sendHandshake th $ NtfServerHandshake {sessionId, ntfVersionRange, authPubKey = Just sk}
  getHandshake th >>= \case
    NtfClientHandshake {ntfVersion = v, keyHash}
      | keyHash /= kh ->
          throwE $ TEHandshake IDENTITY
      | otherwise ->
          case compatibleVRange' ntfVersionRange v of
            Just (Compatible vr) -> pure $ ntfThHandleServer th v vr pk
            Nothing -> throwE TEVersion

-- | Notifcations server client transport handshake.
ntfClientHandshake :: forall c. Transport c => c 'TClient -> C.KeyHash -> VersionRangeNTF -> Bool -> Maybe (ServiceCredentials, C.KeyPairEd25519) -> ExceptT TransportError IO (THandleNTF c 'TClient)
ntfClientHandshake c keyHash ntfVRange _proxyServer _serviceKeys = do
  let th@THandle {params = THandleParams {sessionId}} = ntfTHandle c
  NtfServerHandshake {sessionId = sessId, ntfVersionRange, authPubKey = sk'} <- getHandshake th
  if sessionId /= sessId
    then throwE TEBadSession
    else case ntfVersionRange `compatibleVRange` ntfVRange of
      Just (Compatible vr) -> do
        ck_ <- forM sk' $ \signedKey -> liftEitherWith (const $ TEHandshake BAD_AUTH) $ do
          serverKey <- getServerVerifyKey c
          pubKey <- C.verifyX509 serverKey signedKey
          (,CertChainPubKey (getPeerCertChain c) signedKey) <$> C.x509ToPublic' pubKey
        let v = maxVersion vr
        sendHandshake th $ NtfClientHandshake {ntfVersion = v, keyHash}
        pure $ ntfThHandleClient th v vr ck_
      Nothing -> throwE TEVersion

ntfThHandleServer :: forall c. THandleNTF c 'TServer -> VersionNTF -> VersionRangeNTF -> C.PrivateKeyX25519 -> THandleNTF c 'TServer
ntfThHandleServer th v vr pk =
  let thAuth = THAuthServer {serverPrivKey = pk, peerClientService = Nothing, sessSecret' = Nothing}
   in ntfThHandle_ th v vr (Just thAuth)

ntfThHandleClient :: forall c. THandleNTF c 'TClient -> VersionNTF -> VersionRangeNTF -> Maybe (C.PublicKeyX25519, CertChainPubKey) -> THandleNTF c 'TClient
ntfThHandleClient th v vr ck_ =
  let thAuth = clientTHParams <$> ck_
      clientTHParams (k, ck) = THAuthClient {peerServerPubKey = k, peerServerCertKey = ck, clientService = Nothing, sessSecret = Nothing}
   in ntfThHandle_ th v vr thAuth

ntfThHandle_ :: forall c p. THandleNTF c p -> VersionNTF -> VersionRangeNTF -> Maybe (THandleAuth p) -> THandleNTF c p
ntfThHandle_ th@THandle {params} v vr thAuth =
  -- TODO drop SMP v6: make thAuth non-optional
  let v3 = v >= authBatchCmdsNTFVersion
      params' = params {thVersion = v, thServerVRange = vr, thAuth, implySessId = v3, batch = v3}
   in (th :: THandleNTF c p) {params = params'}

ntfTHandle :: Transport c => c p -> THandleNTF c p
ntfTHandle c = THandle {connection = c, params}
  where
    v = VersionNTF 0
    params =
      THandleParams
        { sessionId = tlsUnique c,
          blockSize = ntfBlockSize,
          thVersion = v,
          thServerVRange = versionToRange v,
          thAuth = Nothing,
          implySessId = False,
          encryptBlock = Nothing,
          batch = False,
          serviceAuth = False
        }