{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}

module Simplex.Chat.Badges
  ( BadgeType (..),
    BadgeStatus (..),
    BadgeInfo (..),
    BadgeCredential (..),
    BadgeProof (..),
    LocalBadge (..),
    JSONBadge (..),
    BBSPublicKeyStr (..),
    localBadgeInfo,
    localBadgeStatus,
    FileSizeLimits (..),
    defaultFileSizeLimits,
    maxXFTPFileSize,
    maxSndXFTPFileSize,
    badgeSndGraceInterval,
    badgeServerCredential,
    maxFileSizeSupporter,
    maxFileSizeLegend,
    ProofPresHeaderTag (..),
    ProofPresHeader (..),
    BadgePurchase (..),
    BadgeMasterKey (..),
    BadgeRequest (..),
    VerifiedBadgeRequest (..),
    bbsBadgeHeader,
    generateMasterKey,
    verifyPayment,
    issueBadge,
    verifyCredential,
    generateBadgeProof,
    badgeProof,
    verifyBadge,
    verifyBadge_,
    mkBadgeStatus,
    BadgeRow,
    BadgeProofKind (..),
    BadgeProofRow,
    badgeProofToRow,
    rowToBadgeProof,
    badgeToRow,
    localBadgeToRow,
    rowToBadge,
  ) where

import Control.Concurrent.STM
import Crypto.Random (ChaChaDRG)
import Data.Aeson (FromJSON (..), ToJSON (..))
import qualified Data.Aeson.TH as JQ
import qualified Data.Attoparsec.ByteString.Char8 as A
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import Data.Either (fromRight)
import Data.Int (Int64)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Data.String
import Data.Text (Text)
import Data.Text.Encoding (encodeUtf8)
import Data.Time.Clock (NominalDiffTime, UTCTime, addUTCTime, nominalDay)
import Data.Time.Clock.System (systemToUTCTime, utcToSystemTime)
import Simplex.FileTransfer.Description (gb, maxFileSize)
import Simplex.Messaging.Agent.Store.DB (Binary (..), BoolInt (..), fromTextField_)
import qualified Simplex.Messaging.Crypto as C
import Simplex.Messaging.Crypto.BBS
import Simplex.Messaging.Crypto.Entitlement (Entitlement (Entitlement), EntitlementCredential (EntitlementCredential), MasterKey (MasterKey), entitlementBBSHeader)
import Simplex.Messaging.Encoding (Encoding (..))
import Simplex.Messaging.Encoding.String
import Simplex.Messaging.Parsers (defaultJSON, dropPrefix, enumJSON)
#if defined(dbPostgres)
import Database.PostgreSQL.Simple.FromField (FromField (..))
import Database.PostgreSQL.Simple.ToField (ToField (..))
#else
import Database.SQLite.Simple.FromField (FromField (..))
import Database.SQLite.Simple.ToField (ToField (..))
#endif

-- Badge type

data BadgeType
  = BTSupporter
  | BTLegend
  | BTInvestor
  | BTUnknown Text
  deriving (Eq, Show)

instance TextEncoding BadgeType where
  textEncode = \case
    BTSupporter -> "supporter"
    BTLegend -> "legend"
    BTInvestor -> "investor"
    BTUnknown tag -> tag
  textDecode s = Just $ case s of
    "supporter" -> BTSupporter
    "legend" -> BTLegend
    "investor" -> BTInvestor
    tag -> BTUnknown tag

instance ToJSON BadgeType where
  toJSON = textToJSON
  toEncoding = textToEncoding

instance FromJSON BadgeType where
  parseJSON = textParseJSON "BadgeType"

-- Badge status

data BadgeStatus = BSActive | BSExpired | BSExpiredOld | BSFailed | BSUnknownKey
  deriving (Eq, Show)

instance TextEncoding BadgeStatus where
  textEncode = \case
    BSActive -> "active"
    BSExpired -> "expired"
    BSExpiredOld -> "expired_old"
    BSFailed -> "failed"
    BSUnknownKey -> "unknown_key"
  textDecode = \case
    "active" -> Just BSActive
    "expired" -> Just BSExpired
    "expired_old" -> Just BSExpiredOld
    "failed" -> Just BSFailed
    "unknown_key" -> Just BSUnknownKey
    _ -> Nothing

