-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Test.Wire.API.MLS where

import Control.Concurrent.Async
import Crypto.PubKey.Ed25519 qualified as Ed25519
import Data.ByteArray hiding (length)
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as B8
import Data.Domain
import Data.Id
import Data.Json.Util (toBase64Text)
import Data.Qualified
import Data.Text qualified as T
import Data.Text qualified as Text
import Data.UUID.V4 qualified as UUID
import Imports
import System.Exit
import System.FilePath ((</>))
import System.Process
import System.Random
import Test.Tasty
import Test.Tasty.HUnit
import UnliftIO (withSystemTempDirectory)
import Wire.API.MLS.AuthenticatedContent
import Wire.API.MLS.CipherSuite
import Wire.API.MLS.Commit
import Wire.API.MLS.Credential
import Wire.API.MLS.Epoch
import Wire.API.MLS.Group
import Wire.API.MLS.GroupInfo
import Wire.API.MLS.HPKEPublicKey
import Wire.API.MLS.KeyPackage
import Wire.API.MLS.Message
import Wire.API.MLS.Proposal
import Wire.API.MLS.ProtocolVersion
import Wire.API.MLS.Serialisation
import Wire.API.MLS.Welcome

tests :: TestTree
tests =
  testGroup "MLS" $
    [ testCase "parse key package" testParseKeyPackage,
      testCase "parse capabilities in key package" testParseKeyPackageWithCapabilities,
      testCase "parse commit message" testParseCommit,
      testCase "parse application message" testParseApplication,
      testCase "parse welcome and groupinfo message" testParseWelcomeAndGroupInfo,
      testCase "key package ref" testKeyPackageRef,
      testCase "create signed remove proposal" testRemoveProposalMessageSignature,
      testCase "parse client identity" testParseClientIdentity
    ]

testParseClientIdentity :: IO ()
testParseClientIdentity = do
  let cid = "wireapp://qHiDLsbkT2-p9uSJsmrZ_A%217f39900830740008@wire.com"
  let actual = sanIdentity cid
  show <$> actual @?= Right "a878832e-c6e4-4f6f-a9f6-e489b26ad9fc:7f39900830740008@wire.com"

