{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Use newtype instead of data" #-}

module Simplex.Messaging.Notifications.Server.Push.APNS
  ( PushNotification (..),
    APNSNotification (..),
    APNSNotificationBody (..),
    APNSAlertBody (..),
    APNSPushClientConfig (..),
    PushProviderError (..),
    PushProviderClient,
    APNSErrorResponse (..),
    apnsProviderHost,
    defaultAPNSPushClientConfig,
    createAPNSPushClient,
    apnsPushProviderClient,
  ) where

import Control.Exception (Exception)
import Control.Logger.Simple
import Control.Monad
import Control.Monad.Except
import Control.Monad.IO.Class
import Control.Monad.Trans.Except
import Crypto.Hash.Algorithms (SHA256 (..))
import qualified Crypto.PubKey.ECC.ECDSA as EC
import qualified Crypto.PubKey.ECC.Types as ECT
import Crypto.Random (ChaChaDRG)
import qualified Crypto.Store.PKCS8 as PK
import Data.ASN1.BinaryEncoding (DER (..))
import Data.ASN1.Encoding
import Data.ASN1.Types
import Data.Aeson (ToJSON, (.=))
import qualified Data.Aeson as J
import qualified Data.Aeson.Encoding as JE
import qualified Data.Aeson.TH as JQ
import Data.Bifunctor (first)
import qualified Data.ByteString.Base64.URL as U
import Data.ByteString.Builder (lazyByteString)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Lazy.Char8 as LB
import qualified Data.CaseInsensitive as CI
import Data.Int (Int64)
import Data.List (find)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map.Strict (Map)
import Data.Maybe (isNothing)
import Data.Text (Text)
import qualified Data.Text as T
import Data.Text.Encoding (encodeUtf8)
import Data.Time.Clock.System
import qualified Data.X509 as X
import qualified Data.X509.CertificateStore as XS
import Network.HPACK.Token as HT
import Network.HTTP.Types (Status)
import qualified Network.HTTP.Types as N
import Network.HTTP2.Client (Request)
import qualified Network.HTTP2.Client as H
import Network.Socket (HostName, ServiceName)
import qualified Simplex.Messaging.Crypto as C
import Simplex.Messaging.Notifications.Protocol
import Simplex.Messaging.Notifications.Server.Push.APNS.Internal
import Simplex.Messaging.Notifications.Server.Store.Types (NtfTknRec (..))
import Simplex.Messaging.Parsers (defaultJSON)
import Simplex.Messaging.Transport.HTTP2 (HTTP2Body (..))
import Simplex.Messaging.Transport.HTTP2.Client
import Simplex.Messaging.Util (safeDecodeUtf8, tshow)
import System.Environment (getEnv)
import UnliftIO.STM

data JWTHeader = JWTHeader
  { alg :: Text, -- key algorithm, ES256 for APNS
    kid :: Text -- key ID
  }
  deriving (Show)

data JWTClaims = JWTClaims
  { iss :: Text, -- issuer, team ID for APNS
    iat :: Int64 -- issue time, seconds from epoch
  }
  deriving (Show)

data JWTToken = JWTToken JWTHeader JWTClaims
  deriving (Show)

mkJWTToken :: JWTHeader -> Text -> IO JWTToken
mkJWTToken hdr iss = do
  iat <- systemSeconds <$> getSystemTime
  pure $ JWTToken hdr JWTClaims {iss, iat}

type SignedJWTToken = ByteString

$(JQ.deriveToJSON defaultJSON ''JWTHeader)

$(JQ.deriveToJSON defaultJSON ''JWTClaims)

signedJWTToken :: EC.PrivateKey -> JWTToken -> IO SignedJWTToken
signedJWTToken pk (JWTToken hdr claims) = do
  let hc = jwtEncode hdr <> "." <> jwtEncode claims
  sig <- EC.sign pk SHA256 hc
  pure $ hc <> "." <> serialize sig
  where
    jwtEncode :: ToJSON a => a -> ByteString
    jwtEncode = U.encodeUnpadded . LB.toStrict . J.encode
    serialize sig = U.encodeUnpadded $ encodeASN1' DER [Start Sequence, IntVal (EC.sign_r sig), IntVal (EC.sign_s sig), End Sequence]

readECPrivateKey :: FilePath -> IO EC.PrivateKey
readECPrivateKey f = do
  -- this pattern match is specific to APNS key type, it may need to be extended for other push providers
  [PK.Unprotected (X.PrivKeyEC X.PrivKeyEC_Named {privkeyEC_name, privkeyEC_priv})] <- PK.readKeyFile f
  pure EC.PrivateKey {private_curve = ECT.getCurveByName privkeyEC_name, private_d = privkeyEC_priv}

data PushNotification
  = PNVerification NtfRegCode
  | PNMessage (NonEmpty PNMessageData)
  | -- | PNAlert Text
    PNCheckMessages
  deriving (Show)

data APNSNotification = APNSNotification {aps :: APNSNotificationBody, notificationData :: Maybe J.Value}
  deriving (Show)

data APNSNotificationBody
  = APNSBackground {contentAvailable :: Int}
  | APNSMutableContent {mutableContent :: Int, alert :: APNSAlertBody, category :: Maybe Text}
  | APNSAlert {alert :: APNSAlertBody, badge :: Maybe Int, sound :: Maybe Text, category :: Maybe Text}
  deriving (Show)

type APNSNotificationData = Map Text Text

data APNSAlertBody = APNSAlertObject {title :: Text, subtitle :: Text, body :: Text} | APNSAlertText Text
  deriving (Show)

instance ToJSON APNSAlertBody where
  toEncoding = \case
    APNSAlertObject {title, subtitle, body} -> J.pairs $ "title" .= title <> "subtitle" .= subtitle <> "body" .= body
    APNSAlertText t -> JE.text t
  toJSON = \case
    APNSAlertObject {title, subtitle, body} -> J.object ["title" .= title, "subtitle" .= subtitle, "body" .= body]
    APNSAlertText t -> J.String t

-- APNS notification types
--
-- Visible alerts:
-- {
--    "aps" : {
--       "alert" : {
--          "title" : "Game Request",
--          "subtitle" : "Five Card Draw",
--          "body" : "Bob wants to play poker"
--       },
--       "badge" : 9,
--       "sound" : "bingbong.aiff",
--       "category" : "GAME_INVITATION"
--    },
--    "gameID" : "12345678"
-- }
--
-- Simple text alert:
-- {"aps":{"alert":"you have a new message"}}
--
-- Background notification to fetch content
-- {"aps":{"content-available":1}}
--
-- Mutable content notification that must be shown but can be processed before before being shown (up to 30 sec)
-- {
--    "aps" : {
--       "category" : "SECRET",
--       "mutable-content" : 1,
--       "alert" : {
--          "title" : "Secret Message!",
--          "body"  : "(Encrypted)"
--      },
--    },
--    "ENCRYPTED_DATA" : "Salted__·öîQÊ$UDì_¶Ù∞èΩ^¬%gq∞NÿÒQùw"
-- }

data APNSPushClientConfig = APNSPushClientConfig
  { tokenTTL :: Int64,
    authKeyFileEnv :: String,
    authKeyAlg :: Text,
    authKeyIdEnv :: String,
    paddedNtfLength :: Int,
    appName :: ByteString,
    appTeamId :: Text,
    apnsPort :: ServiceName,
    http2cfg :: HTTP2ClientConfig,
    caStoreFile :: FilePath
  }

apnsProviderHost :: PushProvider -> Maybe HostName
apnsProviderHost = \case
  PPApnsNull -> Nothing
  PPApnsTest -> Just "localhost"
  PPApnsDev -> Just "api.sandbox.push.apple.com"
  PPApnsProd -> Just "api.push.apple.com"

defaultAPNSPushClientConfig :: APNSPushClientConfig
defaultAPNSPushClientConfig =
  APNSPushClientConfig
    { tokenTTL = 1800, -- 30 minutes
      authKeyFileEnv = "APNS_KEY_FILE", -- the environment variables APNS_KEY_FILE and APNS_KEY_ID must be set, or the server would fail to start
      authKeyAlg = "ES256",
      authKeyIdEnv = "APNS_KEY_ID",
      paddedNtfLength = 3072,
      appName = "chat.simplex.app",
      appTeamId = "5NN7GUYB6T",
      apnsPort = "443",
      http2cfg = defaultHTTP2ClientConfig {bufferSize = 16384},
      caStoreFile = "/etc/ssl/cert.pem"
    }

data APNSPushClient = APNSPushClient
  { https2Client :: TVar (Maybe HTTP2Client),
    privateKey :: EC.PrivateKey,
    jwtHeader :: JWTHeader,
    jwtToken :: TVar (JWTToken, SignedJWTToken),
    nonceDrg :: TVar ChaChaDRG,
    apnsHost :: HostName,
    apnsCfg :: APNSPushClientConfig
  }

createAPNSPushClient :: HostName -> APNSPushClientConfig -> IO APNSPushClient
createAPNSPushClient apnsHost apnsCfg@APNSPushClientConfig {authKeyFileEnv, authKeyAlg, authKeyIdEnv, appTeamId} = do
  https2Client <- newTVarIO Nothing
  void $ connectHTTPS2 apnsHost apnsCfg https2Client
  privateKey <- readECPrivateKey =<< getEnv authKeyFileEnv
  authKeyId <- T.pack <$> getEnv authKeyIdEnv
  let jwtHeader = JWTHeader {alg = authKeyAlg, kid = authKeyId}
  jwtToken <- newTVarIO =<< mkApnsJWTToken appTeamId jwtHeader privateKey
  nonceDrg <- C.newRandom
  pure APNSPushClient {https2Client, privateKey, jwtHeader, jwtToken, nonceDrg, apnsHost, apnsCfg}

getApnsJWTToken :: APNSPushClient -> IO SignedJWTToken
getApnsJWTToken APNSPushClient {apnsCfg = APNSPushClientConfig {appTeamId, tokenTTL}, privateKey, jwtHeader, jwtToken} = do
  (jwt, signedJWT) <- readTVarIO jwtToken
  age <- jwtTokenAge jwt
  if age < tokenTTL
    then pure signedJWT
    else do
      t@(_, signedJWT') <- mkApnsJWTToken appTeamId jwtHeader privateKey
      atomically $ writeTVar jwtToken t
      pure signedJWT'
  where
    jwtTokenAge (JWTToken _ JWTClaims {iat}) = subtract iat . systemSeconds <$> getSystemTime

mkApnsJWTToken :: Text -> JWTHeader -> EC.PrivateKey -> IO (JWTToken, SignedJWTToken)
mkApnsJWTToken appTeamId jwtHeader privateKey = do
  jwt <- mkJWTToken jwtHeader appTeamId
  signedJWT <- signedJWTToken privateKey jwt
  pure (jwt, signedJWT)

connectHTTPS2 :: HostName -> APNSPushClientConfig -> TVar (Maybe HTTP2Client) -> IO (Either HTTP2ClientError HTTP2Client)
connectHTTPS2 apnsHost APNSPushClientConfig {apnsPort, http2cfg, caStoreFile} https2Client = do
  caStore_ <- XS.readCertificateStore caStoreFile
  when (isNothing caStore_) $ logError $ "Error loading CertificateStore from " <> T.pack caStoreFile
  r <- getHTTP2Client apnsHost apnsPort caStore_ http2cfg disconnected
  case r of
    Right client -> atomically . writeTVar https2Client $ Just client
    Left e -> logError $ "Error connecting to APNS: " <> tshow e
  pure r
  where
    disconnected = atomically $ writeTVar https2Client Nothing

getApnsHTTP2Client :: APNSPushClient -> IO (Either HTTP2ClientError HTTP2Client)
getApnsHTTP2Client APNSPushClient {https2Client, apnsHost, apnsCfg} =
  readTVarIO https2Client >>= maybe (connectHTTPS2 apnsHost apnsCfg https2Client) (pure . Right)

disconnectApnsHTTP2Client :: APNSPushClient -> IO ()
disconnectApnsHTTP2Client APNSPushClient {https2Client} =
  readTVarIO https2Client >>= mapM_ closeHTTP2Client >> atomically (writeTVar https2Client Nothing)

ntfCategoryCheckMessage :: Text
ntfCategoryCheckMessage = "NTF_CAT_CHECK_MESSAGE"

apnsNotification :: NtfTknRec -> C.CbNonce -> Int -> PushNotification -> Either C.CryptoError APNSNotification
apnsNotification NtfTknRec {tknDhSecret} nonce paddedLen = \case
  PNVerification (NtfRegCode code) ->
    encrypt code $ \code' ->
      apn APNSBackground {contentAvailable = 1} . Just $ J.object ["nonce" .= nonce, "verification" .= code']
  PNMessage pnMessageData ->
    encrypt (encodePNMessages pnMessageData) $ \ntfData ->
      apn apnMutableContent . Just $ J.object ["nonce" .= nonce, "message" .= ntfData]
  -- PNAlert text -> Right $ apn (apnAlert $ APNSAlertText text) Nothing
  PNCheckMessages -> Right $ apn APNSBackground {contentAvailable = 1} . Just $ J.object ["checkMessages" .= True]
  where
    encrypt :: ByteString -> (Text -> APNSNotification) -> Either C.CryptoError APNSNotification
    encrypt ntfData f = f . safeDecodeUtf8 . U.encode <$> C.cbEncrypt tknDhSecret nonce ntfData paddedLen
    apn aps notificationData = APNSNotification {aps, notificationData}
    apnMutableContent = APNSMutableContent {mutableContent = 1, alert = APNSAlertText "Encrypted message or another app event", category = Just ntfCategoryCheckMessage}

-- apnAlert alert = APNSAlert {alert, badge = Nothing, sound = Nothing, category = Nothing}

$(JQ.deriveToJSON apnsJSONOptions ''APNSNotificationBody)

$(JQ.deriveToJSON defaultJSON ''APNSNotification)

apnsRequest :: APNSPushClient -> ByteString -> APNSNotification -> IO Request
apnsRequest c tkn ntf@APNSNotification {aps} = do
  signedJWT <- getApnsJWTToken c
  pure $ H.requestBuilder N.methodPost path (headers signedJWT) (lazyByteString $ J.encode ntf)
  where
    path = "/3/device/" <> tkn
    headers signedJWT =
      [ (hApnsTopic, appName $ apnsCfg (c :: APNSPushClient)),
        (hApnsPushType, pushType aps),
        (N.hAuthorization, "bearer " <> signedJWT)
      ]
        <> [(hApnsPriority, "5") | isBackground aps]
    isBackground = \case
      APNSBackground {} -> True
      _ -> False
    pushType = \case
      APNSBackground {} -> "background"
      _ -> "alert"

data PushProviderError
  = PPConnection HTTP2ClientError
  | PPCryptoError C.CryptoError
  | PPResponseError (Maybe Status) Text
  | PPTokenInvalid NTInvalidReason
  | PPRetryLater
  | PPPermanentError
  deriving (Show, Exception)

type PushProviderClient = NtfTknRec -> PushNotification -> ExceptT PushProviderError IO ()

-- this is not a newtype on purpose to have a correct JSON encoding as a record
data APNSErrorResponse = APNSErrorResponse {reason :: Text}

$(JQ.deriveFromJSON defaultJSON ''APNSErrorResponse)

apnsPushProviderClient :: APNSPushClient -> PushProviderClient
apnsPushProviderClient c@APNSPushClient {nonceDrg, apnsCfg} tkn@NtfTknRec {token = DeviceToken _ tknStr} pn = do
  http2 <- liftHTTPS2 $ getApnsHTTP2Client c
  nonce <- atomically $ C.randomCbNonce nonceDrg
  apnsNtf <- liftEither $ first PPCryptoError $ apnsNotification tkn nonce (paddedNtfLength apnsCfg) pn
  req <- liftIO $ apnsRequest c tknStr apnsNtf
  -- TODO when HTTP2 client is thread-safe, we can use sendRequestDirect
  HTTP2Response {response, respBody = HTTP2Body {bodyHead}} <- liftHTTPS2 $ sendRequest http2 req Nothing
  let status = H.responseStatus response
      reason' = maybe "" reason $ J.decodeStrict' bodyHead
  if status == Just N.ok200
    then logDebug $ "APNS response: ok" <> apnsIds response
    else logWarn $ "APNS error: " <> T.pack (show status) <> " " <> reason' <> apnsIds response
  result status reason'
  where
    apnsIds response = headerStr "apns-id" <> headerStr "apns-unique-id"
      where
        headerStr name =
          maybe "" (\(_, v) -> ", " <> name <> ": " <> safeDecodeUtf8 v) $
            find (\(t, _) -> HT.tokenKey t == CI.mk (encodeUtf8 name)) (fst (H.responseHeaders response))
    result :: Maybe Status -> Text -> ExceptT PushProviderError IO ()
    result status reason'
      | status == Just N.ok200 = pure ()
      | status == Just N.badRequest400 =
          throwE $ case reason' of
            "BadDeviceToken" -> PPTokenInvalid NTIRBadToken
            "DeviceTokenNotForTopic" -> PPTokenInvalid NTIRTokenNotForTopic
            "TopicDisallowed" -> PPPermanentError
            _ -> PPResponseError status reason'
      | status == Just N.forbidden403 = throwE $ case reason' of
          "ExpiredProviderToken" -> PPPermanentError -- there should be no point retrying it as the token was refreshed
          "InvalidProviderToken" -> PPPermanentError
          _ -> PPResponseError status reason'
      | status == Just N.gone410 = throwE $ case reason' of
          "ExpiredToken" -> PPTokenInvalid NTIRExpiredToken
          "Unregistered" -> PPTokenInvalid NTIRUnregistered
          _ -> PPRetryLater
      | status == Just N.serviceUnavailable503 = liftIO (disconnectApnsHTTP2Client c) >> throwE PPRetryLater
      -- Just tooManyRequests429 -> TooManyRequests - too many requests for the same token
      | otherwise = throwE $ PPResponseError status reason'
    liftHTTPS2 a = ExceptT $ first PPConnection <$> a