-- Badge proof kind - a file has at most one proof of each kind

data BadgeProofKind = BPKInvitation | BPKDescription
  deriving (Eq, Show)

instance TextEncoding BadgeProofKind where
  textEncode = \case
    BPKInvitation -> "inv"
    BPKDescription -> "descr"
  textDecode = \case
    "inv" -> Just BPKInvitation
    "descr" -> Just BPKDescription
    _ -> Nothing

-- Disclosed badge content (BBS messages 1, 2, 3)

data BadgeInfo = BadgeInfo
  { badgeType :: BadgeType,
    badgeExpiry :: UTCTime,
    badgeExtra :: Text
  }
  deriving (Eq, Show)

-- a badge stays active for this long after its expiry, to allow for a delayed renewal
badgeGraceInterval :: NominalDiffTime
badgeGraceInterval = 7 * nominalDay

-- a badge expired longer than this ago is BSExpiredOld and is not shown in the UI
badgeOldInterval :: NominalDiffTime
badgeOldInterval = badgeGraceInterval + 31 * nominalDay

-- the verification outcome of a received proof: Just True = verified, Just False = failed,
-- Nothing = the proof's key index is not among this app version's configured keys (BSUnknownKey).
mkBadgeStatus :: UTCTime -> Maybe Bool -> BadgeInfo -> BadgeStatus
mkBadgeStatus now verified BadgeInfo {badgeExpiry = e} = case verified of
  Nothing -> BSUnknownKey
  Just False -> BSFailed
  Just True
    | addUTCTime badgeOldInterval e < now -> BSExpiredOld
    | addUTCTime badgeGraceInterval e < now -> BSExpired
    | otherwise -> BSActive

-- A badge credential (own, secret) and a proof (a presentation) are independent records.
-- badgeKeyIdx is the issuer key index: it tells verifiers which configured key to use.
-- Only proofs ride the wire (in a profile); credentials come from the badge service. Neither is
-- ever serialized as a sum - each travels as its own record, so the JSON carries no credential/proof tag.

data BadgeCredential = BadgeCredential
  { badgeKeyIdx :: Int,
    masterKey :: BadgeMasterKey,
    signature :: BBSSignature,
    badgeInfo :: BadgeInfo
  }
  deriving (Eq, Show)

data BadgeProof = BadgeProof
  { badgeKeyIdx :: Int,
    presHeader :: BBSPresHeader,
    proof :: BBSProof,
    badgeInfo :: BadgeInfo
  }
  deriving (Eq, Show)

-- Local badge: a stored badge plus its display status (the in-memory sum; never serialized as a sum).
--   OwnBadge   - the user's own credential (loaded from the DB).
--   PeerBadge  - a verified peer proof (from the DB, or received over the wire).
--   ShownBadge - decoded from a crypto-free profile JSON for display only: no crypto, so it cannot be sent.
data LocalBadge
  = OwnBadge BadgeCredential BadgeStatus
  | PeerBadge BadgeProof BadgeStatus
  | ShownBadge BadgeInfo BadgeStatus
  deriving (Eq, Show)

localBadgeInfo :: LocalBadge -> BadgeInfo
localBadgeInfo = \case
  OwnBadge BadgeCredential {badgeInfo} _ -> badgeInfo
  PeerBadge BadgeProof {badgeInfo} _ -> badgeInfo
  ShownBadge i _ -> i

localBadgeStatus :: LocalBadge -> BadgeStatus
localBadgeStatus = \case
  OwnBadge _ st -> st
  PeerBadge _ st -> st
  ShownBadge _ st -> st

-- XFTP file size limit raised by an active badge: a legend badge to 5GB, any other to 2GB, otherwise the default.
maxFileSizeSupporter :: Integer
maxFileSizeSupporter = gb 2

maxFileSizeLegend :: Integer
maxFileSizeLegend = gb 5

