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

module Simplex.FileTransfer.Transport
  ( supportedFileServerVRange,
    authCmdsXFTPVersion,
    blockedFilesXFTPVersion,
    xftpClientHandshakeStub,
    alpnSupportedXFTPhandshakes,
    xftpALPNv1,
    XFTPClientHandshake (..),
    -- xftpClientHandshake,
    XFTPServerHandshake (..),
    -- xftpServerHandshake,
    XFTPClientHello (..),
    THandleXFTP,
    THandleParamsXFTP,
    VersionXFTP,
    VersionRangeXFTP,
    XFTPVersion,
    pattern VersionXFTP,
    XFTPErrorType (..),
    XFTPRcvChunkSpec (..),
    ReceiveFileError (..),
    receiveFile,
    sendEncFile,
    receiveEncFile,
    receiveSbFile,
  )
where

import Control.Applicative (optional)
import qualified Control.Exception as E
import Control.Logger.Simple
import Control.Monad
import Control.Monad.Except
import Control.Monad.IO.Class
import Control.Monad.Trans.Except
import qualified Data.Aeson.TH as J
import qualified Data.Attoparsec.ByteString.Char8 as A
import Data.Bifunctor (first)
import qualified Data.ByteArray as BA
import Data.ByteString.Builder (Builder, byteString)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy.Char8 as LB
import Data.Functor (($>))
import Data.Word (Word16, Word32)
import Network.HTTP2.Client (HTTP2Error)
import qualified Simplex.Messaging.Crypto as C
import qualified Simplex.Messaging.Crypto.Lazy as LC
import Simplex.Messaging.Encoding
import Simplex.Messaging.Encoding.String
import Simplex.Messaging.Parsers
import Simplex.Messaging.Protocol (BlockingInfo, CommandError)
import Simplex.Messaging.Transport (ALPN, CertChainPubKey, ServiceCredentials, SessionId, THandle (..), THandleParams (..), TransportError (..), TransportPeer (..))
import Simplex.Messaging.Transport.HTTP2.File
import Simplex.Messaging.Util (bshow, tshow, (<$?>))
import Simplex.Messaging.Version
import Simplex.Messaging.Version.Internal
import System.IO (Handle, IOMode (..), withFile)

data XFTPRcvChunkSpec = XFTPRcvChunkSpec
  { filePath :: FilePath,
    chunkSize :: Word32,
    chunkDigest :: ByteString
  }
  deriving (Show)

data XFTPVersion

instance VersionScope XFTPVersion

type VersionXFTP = Version XFTPVersion

type VersionRangeXFTP = VersionRange XFTPVersion

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

type THandleXFTP c p = THandle XFTPVersion c p
type THandleParamsXFTP p = THandleParams XFTPVersion p

initialXFTPVersion :: VersionXFTP
initialXFTPVersion = VersionXFTP 1

authCmdsXFTPVersion :: VersionXFTP
authCmdsXFTPVersion = VersionXFTP 2

blockedFilesXFTPVersion :: VersionXFTP
blockedFilesXFTPVersion = VersionXFTP 3

currentXFTPVersion :: VersionXFTP
currentXFTPVersion = VersionXFTP 3

supportedFileServerVRange :: VersionRangeXFTP
supportedFileServerVRange = mkVersionRange initialXFTPVersion currentXFTPVersion

-- spec: spec/modules/Simplex/FileTransfer/Transport.md#xftpclienthandshakestub--xftp-doesnt-use-tls-handshake
-- XFTP protocol does not use this handshake method
xftpClientHandshakeStub :: c 'TClient -> Maybe C.KeyPairX25519 -> C.KeyHash -> VersionRangeXFTP -> Bool -> Maybe (ServiceCredentials, C.KeyPairEd25519) -> ExceptT TransportError IO (THandle XFTPVersion c 'TClient)
xftpClientHandshakeStub _c _ks _keyHash _xftpVRange _proxyServer _serviceKeys = throwE TEVersion

alpnSupportedXFTPhandshakes :: [ALPN]
alpnSupportedXFTPhandshakes = [xftpALPNv1]

xftpALPNv1 :: ALPN
xftpALPNv1 = "xftp/1"

data XFTPClientHello = XFTPClientHello
  { -- | a random string sent by the client to the server to prove that server has identity certificate
    webChallenge :: Maybe ByteString
  }

data XFTPServerHandshake = XFTPServerHandshake
  { xftpVersionRange :: VersionRangeXFTP,
    sessionId :: SessionId,
    -- | pub key to agree shared secrets for command authorization and entity ID encryption.
    authPubKey :: CertChainPubKey,
    -- | signed identity challenge from  XFTPClientHello
    webIdentityProof :: Maybe C.ASignature
  }

