{-# LANGUAGE RecordWildCards #-}

-- 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 Assets.Lib where

import Cassandra as C
import Cassandra.Settings as C
import Control.Lens
import qualified Data.Attoparsec.ByteString.Char8 as Atto (Parser)
import Data.ByteString.Conversion
import Data.Conduit
import qualified Data.Conduit.Combinators as Conduit
import Data.Id (UserId)
import qualified Data.Text.Encoding as T
import Data.Text.Strict.Lens
import Imports
import Options.Applicative
import System.IO (hPutStr)
import qualified System.Logger as Log
import Wire.API.Asset (AssetKey)

data Opts = Opts
  { cHost :: String,
    cPort :: Int,
    cKeyspace :: C.Keyspace
  }

sampleParser :: Parser Opts
sampleParser =
  Opts
    <$> strOption
      ( long "cassandra-host"
          <> short 's'
          <> metavar "HOST"
          <> help "Cassandra Host"
          <> value "localhost"
          <> showDefault
      )
    <*> option
      auto
      ( long "cassandra-port"
          <> short 'p'
          <> metavar "PORT"
          <> help "Cassandra Port"
          <> value 9042
          <> showDefault
      )
    <*> ( C.Keyspace . view packed
            <$> strOption
              ( long "cassandra-keyspace"
                  <> short 'k'
                  <> metavar "STRING"
                  <> help "Cassandra Keyspace"
                  <> value "brig_test"
                  <> showDefault
              )
        )

main :: IO ()
main = do
  putStrLn "starting to read users ..."
  opts <- execParser (info (helper <*> sampleParser) desc)
  logger <- initLogger
  client <- initCas opts logger
  res <- process client
  putStrLn "\n"
  print res
  where
    initLogger = Log.new . Log.setOutput Log.StdOut . Log.setFormat Nothing . Log.setBufSize 0 $ Log.defSettings
    initCas Opts {..} l =
      C.init
        . C.setLogger (C.mkLogger Nothing l)
        . C.setContacts cHost []
        . C.setPortNumber (fromIntegral cPort)
        . C.setKeyspace cKeyspace
        . C.setProtocolVersion C.V4
        $ C.defSettings
    desc = header "assets" <> progDesc "find invalid asset keys in cassandra brig" <> fullDesc

selectUsersAll :: C.PrepQuery C.R () UserRow
selectUsersAll = "SELECT id, assets FROM user"

readUsers :: ClientState -> ConduitM () [UserRow] IO ()
readUsers client =
  transPipe (runClient client) $
    paginateC selectUsersAll (paramsP LocalQuorum () 500) x5

process :: ClientState -> IO Result
process client =
  runConduit $
    readUsers client
      .| Conduit.mapM (\chunk -> hPutStr stderr "." $> chunk)
      .| Conduit.concat
      .| Conduit.foldMap checkAssets
  where
    isInvalid :: AssetText -> Bool
    isInvalid asset = isLeft $ runParser (parser :: Atto.Parser AssetKey) $ T.encodeUtf8 (txtAssetKey asset)

    checkAssets :: UserRow -> Result
    checkAssets (_, Nothing) = Result 1 0 []
    checkAssets (_, Just []) = Result 1 0 []
    checkAssets row@(_, Just assets) = if any isInvalid assets then Result 0 0 [row] else Result 0 1 []

type UserRow = (UserId, Maybe [AssetText])

data Result = Result
  { noAsset :: Int,
    validAsset :: Int,
    invalidAsset :: [UserRow]
  }
  deriving stock (Eq, Generic)

newtype AssetText = ImageAssetText
  { txtAssetKey :: Text
  }
  deriving stock (Eq, Generic)

instance Show AssetText where
  show (ImageAssetText ak) = show ak

instance Cql AssetText where
  ctype =
    Tagged
      ( UdtColumn
          "asset"
          [ ("typ", IntColumn),
            ("key", TextColumn)
          ]
      )

  fromCql (CqlUdt fs) = do
    t <- required "typ"
    k <- required "key"
    case (t :: Int32) of
      0 -> pure $! ImageAssetText k
      _ -> Left $ "unexpected user asset type: " ++ show t
    where
      required :: (Cql r) => Text -> Either String r
      required f =
        maybe
          (Left ("Asset: Missing required field '" ++ show f ++ "'"))
          fromCql
          (lookup f fs)
  fromCql _ = Left "UserAsset: UDT expected"

  -- Note: Order must match up with the 'ctype' definition.
  toCql (ImageAssetText k) =
    CqlUdt
      [ ("typ", CqlInt 0),
        ("key", toCql k)
      ]

instance Show Result where
  show (Result n v i) =
    "num_no_assets: "
      <> show n
      <> "\nnum_valid_assets: "
      <> show v
      <> "\nnum_invalid_assets: "
      <> show (length i)
      <> "\ninvalid_assets:\n"
      <> concatMap showRow i
    where
      showRow (uid, Nothing) = "  - user_id: " <> show uid <> "\n"
      showRow (uid, Just as) = "  - user_id: " <> show uid <> "\n" <> showAssets as
      showAsset a = "    key: " <> show (txtAssetKey a) <> "\n"
      showAssets assets = concatMap showAsset assets

instance Semigroup Result where
  (<>) (Result n1 v1 i1) (Result n2 v2 i2) =
    Result (n1 + n2) (v1 + v2) (i1 <> i2)

instance Monoid Result where
  mempty = Result 0 0 []
  mappend = (<>)