badgeServerCredential :: Maybe LocalBadge -> Maybe EntitlementCredential
badgeServerCredential = \case
  Just (OwnBadge (BadgeCredential idx (BadgeMasterKey mk) sig BadgeInfo {badgeType, badgeExpiry, badgeExtra}) _) ->
    Just $ EntitlementCredential (fromIntegral idx) (MasterKey mk) (Entitlement badgeExpiry (textEncode badgeType) badgeExtra) sig
  _ -> Nothing

data FileSizeLimits = FileSizeLimits
  { noBadge :: Integer,
    supporter :: Integer,
    legend :: Integer
  }
  deriving (Eq, Show)

defaultFileSizeLimits :: FileSizeLimits
defaultFileSizeLimits = FileSizeLimits {noBadge = toInteger maxFileSize, supporter = maxFileSizeSupporter, legend = maxFileSizeLegend}

-- a badge raises the size limit at send for this long after its expiry, shorter than badgeGraceInterval so the receiver still accepts the size
badgeSndGraceInterval :: NominalDiffTime
badgeSndGraceInterval = nominalDay

badgeFileSize :: FileSizeLimits -> LocalBadge -> Integer
badgeFileSize FileSizeLimits {supporter, legend} b = case badgeType (localBadgeInfo b) of
  BTLegend -> legend
  _ -> supporter

maxXFTPFileSize :: FileSizeLimits -> Maybe LocalBadge -> Integer
maxXFTPFileSize lims = \case
  Just b | localBadgeStatus b == BSActive -> badgeFileSize lims b
  _ -> noBadge lims

maxSndXFTPFileSize :: FileSizeLimits -> UTCTime -> Maybe LocalBadge -> Integer
maxSndXFTPFileSize lims now = \case
  Just b | localBadgeStatus b == BSActive && addUTCTime badgeSndGraceInterval (badgeExpiry (localBadgeInfo b)) >= now -> badgeFileSize lims b
  _ -> noBadge lims

-- Presentation header: a tag char + payload. PHTest is unbound - a fresh random nonce per
-- presentation, not bound to any context; the 'T' tag marks it so master rejects it.
-- PHUnknown is the forward-compat catch-all for tags this version does not interpret.

data ProofPresHeaderTag = PHTestTag | PHChatTag | PHFileInvTag | PHFileDescrTag | PHUnknownTag Char

instance StrEncoding ProofPresHeaderTag where
  strEncode = B.singleton . \case
    PHTestTag -> 'T'
    PHChatTag -> 'C'
    PHFileInvTag -> 'F'
    PHFileDescrTag -> 'D'
    PHUnknownTag c -> c
  strP = tag <$> A.anyChar
    where
      tag = \case
        'T' -> PHTestTag
        'C' -> PHChatTag
        'F' -> PHFileInvTag
        'D' -> PHFileDescrTag
        c -> PHUnknownTag c

data ProofPresHeader
  = PHTest ByteString
  | PHChat ByteString
  | PHFileInv {chatBinding :: ByteString, fileSize :: Int64}
  | PHFileDescr {chatBinding :: ByteString, fileSize :: Int64, descrHash :: ByteString, fileExpires :: Maybe UTCTime}
  | PHUnknown Char ByteString
  deriving (Eq, Show)
  deriving (ToJSON, FromJSON) via (StrJSON "ProofPresHeader" ProofPresHeader)

instance StrEncoding ProofPresHeader where
  strEncode = \case
    PHTest nonce -> strEncode PHTestTag <> nonce
    PHChat binding -> strEncode PHChatTag <> binding
    PHFileInv {chatBinding, fileSize} ->
      strEncode PHFileInvTag <> smpEncode (chatBinding, fileSize)
    PHFileDescr {chatBinding, fileSize, descrHash, fileExpires} ->
      strEncode PHFileDescrTag <> smpEncode (chatBinding, fileSize, descrHash, utcToSystemTime <$> fileExpires)
    PHUnknown c b -> strEncode (PHUnknownTag c) <> b
  strP =
    strP >>= \case
      PHTestTag -> PHTest <$> A.takeByteString
      PHChatTag -> PHChat <$> A.takeByteString
      PHFileInvTag -> do
        (chatBinding, fileSize) <- smpP
        pure PHFileInv {chatBinding, fileSize}
      PHFileDescrTag -> do
        (chatBinding, fileSize, descrHash, expires_) <- smpP
        pure PHFileDescr {chatBinding, fileSize, descrHash, fileExpires = systemToUTCTime <$> expires_}
      PHUnknownTag c -> PHUnknown c <$> A.takeByteString