data XFTPClientHandshake = XFTPClientHandshake
  { -- | agreed XFTP server protocol version
    xftpVersion :: VersionXFTP,
    -- | server identity - CA certificate fingerprint
    keyHash :: C.KeyHash
  }

instance Encoding XFTPClientHello where
  smpEncode XFTPClientHello {webChallenge} = smpEncode webChallenge
  smpP = do
    webChallenge <- smpP
    forM_ webChallenge $ \challenge -> unless (B.length challenge == 32) $ fail "bad XFTPClientHello webChallenge"
    Tail _compat <- smpP
    pure XFTPClientHello {webChallenge}

instance Encoding XFTPClientHandshake where
  smpEncode XFTPClientHandshake {xftpVersion, keyHash} =
    smpEncode (xftpVersion, keyHash)
  smpP = do
    (xftpVersion, keyHash) <- smpP
    Tail _compat <- smpP
    pure XFTPClientHandshake {xftpVersion, keyHash}

instance Encoding XFTPServerHandshake where
  smpEncode XFTPServerHandshake {xftpVersionRange, sessionId, authPubKey, webIdentityProof} =
    smpEncode (xftpVersionRange, sessionId, authPubKey, C.signatureBytes webIdentityProof)
  smpP = do
    (xftpVersionRange, sessionId, authPubKey) <- smpP
    webIdentityProof <- optional $ C.decodeSignature <$?> smpP
    Tail _compat <- smpP
    pure XFTPServerHandshake {xftpVersionRange, sessionId, authPubKey, webIdentityProof}

