{-# LANGUAGE TupleSections #-}
{-# OPTIONS_GHC -Wno-ambiguous-fields #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- 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 MlsStats.Run
  ( run,
  )
where

import Amazonka hiding (await)
import Amazonka.S3
import Amazonka.S3.CreateMultipartUpload
import Amazonka.S3.Lens
import Cassandra as C
import Cassandra.Settings as C
import Conduit
import Control.Exception
import Control.Lens ((.~), (?~), (^.))
import Data.Aeson qualified as A
import Data.ByteString.Base64 qualified as BS64
import Data.Conduit.Combinators hiding (foldMap, stderr, stdout)
import Data.Domain
import Data.Id
import Data.List.NonEmpty (nonEmpty)
import Data.Schema
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Data.Time
import Data.Time.Format.ISO8601
import Imports hiding (concat, filter, print)
import MlsStats.Options
import Network.HTTP.Types
import System.FilePath.Posix
import System.Logger qualified as Log
import Util.Options
import Wire.API.Conversation.Protocol
import Wire.API.MLS.Group

run :: Opts -> IO ()
run o = do
  logger' <- initLogger
  let settings = o.cassandraSettings
  galleyTables <- initCas settings.galleyHost settings.galleyPort settings.galleyKeyspace logger'
  brigTables <- initCas settings.brigHost settings.brigPort settings.brigKeyspace logger'
  runCommand o.s3Settings galleyTables brigTables o.cassandraSettings.pageSize
  where
    initLogger =
      Log.new
        . Log.setOutput Log.StdOut
        . Log.setFormat Nothing
        . Log.setBufSize 0
        $ Log.defSettings
    initCas casHost casPort casKeyspace l =
      C.init
        . C.setLogger (C.mkLogger Nothing l)
        . C.setContacts casHost []
        . C.setPortNumber (fromIntegral casPort)
        . C.setProtocolVersion C.V4
        . C.setKeyspace casKeyspace
        $ C.defSettings

runCommand :: S3Settings -> ClientState -> ClientState -> Int32 -> IO ()
runCommand s3 galleyTables brigTables queryPageSize = do
  logger <- newLogger Debug stderr
  let service =
        setEndpoint (s3.endpoint ^. awsSecure) (s3.endpoint ^. awsHost) (s3.endpoint ^. awsPort) defaultService
          & service_s3AddressingStyle .~ s3.addressingStyle
  env <-
    maybe id (\reg e -> (e :: Env) {region = Region' reg}) s3.region
      . (\e -> e {logger = logger})
      . configureService service
      <$> newEnv discover
  now <- formatShowM iso8601Format <$> getCurrentTime
  let upload =
        uploadStream env (BucketName s3.bucketName)
          . ObjectKey
          . T.pack
          . maybe id (</>) s3.bucketDir
          . maybe id (</>) now
  runResourceT $ do
    upload "user-client.csv" (userClient brigTables queryPageSize)
    upload "conv-group-team-protocol.csv" (convGroupTeamProtocol galleyTables queryPageSize)
    upload "domain-user-client-group.csv" (domainUserClientGroup galleyTables queryPageSize)
    upload "user-conv.csv" (userConv galleyTables queryPageSize)

userClient :: (MonadIO m) => ClientState -> Int32 -> ConduitT () ByteString m ()
userClient cassandra queryPageSize = do
  yield "user,client\r\n"
  ( transPipe
      (runClient cassandra)
      (paginateC userClientCql (paramsP LocalQuorum () queryPageSize) x1)
      .| concat
      .| mapC (\(u, c) -> T.encodeUtf8 $ T.pack (show u) <> "," <> clientToText c <> "\r\n")
    )
  where
    userClientCql :: PrepQuery R () (UserId, ClientId)
    userClientCql = "SELECT user, client FROM clients"

convGroupTeamProtocol :: (MonadIO m) => ClientState -> Int32 -> ConduitT () ByteString m ()
convGroupTeamProtocol cassandra queryPageSize = do
  yield "conversation,group,team,protocol\r\n"
  ( transPipe
      (runClient cassandra)
      (paginateC convGroupTeamProtocolCql (paramsP LocalQuorum () queryPageSize) x1)
      .| concat
      .| mapC (\(c, g, mt, p) -> fmap (c,g,,p) mt) -- filter out non-team conversations
      .| concat
      .| mapC
        ( \(c, g, t, p) ->
            T.encodeUtf8 (T.pack (show c))
              <> ","
              <> foldMap (BS64.encode . unGroupId) g
              <> ","
              <> T.encodeUtf8 (T.pack (show t))
              <> ","
              <> T.encodeUtf8 (convertProtocol p)
              <> "\r\n"
        )
    )
  where
    convGroupTeamProtocolCql :: PrepQuery R () (ConvId, Maybe GroupId, Maybe TeamId, Maybe ProtocolTag)
    convGroupTeamProtocolCql = "SELECT conv, group_id, team, protocol FROM conversation"
    convertProtocol :: Maybe ProtocolTag -> Text
    convertProtocol p = case schemaToJSON (fromMaybe ProtocolProteusTag p) of
      A.String s -> s
      _ -> "?"

domainUserClientGroup :: (MonadIO m) => ClientState -> Int32 -> ConduitT () ByteString m ()
domainUserClientGroup cassandra queryPageSize = do
  yield "user_domain,user,client,group\r\n"
  ( transPipe
      (runClient cassandra)
      (paginateC domainUserClientGroupCql (paramsP LocalQuorum () queryPageSize) x1)
      .| concat
      .| mapC
        ( \(d, u, c, g) ->
            (T.encodeUtf8 (domainText d))
              <> ","
              <> T.encodeUtf8 (T.pack (show u))
              <> ","
              <> T.encodeUtf8 (clientToText c)
              <> ","
              <> BS64.encode (unGroupId g)
              <> "\r\n"
        )
    )
  where
    domainUserClientGroupCql :: PrepQuery R () (Domain, UserId, ClientId, GroupId)
    domainUserClientGroupCql = "SELECT user_domain, user, client, group_id FROM mls_group_member_client"

userConv :: (MonadIO m) => ClientState -> Int32 -> ConduitT () ByteString m ()
userConv cassandra queryPageSize = do
  yield "user,conversation\r\n"
  ( transPipe
      (runClient cassandra)
      (paginateC userConvCql (paramsP LocalQuorum () queryPageSize) x1)
      .| concat
      .| mapC (\(u, c) -> T.encodeUtf8 $ T.pack (show u) <> "," <> T.pack (show c) <> "\r\n")
    )
  where
    userConvCql :: PrepQuery R () (UserId, ConvId)
    userConvCql = "SELECT user, conv FROM user"

uploadStream ::
  Env ->
  BucketName ->
  ObjectKey ->
  ConduitT () ByteString (ResourceT IO) () ->
  (ResourceT IO) ()
uploadStream env bucket key stream = do
  createMultipartResp <-
    sendEither env (newCreateMultipartUpload bucket key) >>= \case
      Left (ServiceError e) | e.status.statusCode == 404 && e.code == ErrorCode "NoSuchBucket" -> do
        void $ send env (newCreateBucket bucket)
        send env (newCreateMultipartUpload bucket key)
      Left e -> liftIO $ throwIO e
      Right resp -> pure resp
  let uploadId' = createMultipartResp ^. createMultipartUploadResponse_uploadId
  parts <-
    runConduit $
      stream
        .| chunksOfE chunkSize
        .| uploadParts env bucket key uploadId' 1
        .| mapC (uncurry newCompletedPart)
        .| sinkList
  void $
    send env $
      newCompleteMultipartUpload bucket key uploadId'
        & completeMultipartUpload_multipartUpload
          ?~ ( newCompletedMultipartUpload
                 & completedMultipartUpload_parts .~ nonEmpty parts
             )
  where
    chunkSize = 5 * 1024 * 1024

uploadParts ::
  Env ->
  BucketName ->
  ObjectKey ->
  Text ->
  Int ->
  ConduitT ByteString (Int, ETag) (ResourceT IO) ()
uploadParts env bucket key uploadId partNum = do
  chunkM <- await
  case chunkM of
    Just chunk -> do
      let req = newUploadPart bucket key partNum uploadId $ toBody chunk
      resp <- send env req
      for_ (resp ^. uploadPartResponse_eTag) $ \etag ->
        yield (partNum, etag)
      uploadParts env bucket key uploadId (partNum + 1)
    Nothing -> pure ()