-- v6.5.x accepts both; v7 will reject PHTest/PHUnknown
proofPresHeaderAccepted :: ProofPresHeader -> Bool
proofPresHeaderAccepted = \case
  PHTest _ -> True
  PHChat _ -> True
  PHFileInv {} -> True
  PHFileDescr {} -> True
  PHUnknown _ _ -> True

-- Payment proof

data BadgePurchase
  = BPAppleReceipt Text
  | BPGoogleReceipt Text
  | BPStripeSession
  | BPRedeemCode Text
  deriving (Eq, Show)

-- Master key

newtype BadgeMasterKey = BadgeMasterKey ByteString
  deriving newtype (Eq, Show, StrEncoding)

instance ToJSON BadgeMasterKey where
  toJSON = strToJSON
  toEncoding = strToJEncoding

instance FromJSON BadgeMasterKey where
  parseJSON = strParseJSON "BadgeMasterKey"

generateMasterKey :: TVar ChaChaDRG -> IO BadgeMasterKey
generateMasterKey drg = BadgeMasterKey <$> atomically (C.randomBytes 32 drg)

-- Workflow types

data BadgeRequest = BadgeRequest
  { masterKey :: BadgeMasterKey,
    badgeInfo :: BadgeInfo
  }
  deriving (Show)

newtype VerifiedBadgeRequest = VerifiedBadgeRequest BadgeRequest
  deriving (Show)

-- Constants

bbsBadgeHeader :: BBSHeader
bbsBadgeHeader = entitlementBBSHeader

bbsBadgeMessageCount :: Int
bbsBadgeMessageCount = 4

bbsBadgeDisclosedIndexes :: [Int]
bbsBadgeDisclosedIndexes = [1, 2, 3]

-- Message encoding

badgeMessages :: BadgeMasterKey -> BadgeInfo -> [ByteString]
badgeMessages (BadgeMasterKey ms) info = ms : badgeInfoMessages info

badgeInfoMessages :: BadgeInfo -> [ByteString]
badgeInfoMessages BadgeInfo {badgeType, badgeExpiry, badgeExtra} =
  [strEncode badgeExpiry, encodeUtf8 (textEncode badgeType), encodeUtf8 badgeExtra]

-- Payment verification (stub - always passes)

verifyPayment :: BadgePurchase -> BadgeRequest -> IO (Maybe VerifiedBadgeRequest)
verifyPayment _payment req = pure $ Just (VerifiedBadgeRequest req)

-- Server-side: issue a badge credential, recording which issuer key signed it

issueBadge :: Int -> BBSSecretKey -> VerifiedBadgeRequest -> IO (Either String BadgeCredential)
issueBadge keyIdx sk (VerifiedBadgeRequest BadgeRequest {masterKey, badgeInfo})
  | badgeExtra badgeInfo /= "" = pure $ Left "badgeExtra must be empty (reserved)"
  | otherwise = fmap (\sig -> BadgeCredential keyIdx masterKey sig badgeInfo) <$> bbsSign sk bbsBadgeHeader (badgeMessages masterKey badgeInfo)

-- Client-side: verify the credential received from server

verifyCredential :: BBSPublicKey -> BadgeCredential -> IO Bool
verifyCredential pk (BadgeCredential _ masterKey signature badgeInfo) =
  bbsVerify pk signature bbsBadgeHeader (badgeMessages masterKey badgeInfo)

-- Client-side: generate a proof for a contact/group; the proof carries the credential's key index

