-- 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/>.
{-# LANGUAGE RecordWildCards #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 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 API.MLS.Util where

import Bilge
import Bilge.Assert
import Data.Aeson (object, toJSON, (.=))
import Data.ByteString.Conversion
import Data.Default
import Data.Id
import Data.Json.Util
import Data.Map qualified as Map
import Data.Qualified
import Data.Text qualified as Text
import Data.Timeout
import Imports
import System.FilePath
import System.Process
import Test.Tasty.HUnit
import Util
import Wire.API.MLS.CipherSuite
import Wire.API.MLS.Credential
import Wire.API.MLS.KeyPackage
import Wire.API.MLS.Serialisation
import Wire.API.User.Client

data SetKey = SetKey | DontSetKey
  deriving (Eq)

data KeyingInfo = KeyingInfo
  { kiSetKey :: SetKey,
    kiLifetime :: Maybe Timeout
  }

instance Default KeyingInfo where
  def = KeyingInfo SetKey Nothing

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

initStore ::
  (HasCallStack) =>
  (MonadIO m) =>
  FilePath ->
  ClientIdentity ->
  m ()
initStore tmp qcid = do
  void . liftIO . flip spawn Nothing $
    cliCmd tmp qcid ["init", show qcid]

generateKeyPackage ::
  (HasCallStack) =>
  (MonadIO m) =>
  FilePath ->
  ClientIdentity ->
  Maybe Timeout ->
  m (RawMLS KeyPackage, KeyPackageRef)
generateKeyPackage tmp qcid lifetime = do
  kp <-
    liftIO $
      decodeMLSError <=< flip spawn Nothing $
        cliCmd tmp qcid $
          ["key-package", "create"]
            <> (("--lifetime " <>) . show . (#> Second) <$> maybeToList lifetime)
  let ref = fromJust (kpRef' kp)
  pure (kp, ref)

uploadKeyPackages ::
  (HasCallStack) =>
  Brig ->
  FilePath ->
  KeyingInfo ->
  Qualified UserId ->
  ClientId ->
  Int ->
  Http ()
uploadKeyPackages brig tmp KeyingInfo {..} u c n = do
  let cid = mkClientIdentity u c
  initStore tmp cid
  kps <- replicateM n (fst <$> generateKeyPackage tmp cid kiLifetime)
  when (kiSetKey == SetKey) $
    do
      pk <-
        liftIO . flip spawn Nothing $
          cliCmd tmp cid ["public-key"]
      put
        ( brig
            . paths ["clients", toByteString' c]
            . zUser (qUnqualified u)
            . json defUpdateClient {updateClientMLSPublicKeys = Map.fromList [(Ed25519, pk)]}
        )
      !!! const 200
        === statusCode
  let upload = object ["key_packages" .= toJSON (map (Base64ByteString . raw) kps)]
  post
    ( brig
        . paths ["mls", "key-packages", "self", toByteString' c]
        . zUser (qUnqualified u)
        . json upload
    )
    !!! const (case kiSetKey of SetKey -> 201; DontSetKey -> 400)
      === statusCode

getKeyPackageCount :: (HasCallStack) => Brig -> Qualified UserId -> ClientId -> Http KeyPackageCount
getKeyPackageCount brig u c =
  responseJsonError
    =<< get
      ( brig
          . paths ["mls", "key-packages", "self", toByteString' c, "count"]
          . queryItem "ciphersuite" "0x0001"
          . zUser (qUnqualified u)
      )
      <!! const 200
        === statusCode

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