{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-unticked-promoted-constructors #-}

module Simplex.FileTransfer.Protocol
  ( FileParty (..),
    SFileParty (..),
    AFileParty (..),
    FilePartyI (..),
    FileCommand (..),
    FileCmd (..),
    FileInfo (..),
    XFTPFileId,
    FileResponse (..),
    xftpBlockSize,
    toFileParty,
    aFileParty,
    checkParty,
    xftpEncodeAuthTransmission,
    xftpEncodeTransmission,
    xftpDecodeTServer,
    xftpDecodeTClient,
  ) where

import qualified Data.Aeson.TH as J
import Data.Bifunctor (first)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import Data.Kind (Type)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Maybe (isNothing)
import Data.Type.Equality
import Data.Word (Word32)
import Simplex.FileTransfer.Transport (XFTPErrorType (..), XFTPVersion, blockedFilesXFTPVersion, xftpClientHandshakeStub)
import Simplex.Messaging.Client (authTransmission)
import qualified Simplex.Messaging.Crypto as C
import Simplex.Messaging.Encoding
import Simplex.Messaging.Encoding.String
import Simplex.Messaging.Parsers
import Simplex.Messaging.Protocol
  ( BasicAuth,
    CommandError (..),
    Protocol (..),
    ProtocolEncoding (..),
    ProtocolErrorType (..),
    ProtocolMsgTag (..),
    ProtocolType (..),
    RcvPublicAuthKey,
    RcvPublicDhKey,
    EntityId (..),
    RecipientId,
    SenderId,
    RawTransmission,
    SentRawTransmission,
    SignedTransmissionOrError,
    SndPublicAuthKey,
    Transmission,
    TransmissionForAuth (..),
    CorrId (..),
    encodeTransmission,
    encodeTransmissionForAuth,
    messageTagP,
    tDecodeServer,
    tDecodeClient,
    tEncodeBatch1,
    tParse,
  )
import Simplex.Messaging.Transport (THandleParams (..), TransportError (..), TransportPeer (..))
import Simplex.Messaging.Util ((<$?>))

xftpBlockSize :: Int
xftpBlockSize = 16384

-- | File protocol clients
data FileParty = FRecipient | FSender
  deriving (Eq, Show)

data SFileParty :: FileParty -> Type where
  SFRecipient :: SFileParty FRecipient
  SFSender :: SFileParty FSender

instance TestEquality SFileParty where
  testEquality SFRecipient SFRecipient = Just Refl
  testEquality SFSender SFSender = Just Refl
  testEquality _ _ = Nothing

deriving instance Eq (SFileParty p)

deriving instance Show (SFileParty p)

data AFileParty = forall p. FilePartyI p => AFP (SFileParty p)

toFileParty :: SFileParty p -> FileParty
toFileParty = \case
  SFRecipient -> FRecipient
  SFSender -> FSender

aFileParty :: FileParty -> AFileParty
aFileParty = \case
  FRecipient -> AFP SFRecipient
  FSender -> AFP SFSender

class FilePartyI (p :: FileParty) where sFileParty :: SFileParty p

instance FilePartyI FRecipient where sFileParty = SFRecipient

instance FilePartyI FSender where sFileParty = SFSender

data FileCommandTag (p :: FileParty) where
  FNEW_ :: FileCommandTag FSender
  FADD_ :: FileCommandTag FSender
  FPUT_ :: FileCommandTag FSender
  FDEL_ :: FileCommandTag FSender
  FGET_ :: FileCommandTag FRecipient
  FACK_ :: FileCommandTag FRecipient
  PING_ :: FileCommandTag FRecipient

deriving instance Show (FileCommandTag p)

data FileCmdTag = forall p. FilePartyI p => FCT (SFileParty p) (FileCommandTag p)

instance FilePartyI p => Encoding (FileCommandTag p) where
  smpEncode = \case
    FNEW_ -> "FNEW"
    FADD_ -> "FADD"
    FPUT_ -> "FPUT"
    FDEL_ -> "FDEL"
    FGET_ -> "FGET"
    FACK_ -> "FACK"
    PING_ -> "PING"
  smpP = messageTagP

instance Encoding FileCmdTag where
  smpEncode (FCT _ t) = smpEncode t
  smpP = messageTagP

instance ProtocolMsgTag FileCmdTag where
  decodeTag = \case
    "FNEW" -> Just $ FCT SFSender FNEW_
    "FADD" -> Just $ FCT SFSender FADD_
    "FPUT" -> Just $ FCT SFSender FPUT_
    "FDEL" -> Just $ FCT SFSender FDEL_
    "FGET" -> Just $ FCT SFRecipient FGET_
    "FACK" -> Just $ FCT SFRecipient FACK_
    "PING" -> Just $ FCT SFRecipient PING_
    _ -> Nothing

instance FilePartyI p => ProtocolMsgTag (FileCommandTag p) where
  decodeTag s = decodeTag s >>= (\(FCT _ t) -> checkParty' t)

instance Protocol XFTPVersion XFTPErrorType FileResponse where
  type ProtoCommand FileResponse = FileCmd
  type ProtoType FileResponse = 'PXFTP
  protocolClientHandshake = xftpClientHandshakeStub
  {-# INLINE protocolClientHandshake #-}
  useServiceAuth _ = False
  {-# INLINE useServiceAuth #-}
  protocolPing = FileCmd SFRecipient PING
  {-# INLINE protocolPing #-}
  protocolError = \case
    FRErr e -> Just e
    _ -> Nothing
  {-# INLINE protocolError #-}

data FileCommand (p :: FileParty) where
  FNEW :: FileInfo -> NonEmpty RcvPublicAuthKey -> Maybe BasicAuth -> FileCommand FSender
  FADD :: NonEmpty RcvPublicAuthKey -> FileCommand FSender
  FPUT :: FileCommand FSender
  FDEL :: FileCommand FSender
  FGET :: RcvPublicDhKey -> FileCommand FRecipient
  FACK :: FileCommand FRecipient
  PING :: FileCommand FRecipient

deriving instance Show (FileCommand p)

data FileCmd = forall p. FilePartyI p => FileCmd (SFileParty p) (FileCommand p)

deriving instance Show FileCmd

data FileInfo = FileInfo
  { sndKey :: SndPublicAuthKey,
    size :: Word32,
    digest :: ByteString
  }
  deriving (Show)

type XFTPFileId = EntityId

instance FilePartyI p => ProtocolEncoding XFTPVersion XFTPErrorType (FileCommand p) where
  type Tag (FileCommand p) = FileCommandTag p
  encodeProtocol _v = \case
    FNEW file rKeys auth_ -> e (FNEW_, ' ', file, rKeys, auth_)
    FADD rKeys -> e (FADD_, ' ', rKeys)
    FPUT -> e FPUT_
    FDEL -> e FDEL_
    FGET rKey -> e (FGET_, ' ', rKey)
    FACK -> e FACK_
    PING -> e PING_
    where
      e :: Encoding a => a -> ByteString
      e = smpEncode

  protocolP v tag = (\(FileCmd _ c) -> checkParty c) <$?> protocolP v (FCT (sFileParty @p) tag)

  fromProtocolError = fromProtocolError @XFTPVersion @XFTPErrorType @FileResponse
  {-# INLINE fromProtocolError #-}

  checkCredentials auth (EntityId fileId) cmd = case cmd of
    -- FNEW must not have signature and chunk ID
    FNEW {}
      | isNothing auth -> Left $ CMD NO_AUTH
      | not (B.null fileId) -> Left $ CMD HAS_AUTH
      | otherwise -> Right cmd
    PING
      | isNothing auth && B.null fileId -> Right cmd
      | otherwise -> Left $ CMD HAS_AUTH
    -- other client commands must have both signature and queue ID
    _
      | isNothing auth || B.null fileId -> Left $ CMD NO_AUTH
      | otherwise -> Right cmd

instance ProtocolEncoding XFTPVersion XFTPErrorType FileCmd where
  type Tag FileCmd = FileCmdTag
  encodeProtocol _v (FileCmd _ c) = encodeProtocol _v c

  protocolP _v = \case
    FCT SFSender tag ->
      FileCmd SFSender <$> case tag of
        FNEW_ -> FNEW <$> _smpP <*> smpP <*> smpP
        FADD_ -> FADD <$> _smpP
        FPUT_ -> pure FPUT
        FDEL_ -> pure FDEL
    FCT SFRecipient tag ->
      FileCmd SFRecipient <$> case tag of
        FGET_ -> FGET <$> _smpP
        FACK_ -> pure FACK
        PING_ -> pure PING

  fromProtocolError = fromProtocolError @XFTPVersion @XFTPErrorType @FileResponse
  {-# INLINE fromProtocolError #-}

  checkCredentials tAuth entId (FileCmd p c) = FileCmd p <$> checkCredentials tAuth entId c
  {-# INLINE checkCredentials #-}

instance Encoding FileInfo where
  smpEncode FileInfo {sndKey, size, digest} = smpEncode (sndKey, size, digest)
  smpP = FileInfo <$> smpP <*> smpP <*> smpP

instance StrEncoding FileInfo where
  strEncode FileInfo {sndKey, size, digest} = strEncode (sndKey, size, digest)
  strP = FileInfo <$> strP_ <*> strP_ <*> strP

data FileResponseTag
  = FRSndIds_
  | FRRcvIds_
  | FRFile_
  | FROk_
  | FRErr_
  | FRPong_
  deriving (Show)

instance Encoding FileResponseTag where
  smpEncode = \case
    FRSndIds_ -> "SIDS"
    FRRcvIds_ -> "RIDS"
    FRFile_ -> "FILE"
    FROk_ -> "OK"
    FRErr_ -> "ERR"
    FRPong_ -> "PONG"
  smpP = messageTagP

instance ProtocolMsgTag FileResponseTag where
  decodeTag = \case
    "SIDS" -> Just FRSndIds_
    "RIDS" -> Just FRRcvIds_
    "FILE" -> Just FRFile_
    "OK" -> Just FROk_
    "ERR" -> Just FRErr_
    "PONG" -> Just FRPong_
    _ -> Nothing

data FileResponse
  = FRSndIds SenderId (NonEmpty RecipientId)
  | FRRcvIds (NonEmpty RecipientId)
  | FRFile RcvPublicDhKey C.CbNonce
  | FROk
  | FRErr XFTPErrorType
  | FRPong
  deriving (Show)

instance ProtocolEncoding XFTPVersion XFTPErrorType FileResponse where
  type Tag FileResponse = FileResponseTag
  encodeProtocol v = \case
    FRSndIds fId rIds -> e (FRSndIds_, ' ', fId, rIds)
    FRRcvIds rIds -> e (FRRcvIds_, ' ', rIds)
    FRFile rDhKey nonce -> e (FRFile_, ' ', rDhKey, nonce)
    FROk -> e FROk_
    FRErr err -> case err of
      BLOCKED _ | v < blockedFilesXFTPVersion -> e (FRErr_, ' ', AUTH)
      _ -> e (FRErr_, ' ', err)
    FRPong -> e FRPong_
    where
      e :: Encoding a => a -> ByteString
      e = smpEncode

  protocolP _v = \case
    FRSndIds_ -> FRSndIds <$> _smpP <*> smpP
    FRRcvIds_ -> FRRcvIds <$> _smpP
    FRFile_ -> FRFile <$> _smpP <*> smpP
    FROk_ -> pure FROk
    FRErr_ -> FRErr <$> _smpP
    FRPong_ -> pure FRPong

  fromProtocolError = \case
    PECmdSyntax -> CMD SYNTAX
    PECmdUnknown -> CMD UNKNOWN
    PESession -> SESSION
    PEBlock -> BLOCK
  {-# INLINE fromProtocolError #-}

  checkCredentials _ (EntityId entId) cmd = case cmd of
    FRSndIds {} -> noEntity
    -- ERR response does not always have entity ID
    FRErr _ -> Right cmd
    -- PONG response must not have queue ID
    FRPong -> noEntity
    -- other server responses must have entity ID
    _
      | B.null entId -> Left $ CMD NO_ENTITY
      | otherwise -> Right cmd
    where
      noEntity
        | B.null entId = Right cmd
        | otherwise = Left $ CMD HAS_AUTH

checkParty :: forall t p p'. (FilePartyI p, FilePartyI p') => t p' -> Either String (t p)
checkParty c = case testEquality (sFileParty @p) (sFileParty @p') of
  Just Refl -> Right c
  Nothing -> Left "incorrect XFTP party"

checkParty' :: forall t p p'. (FilePartyI p, FilePartyI p') => t p' -> Maybe (t p)
checkParty' c = case testEquality (sFileParty @p) (sFileParty @p') of
  Just Refl -> Just c
  _ -> Nothing

xftpEncodeAuthTransmission :: ProtocolEncoding XFTPVersion XFTPErrorType c => THandleParams XFTPVersion 'TClient -> C.APrivateAuthKey -> Transmission c -> Either TransportError ByteString
xftpEncodeAuthTransmission thParams@THandleParams {thAuth} pKey t@(corrId, _, _) = do
  let TransmissionForAuth {tForAuth, tToSend} = encodeTransmissionForAuth thParams t
  xftpEncodeBatch1 . (,tToSend) =<< authTransmission thAuth False (Just pKey) (C.cbNonce $ bs corrId) tForAuth

xftpEncodeTransmission :: ProtocolEncoding XFTPVersion XFTPErrorType c => THandleParams XFTPVersion p -> Transmission c -> Either TransportError ByteString
xftpEncodeTransmission thParams t = xftpEncodeBatch1 (Nothing, encodeTransmission thParams t)

-- this function uses batch syntax but puts only one transmission in the batch
xftpEncodeBatch1 :: SentRawTransmission -> Either TransportError ByteString
xftpEncodeBatch1 t = first (const TELargeMsg) $ C.pad (tEncodeBatch1 False t) xftpBlockSize

xftpDecodeTServer :: THandleParams XFTPVersion 'TServer -> ByteString -> Either XFTPErrorType (SignedTransmissionOrError XFTPErrorType FileCmd)
xftpDecodeTServer = xftpDecodeTransmission tDecodeServer
{-# INLINE xftpDecodeTServer #-}

xftpDecodeTClient :: THandleParams XFTPVersion 'TClient -> ByteString -> Either XFTPErrorType (Transmission (Either XFTPErrorType FileResponse))
xftpDecodeTClient = xftpDecodeTransmission tDecodeClient
{-# INLINE xftpDecodeTClient #-}

xftpDecodeTransmission ::
  (THandleParams XFTPVersion p -> Either TransportError RawTransmission -> r) ->
  THandleParams XFTPVersion p ->
  ByteString ->
  Either XFTPErrorType r
xftpDecodeTransmission tDecode thParams t = do
  t' <- first (const BLOCK) $ C.unPad t
  case tParse thParams t' of
    t'' :| [] -> Right $ tDecode thParams t''
    _ -> Left BLOCK

$(J.deriveJSON (enumJSON $ dropPrefix "F") ''FileParty)