generateBadgeProof :: BBSPublicKey -> BadgeCredential -> BBSPresHeader -> IO (Either String BadgeProof)
generateBadgeProof pk (BadgeCredential keyIdx masterKey signature badgeInfo) ph =
  fmap (\p -> BadgeProof keyIdx ph p badgeInfo) <$> bbsProofGen pk signature bbsBadgeHeader ph bbsBadgeDisclosedIndexes (badgeMessages masterKey badgeInfo)

-- application-level proof generation with a semantic presentation header
badgeProof :: BBSPublicKey -> BadgeCredential -> ProofPresHeader -> IO (Either String BadgeProof)
badgeProof pk cred ph = generateBadgeProof pk cred (BBSPresHeader $ strEncode ph)

-- Recipient-side: verify a badge proof with the configured key its index points to.
-- Nothing means the key index is not in the configured keys (this app version can't verify it).

verifyBadge :: Map Int BBSPublicKey -> BadgeProof -> IO (Maybe Bool)
verifyBadge keys b@(BadgeProof keyIdx _ _ _) = case M.lookup keyIdx keys of
  Nothing -> pure Nothing
  Just pk -> Just <$> verifyBadgeWith pk b

verifyBadgeWith :: BBSPublicKey -> BadgeProof -> IO Bool
verifyBadgeWith pk (BadgeProof _ ph@(BBSPresHeader phBytes) proof badgeInfo)
  | either (const False) proofPresHeaderAccepted (strDecode phBytes) =
      bbsProofVerify pk proof bbsBadgeHeader ph bbsBadgeDisclosedIndexes bbsBadgeMessageCount (badgeInfoMessages badgeInfo)
  | otherwise = pure False

verifyBadge_ :: Map Int BBSPublicKey -> Maybe BadgeProof -> IO (Maybe Bool)
verifyBadge_ keys = maybe (pure (Just False)) (verifyBadge keys)

-- DB

instance FromField BadgeType where fromField = fromTextField_ textDecode

instance ToField BadgeType where toField = toField . textEncode

instance FromField BadgeStatus where fromField = fromTextField_ textDecode

instance ToField BadgeStatus where toField = toField . textEncode

instance FromField BadgeProofKind where fromField = fromTextField_ textDecode

instance ToField BadgeProofKind where toField = toField . textEncode

-- (proof, pres_header, key_idx, type, expiry, extra) - the fields of BadgeProof as stored in file_badge_proofs
type BadgeProofRow = (Binary ByteString, Binary ByteString, Int, Text, UTCTime, Text)

badgeProofToRow :: BadgeProof -> BadgeProofRow
badgeProofToRow (BadgeProof idx (BBSPresHeader ph) (BBSProof p) BadgeInfo {badgeType, badgeExpiry, badgeExtra}) =
  (Binary p, Binary ph, idx, textEncode badgeType, badgeExpiry, badgeExtra)

rowToBadgeProof :: BadgeProofRow -> Maybe BadgeProof
rowToBadgeProof (Binary p, Binary ph, idx, type_, badgeExpiry, badgeExtra) = do
  badgeType <- textDecode type_
  pure $ BadgeProof idx (BBSPresHeader ph) (BBSProof p) BadgeInfo {badgeType, badgeExpiry, badgeExtra}

-- (proof, pres_header, expiry, type, verified, extra, master_key, signature, key_idx) - binary columns wrapped in Binary (BLOB/bytea)
type BadgeRow = (Maybe (Binary ByteString), Maybe (Binary ByteString), Maybe UTCTime, Maybe Text, Maybe BoolInt, Maybe Text, Maybe (Binary ByteString), Maybe (Binary ByteString), Maybe Int)

-- receive/store sites have a wire proof + a computed verification outcome;
-- the status here only drives the stored verified flag, the display status is recomputed on load
badgeToRow :: Maybe BadgeProof -> Maybe Bool -> BadgeRow
badgeToRow badge verified = localBadgeToRow $ (`PeerBadge` st) <$> badge
  where
    st = case verified of
      Just True -> BSActive
      Just False -> BSFailed
      Nothing -> BSUnknownKey

localBadgeToRow :: Maybe LocalBadge -> BadgeRow
localBadgeToRow (Just lb) = case lb of
  OwnBadge (BadgeCredential idx (BadgeMasterKey mk) (BBSSignature sg) BadgeInfo {badgeType, badgeExpiry, badgeExtra}) st ->
    (Nothing, Nothing, Just badgeExpiry, Just (textEncode badgeType), verifiedField st, Just badgeExtra, Just (Binary mk), Just (Binary sg), Just idx)
  PeerBadge (BadgeProof idx (BBSPresHeader ph) (BBSProof p) BadgeInfo {badgeType, badgeExpiry, badgeExtra}) st ->
    (Just (Binary p), Just (Binary ph), Just badgeExpiry, Just (textEncode badgeType), verifiedField st, Just badgeExtra, Nothing, Nothing, Just idx)
  ShownBadge BadgeInfo {badgeType, badgeExpiry, badgeExtra} st ->
    (Nothing, Nothing, Just badgeExpiry, Just (textEncode badgeType), verifiedField st, Just badgeExtra, Nothing, Nothing, Nothing)
  where
    verifiedField st = case st of
      BSFailed -> Just (BI False)
      BSUnknownKey -> Nothing
      _ -> Just (BI True)
localBadgeToRow Nothing = (Nothing, Nothing, Nothing, Nothing, Just (BI False), Nothing, Nothing, Nothing, Nothing)

rowToBadge :: UTCTime -> BadgeRow -> Maybe LocalBadge
rowToBadge now (p_, ph_, expiry_, type_, verified_, extra_, mk_, sg_, idx_) = do
  btText <- type_
  bt <- textDecode btText
  badgeExpiry <- expiry_
  let info = BadgeInfo {badgeType = bt, badgeExpiry, badgeExtra = maybe "" id extra_}
      -- NULL badge_verified means the key index was unknown when stored (Nothing)
      st = mkBadgeStatus now (unBI <$> verified_) info
  case (mk_, sg_, p_, ph_, idx_) of
    (Just (Binary mk), Just (Binary sg), _, _, Just idx) -> Just $ OwnBadge (BadgeCredential idx (BadgeMasterKey mk) (BBSSignature sg) info) st
    (_, _, Just (Binary p), Just (Binary ph), Just idx) -> Just $ PeerBadge (BadgeProof idx (BBSPresHeader ph) (BBSProof p) info) st
    _ -> Just $ ShownBadge info st

-- JSON

$(JQ.deriveJSON (enumJSON $ dropPrefix "BS") ''BadgeStatus)

$(JQ.deriveJSON defaultJSON ''BadgeInfo)

$(JQ.deriveJSON defaultJSON ''BadgeRequest)

-- Each record is a plain JSON object (defaultJSON), platform-independent and with no credential/proof
-- tag - the context (a proof in a profile, a credential from the service) determines which it is.

$(JQ.deriveJSON defaultJSON ''BadgeCredential)

$(JQ.deriveJSON defaultJSON ''BadgeProof)

-- LocalBadge is sent to the UI/clients WITHOUT crypto - only disclosed info + status. The credential/proof
-- bytes stay core-side. FromJSON reconstructs a display-only badge (empty proof) for read-only consumers
-- (remote host, UI echoes); the authoritative badge is loaded from the DB (rowToBadge), never from this JSON.
data JSONBadge = JSONBadge {badge :: BadgeInfo, status :: BadgeStatus}

$(JQ.deriveJSON defaultJSON ''JSONBadge)

instance ToJSON LocalBadge where
  toJSON lb = toJSON $ JSONBadge (localBadgeInfo lb) (localBadgeStatus lb)
  toEncoding lb = toEncoding $ JSONBadge (localBadgeInfo lb) (localBadgeStatus lb)

instance FromJSON LocalBadge where
  parseJSON v = do
    JSONBadge info st <- parseJSON v
    pure $ ShownBadge info st

newtype BBSPublicKeyStr = BBSPublicKeyStr {toBBSPublicKey :: BBSPublicKey}

instance IsString BBSPublicKeyStr where
  fromString = BBSPublicKeyStr . fromRight (error "bad base64 in BBSPublicKey") . strDecode . B.pack