testParseKeyPackage :: IO ()
testParseKeyPackage = do
  alice <- randomIdentity
  let qcid = B8.unpack (encodeMLS' alice)
  kpData <- withSystemTempDirectory "mls" $ \tmp -> do
    void $ spawn (cli qcid tmp ["init", qcid]) Nothing
    spawn (cli qcid tmp ["key-package", "create"]) Nothing

  kp <- case decodeMLS' @KeyPackage kpData of
    Left err -> assertFailure (T.unpack err)
    Right x -> pure x

  pvTag (kp.protocolVersion) @?= Just ProtocolMLS10
  kp.cipherSuite @?= CipherSuite 1
  BS.length (unHPKEPublicKey kp.initKey) @?= 32

  case keyPackageIdentity kp of
    Left err -> assertFailure $ "Failed to parse identity: " <> T.unpack err
    Right identity -> identity @?= RegularClient alice

testParseKeyPackageWithCapabilities :: IO ()
testParseKeyPackageWithCapabilities = do
  kpData <- BS.readFile "test/resources/key_package1.mls"
  case decodeMLS' @KeyPackage kpData of
    Left err -> assertFailure (T.unpack err)
    Right _ -> pure ()

testParseCommit :: IO ()
testParseCommit = do
  qcid <- B8.unpack . encodeMLS' <$> randomIdentity
  commitData <- withSystemTempDirectory "mls" $ \tmp -> do
    void $ spawn (cli qcid tmp ["init", qcid]) Nothing
    groupJSON <- spawn (cli qcid tmp ["group", "create", "Zm9v"]) Nothing
    spawn (cli qcid tmp ["commit", "--group", "-"]) (Just groupJSON)

  msg <- case decodeMLS' @Message commitData of
    Left err -> assertFailure (T.unpack err)
    Right x -> pure x

  pvTag (msg.protocolVersion) @?= Just ProtocolMLS10

  pmsg <- case msg.content of
    MessagePublic x -> pure x
    _ -> assertFailure "expected public message"

  pmsg.content.value.sender @?= SenderMember 0

  commit <- case pmsg.content.value.content of
    FramedContentCommit c -> pure c
    _ -> assertFailure "expected commit"

  commit.value.proposals @?= []

testParseApplication :: IO ()
testParseApplication = do
  qcid <- B8.unpack . encodeMLS' <$> randomIdentity
  msgData <- withSystemTempDirectory "mls" $ \tmp -> do
    void $ spawn (cli qcid tmp ["init", qcid]) Nothing
    groupJSON <- spawn (cli qcid tmp ["group", "create", "Zm9v"]) Nothing
    spawn (cli qcid tmp ["message", "--group-in", "-", "hello"]) (Just groupJSON)

  msg <- case decodeMLS' @Message msgData of
    Left err -> assertFailure (T.unpack err)
    Right x -> pure x

  pvTag (msg.protocolVersion) @?= Just ProtocolMLS10

  pmsg <- case msg.content of
    MessagePrivate x -> pure x.value
    _ -> assertFailure "expected private message"

  pmsg.groupId @?= GroupId "foo"
  pmsg.epoch @?= Epoch 0

testParseWelcomeAndGroupInfo :: IO ()
testParseWelcomeAndGroupInfo = do
  qcid <- B8.unpack . encodeMLS' <$> randomIdentity
  qcid2 <- B8.unpack . encodeMLS' <$> randomIdentity
  (welData, giData) <- withSystemTempDirectory "mls" $ \tmp -> do
    void $ spawn (cli qcid tmp ["init", qcid]) Nothing
    void $ spawn (cli qcid2 tmp ["init", qcid2]) Nothing
    groupJSON <- spawn (cli qcid tmp ["group", "create", "Zm9v"]) Nothing
    kp <- spawn (cli qcid2 tmp ["key-package", "create"]) Nothing
    BS.writeFile (tmp </> "kp") kp
    void $
      spawn
        ( cli
            qcid
            tmp
            [ "member",
              "add",
              "--group",
              "-",
              tmp </> "kp",
              "--welcome-out",
              tmp </> "welcome",
              "--group-info-out",
              tmp </> "gi"
            ]
        )
        (Just groupJSON)
    (,)
      <$> BS.readFile (tmp </> "welcome")
      <*> BS.readFile (tmp </> "gi")

  do
    welcomeMsg <- case decodeMLS' @Message welData of
      Left err -> assertFailure (T.unpack err)
      Right x -> pure x

    pvTag (welcomeMsg.protocolVersion) @?= Just ProtocolMLS10

    wel <- case welcomeMsg.content of
      MessageWelcome x -> pure x.value
      _ -> assertFailure "expected welcome message"

    length (wel.welSecrets) @?= 1

  do
    gi <- case decodeMLS' @GroupInfo giData of
      Left err -> assertFailure (T.unpack err)
      Right x -> pure x

    gi.groupContext.groupId @?= GroupId "foo"
    gi.groupContext.epoch @?= Epoch 1

testKeyPackageRef :: IO ()
testKeyPackageRef = do
  let qcid = "b455a431-9db6-4404-86e7-6a3ebe73fcaf:3ae58155@mls.example.com"
  (kpData, ref) <- withSystemTempDirectory "mls" $ \tmp -> do
    void $ spawn (cli qcid tmp ["init", qcid]) Nothing
    kpData <- spawn (cli qcid tmp ["key-package", "create"]) Nothing
    ref <- spawn (cli qcid tmp ["key-package", "ref", "-"]) (Just kpData)
    pure (kpData, KeyPackageRef ref)

  kpRef MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519 (KeyPackageData kpData) @?= ref

testRemoveProposalMessageSignature :: IO ()
testRemoveProposalMessageSignature = withSystemTempDirectory "mls" $ \tmp -> do
  qcid <- do
    let c = ClientId 0x3ae58155
    usr <- flip Qualified (Domain "example.com") <$> (Id <$> UUID.nextRandom)
    pure (userClientQid usr c)
  void $ spawn (cli qcid tmp ["init", qcid]) Nothing

  qcid2 <- do
    let c = ClientId 0x4ae58157
    usr <- flip Qualified (Domain "example.com") <$> (Id <$> UUID.nextRandom)
    pure (userClientQid usr c)
  void $ spawn (cli qcid2 tmp ["init", qcid2]) Nothing
  kp :: RawMLS KeyPackage <-
    decodeMLSError <$> spawn (cli qcid2 tmp ["key-package", "create"]) Nothing
  BS.writeFile (tmp </> qcid2) (raw kp)

  secretKey <- Ed25519.generateSecretKey
  let groupFilename = "group"
      gid = GroupId "abcd"
      signerKeyFilename = "signer-key.bin"
      publicKey = Ed25519.toPublic secretKey
  BS.writeFile (tmp </> signerKeyFilename) (convert publicKey)
  createGroup tmp qcid groupFilename signerKeyFilename gid

  void $ spawn (cli qcid tmp ["member", "add", "--group", tmp </> groupFilename, "--in-place", tmp </> qcid2]) Nothing

  let proposal = mkRawMLS (RemoveProposal 1)
  pmessage <-
    mkSignedPublicMessage
      @Ed25519
      (secretKey, publicKey)
      gid
      (Epoch 1)
      (TaggedSenderExternal 0)
      (FramedContentProposal proposal)
  let message = mkMessage $ MessagePublic pmessage
      messageFilename = "signed-message.mls"

  BS.writeFile (tmp </> messageFilename) (raw (mkRawMLS message))

  void $
    spawn
      ( cli
          qcid
          tmp
          [ "consume",
            "--group",
            tmp </> groupFilename,
            tmp </> messageFilename
          ]
      )
      Nothing

createGroup :: FilePath -> String -> String -> String -> GroupId -> IO ()
createGroup tmp store groupName removalKey gid = do
  groupJSON <-
    liftIO $
      spawn
        ( cli
            store
            tmp
            [ "group",
              "create",
              "--removal-key",
              tmp </> removalKey,
              T.unpack (toBase64Text (unGroupId gid))
            ]
        )
        Nothing
  liftIO $ BS.writeFile (tmp </> groupName) groupJSON

decodeMLSError :: (ParseMLS a) => ByteString -> a
decodeMLSError s = case decodeMLS' s of
  Left e -> error ("Could not parse MLS object: " <> Text.unpack e)
  Right x -> x

userClientQid :: Qualified UserId -> ClientId -> String
userClientQid usr c =
  show (qUnqualified usr)
    <> ":"
    <> T.unpack (clientToText c)
    <> "@"
    <> T.unpack (domainText (qDomain usr))

spawn :: (HasCallStack) => CreateProcess -> Maybe ByteString -> IO ByteString
spawn cp minput = do
  (mout, ex) <- withCreateProcess
    cp
      { std_out = CreatePipe,
        std_in = CreatePipe
      }
    $ \minh mouth _ ph ->
      let writeInput = for_ minh $ \inh -> do
            forM_ minput $ BS.hPutStr inh
            hClose inh
          readOutput = (,) <$> traverse BS.hGetContents mouth <*> waitForProcess ph
       in snd <$> concurrently writeInput readOutput
  case (mout, ex) of
    (Just out, ExitSuccess) -> pure out
    _ -> assertFailure $ "Failed spawning process\n" <> show mout <> "\n" <> show ex

cli :: String -> FilePath -> [String] -> CreateProcess
cli store tmp args =
  proc "mls-test-cli" $
    ["--store", tmp </> (store <> ".db")] <> args

randomIdentity :: IO ClientIdentity
randomIdentity = do
  uid <- Id <$> UUID.nextRandom
  c <- ClientId <$> randomIO
  pure $ ClientIdentity (Domain "mls.example.com") uid c