sendEncFile :: Handle -> (Builder -> IO ()) -> LC.SbState -> Word32 -> IO ()
sendEncFile h send = go
  where
    go sbState 0 = do
      let authTag = BA.convert (LC.sbAuth sbState)
      send $ byteString authTag
    go sbState sz =
      getFileChunk h sz >>= \ch -> do
        let (encCh, sbState') = LC.sbEncryptChunk sbState ch
        send (byteString encCh) `E.catch` \(e :: E.SomeException) -> print e >> E.throwIO e
        go sbState' $ sz - fromIntegral (B.length ch)

receiveFile :: (Int -> IO ByteString) -> XFTPRcvChunkSpec -> ExceptT XFTPErrorType IO ()
receiveFile getBody chunk = ExceptT $ runExceptT (receiveFile_ receive chunk) `E.catches` handlers
  where
    receive h sz = hReceiveFile getBody h sz >>= \sz' -> pure $ if sz' == 0 then Right () else Left SIZE
    handlers =
      [ E.Handler $ \(e :: HTTP2Error) -> logWarn (err e) $> Left TIMEOUT,
        E.Handler $ \(e :: E.SomeException) -> logError (err e) $> Left FILE_IO
      ]
    err e = "receiveFile error: " <> tshow e

receiveEncFile :: (Int -> IO ByteString) -> LC.SbState -> XFTPRcvChunkSpec -> ExceptT XFTPErrorType IO ()
receiveEncFile getBody = receiveFile_ . receive
  where
    receive sbState h sz = first err <$> receiveSbFile getBody h sbState sz
    err RFESize = SIZE
    err RFECrypto = CRYPTO

data ReceiveFileError = RFESize | RFECrypto

-- spec: spec/modules/Simplex/FileTransfer/Transport.md#receivesbfile--constant-time-auth-tag-verification
receiveSbFile :: (Int -> IO ByteString) -> Handle -> LC.SbState -> Word32 -> IO (Either ReceiveFileError ())
receiveSbFile getBody h = receive
  where
    receive sbState sz = do
      ch <- getBody fileBlockSize
      let chSize = fromIntegral $ B.length ch
      if
        | chSize > sz + authSz -> pure $ Left RFESize
        | chSize > 0 -> do
            let (ch', rest) = B.splitAt (fromIntegral sz) ch
                (decCh, sbState') = LC.sbDecryptChunk sbState ch'
                sz' = sz - fromIntegral (B.length ch')
            B.hPut h decCh
            if sz' > 0
              then receive sbState' sz'
              else do
                let tag' = B.take C.authTagSize rest
                    tagSz = B.length tag'
                    tag = LC.sbAuth sbState'
                tag'' <- if tagSz == C.authTagSize then pure tag' else (tag' <>) <$> getBody (C.authTagSize - tagSz)
                pure $ if BA.constEq tag'' tag then Right () else Left RFECrypto
        | otherwise -> pure $ Left RFESize
    authSz = fromIntegral C.authTagSize

receiveFile_ :: (Handle -> Word32 -> IO (Either XFTPErrorType ())) -> XFTPRcvChunkSpec -> ExceptT XFTPErrorType IO ()
receiveFile_ receive XFTPRcvChunkSpec {filePath, chunkSize, chunkDigest} = do
  ExceptT $ withFile filePath WriteMode (`receive` chunkSize)
  digest' <- liftIO $ LC.sha256Hash <$> LB.readFile filePath
  when (digest' /= chunkDigest) $ throwE DIGEST

data XFTPErrorType
  = -- | incorrect block format, encoding or signature size
    BLOCK
  | -- | incorrect SMP session ID (TLS Finished message / tls-unique binding RFC5929)
    SESSION
  | -- | incorrect handshake command
    HANDSHAKE
  | -- | SMP command is unknown or has invalid syntax
    CMD {cmdErr :: CommandError}
  | -- | command authorization error - bad signature or non-existing SMP queue
    AUTH
  | -- | command with the entity that was blocked
    BLOCKED {blockInfo :: BlockingInfo}
  | -- | incorrent file size
    SIZE
  | -- | storage quota exceeded
    QUOTA
  | -- | incorrent file digest
    DIGEST
  | -- | file encryption/decryption failed
    CRYPTO
  | -- | no expected file body in request/response or no file on the server
    NO_FILE
  | -- | unexpected file body
    HAS_FILE
  | -- | file IO error
    FILE_IO
  | -- | file sending or receiving timeout
    TIMEOUT
  | -- | internal server error
    INTERNAL
  | -- | used internally, never returned by the server (to be removed)
    DUPLICATE_ -- not part of SMP protocol, used internally
  deriving (Eq, Show)

instance StrEncoding XFTPErrorType where
  strEncode = \case
    BLOCK -> "BLOCK"
    SESSION -> "SESSION"
    HANDSHAKE -> "HANDSHAKE"
    CMD e -> "CMD " <> bshow e
    AUTH -> "AUTH"
    BLOCKED info -> "BLOCKED " <> strEncode info
    SIZE -> "SIZE"
    QUOTA -> "QUOTA"
    DIGEST -> "DIGEST"
    CRYPTO -> "CRYPTO"
    NO_FILE -> "NO_FILE"
    HAS_FILE -> "HAS_FILE"
    FILE_IO -> "FILE_IO"
    TIMEOUT -> "TIMEOUT"
    INTERNAL -> "INTERNAL"
    DUPLICATE_ -> "DUPLICATE_"

  strP =
    A.takeTill (== ' ') >>= \case
      "BLOCK" -> pure BLOCK
      "SESSION" -> pure SESSION
      "HANDSHAKE" -> pure HANDSHAKE
      "CMD" -> CMD <$> parseRead1
      "AUTH" -> pure AUTH
      "BLOCKED" -> BLOCKED <$> _strP
      "SIZE" -> pure SIZE
      "QUOTA" -> pure QUOTA
      "DIGEST" -> pure DIGEST
      "CRYPTO" -> pure CRYPTO
      "NO_FILE" -> pure NO_FILE
      "HAS_FILE" -> pure HAS_FILE
      "FILE_IO" -> pure FILE_IO
      "TIMEOUT" -> pure TIMEOUT
      "INTERNAL" -> pure INTERNAL
      "DUPLICATE_" -> pure DUPLICATE_
      _ -> fail "bad error type"

instance Encoding XFTPErrorType where
  smpEncode = \case
    BLOCK -> "BLOCK"
    SESSION -> "SESSION"
    HANDSHAKE -> "HANDSHAKE"
    CMD err -> "CMD " <> smpEncode err
    AUTH -> "AUTH"
    BLOCKED info -> "BLOCKED " <> smpEncode info
    SIZE -> "SIZE"
    QUOTA -> "QUOTA"
    DIGEST -> "DIGEST"
    CRYPTO -> "CRYPTO"
    NO_FILE -> "NO_FILE"
    HAS_FILE -> "HAS_FILE"
    FILE_IO -> "FILE_IO"
    TIMEOUT -> "TIMEOUT"
    INTERNAL -> "INTERNAL"
    DUPLICATE_ -> "DUPLICATE_"

  smpP =
    A.takeTill (== ' ') >>= \case
      "BLOCK" -> pure BLOCK
      "SESSION" -> pure SESSION
      "HANDSHAKE" -> pure HANDSHAKE
      "CMD" -> CMD <$> _smpP
      "AUTH" -> pure AUTH
      "BLOCKED" -> BLOCKED <$> _smpP
      "SIZE" -> pure SIZE
      "QUOTA" -> pure QUOTA
      "DIGEST" -> pure DIGEST
      "CRYPTO" -> pure CRYPTO
      "NO_FILE" -> pure NO_FILE
      "HAS_FILE" -> pure HAS_FILE
      "FILE_IO" -> pure FILE_IO
      "TIMEOUT" -> pure TIMEOUT
      "INTERNAL" -> pure INTERNAL
      "DUPLICATE_" -> pure DUPLICATE_
      _ -> fail "bad error type"

$(J.deriveJSON (sumTypeJSON id) ''XFTPErrorType)