content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE QuasiQuotes #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# OPTIONS_GHC -fno-warn-type-defaults #-}
  
  -- 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 ParseSchema where
  
  import Data.Aeson (ToJSON, object, (.=))
  import Data.Text qualified as T
  import Data.Text.IO qualified as T
  import Data.Text.Lazy.IO qualified as TIO
  import Imports
  import Options.Applicative qualified as OA
  import System.Environment (withArgs)
  import System.Exit (exitFailure)
  import System.FilePath ((</>))
  import System.FilePath.Posix (takeDirectory)
  import System.IO (Handle)
  import System.Process (system)
  import Text.Megaparsec
  import Text.Megaparsec.Char
  import Text.Megaparsec.Char.Lexer qualified as L
  import Text.Mustache
  import Text.RawString.QQ
  
  type Parser = Parsec Void Text
  
  lexeme :: Parser a -> Parser a
  lexeme = L.lexeme (L.space space1 (L.skipLineComment ("--" :: Text)) empty)
  
  data Column = Column
    { colName :: Text,
      colType :: Text
    }
    deriving (Eq, Show, Ord)
  
  data CreateTable = CreateTable
    { ctKeyspace :: Text,
      ctTablename :: Text,
      ctColumns :: [Column]
    }
    deriving (Eq, Show, Ord)
  
  findCreateTable :: Text -> Text -> [CreateTable] -> CreateTable
  findCreateTable ks tn [] = error $ "bad CreateTable name, keyspace: " <> show (ks, tn)
  findCreateTable ks tn (ct@(CreateTable ks' tn' _) : cts) =
    if ks == ks' && tn == tn'
      then ct
      else findCreateTable ks tn cts
  
  name' :: Parser Text
  name' = T.pack <$> some (satisfy (\c -> isLower c || c == '_'))
  
  keyspace :: Parser Text
  keyspace = do
    n <- name'
    pure $ case T.splitOn "_test" n of
      [n', ""] -> n'
      _ -> n
  
  tablename :: Parser Text
  tablename = name'
  
  columnName :: Parser Text
  columnName = name'
  
  columnType :: Parser Text
  columnType = T.pack <$> some (satisfy (\c -> isLower c || c == '_' || c == '<' || c == '>'))
  
  column :: Parser Column
  column = do
    col <- Column <$> lexeme columnName <*> lexeme columnType
    void $ optional (lexeme "PRIMARY KEY")
    pure col
  
  betweenBrackets :: ParsecT Void Text Identity a -> ParsecT Void Text Identity a
  betweenBrackets = between (lexeme "(") (lexeme ")")
  
  primaryKeyAtEndOfColumns :: Parser ()
  primaryKeyAtEndOfColumns = do
    void $ lexeme "PRIMARY KEY"
    void $
      lexeme $
        betweenBrackets $
          (void name' <|> nested) `sepBy1` lexeme ","
    where
      nested = void $ betweenBrackets $ name' `sepBy1` lexeme ","
  
  createTable :: Parser CreateTable
  createTable = do
    void $ lexeme "CREATE TABLE"
    ks <- keyspace
    void "."
    tn <- lexeme tablename
    cols <-
      betweenBrackets $
        catMaybes
          <$> ((Just <$> column) <|> (Nothing <$ primaryKeyAtEndOfColumns))
            `sepBy` lexeme ","
    noise
    pure (CreateTable ks tn cols)
  
  noise :: Parser ()
  noise = void $ some (satisfy (/= ';'))
  
  stmt :: Parser (Maybe CreateTable)
  stmt =
    do
      Just <$> createTable
      <|> (Nothing <$ (notFollowedBy (string "CREATE TABLE") >> noise))
  
  parser :: Parser [CreateTable]
  parser = catMaybes <$> (stmt `sepEndBy` lexeme (string ";"))
  
  -------------------------------------------------------------------------------
  
  -- FUTUREWORK: the special cases in the first 6 lines should be handled by the parser, and this
  -- function should get at least a [Text] with all the qualifiers and the base type. Or
  -- something slightly more sophisticated.
  toHaskellType :: Text -> Text
  toHaskellType (T.split (`elem` ['<', '>']) -> ("frozen" : key')) = toHaskellType (mconcat key')
  toHaskellType (T.split (`elem` ['<', '>']) -> ("list" : key')) = "[" <> toHaskellType (mconcat key') <> "]"
  toHaskellType (T.split (`elem` ['<', '>']) -> ("set" : key')) = "(Cassandra.Set " <> toHaskellType (mconcat key') <> ")"
  toHaskellType (T.splitOn "frozen" -> ["", key']) = toHaskellType key'
  toHaskellType (T.splitOn "list" -> ["", key']) = "[" <> toHaskellType key' <> "]"
  toHaskellType (T.splitOn "set" -> ["", key']) = "(Cassandra.Set " <> toHaskellType key' <> ")"
  toHaskellType "ascii" = "Ascii"
  toHaskellType "asset" = "AssetIgnoreData" -- FUTUREWORK: check if we really dont need data of this type
  toHaskellType "bigint" = "Int64"
  toHaskellType "blob" = "Blob"
  toHaskellType "boolean" = "Bool"
  toHaskellType "double" = "Double"
  toHaskellType "float" = "Float"
  toHaskellType "permissions" = "Permissions"
  toHaskellType "inet" = "IP"
  toHaskellType "int" = "Int32"
  toHaskellType "pubkey" = "Blob"
  toHaskellType "text" = "Text"
  toHaskellType "timestamp" = "UTCTime"
  toHaskellType "timeuuid" = "TimeUuid"
  toHaskellType "uuid" = "UUID"
  toHaskellType st = error (T.unpack ("toHaskellType not implemented for " <> st))
  
  data Arguments = Arguments
    { argsCqlSchemaFile :: FilePath,
      argsCqlOutputFile :: Maybe FilePath
    }
  
  argParser :: OA.Parser Arguments
  argParser =
    Arguments
      <$> OA.argument OA.str (OA.metavar "SCHEMA_FILE")
      <*> optional
        ( OA.strOption
            ( OA.long "output"
                <> OA.short 'o'
                <> OA.metavar "FILE"
                <> OA.help "Write output to FILE"
            )
        )
  
  moduleTemplate :: Text
  moduleTemplate =
    [r|{-# LANGUAGE RecordWildCards #-}
  
  -- This file is part of the Wire Server implementation.
  --
  -- Copyright (C) 2020 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 Schema where
  
  import Cassandra
  import Common
  import Data.Conduit
  import Data.Id
  import Data.Time
  import Data.UUID
  import Wire.ConversationStore.Cassandra.Instances ()
  import Imports
  import Types
  import Wire.API.Team.Permission
  import Data.IP (IP)
  import Data.Handle (Handle)
  import Wire.API.User.Password (PasswordResetKey)
  import System.FilePath.Posix ((</>))
  
  -- This file was autogenerated by move-team-generate
  
  
  {{#chunkTable}}
  
  -- {{keySpace}}.{{tableName}}
  
  type Row{{keySpaceCaml}}{{tableNameCaml}} = ({{{typeOfRow}}})
  
  select{{keySpaceCaml}}{{tableNameCaml}} :: PrepQuery R (Identity ({{lookupKeyType}})) Row{{keySpaceCaml}}{{tableNameCaml}}
  select{{keySpaceCaml}}{{tableNameCaml}} = "SELECT {{columns}} FROM {{tableName}} WHERE {{lookupKeyWhere}}"
  
  read{{keySpaceCaml}}{{tableNameCaml}}:: Env -> {{lookupKeyType}} -> ConduitM () [Row{{keySpaceCaml}}{{tableNameCaml}}] IO ()
  read{{keySpaceCaml}}{{tableNameCaml}} Env {..} {{lookupKeyVar}} =
    transPipe (runClient env{{keySpaceCaml}}) $
      paginateC select{{keySpaceCaml}}{{tableNameCaml}} (paramsP LocalQuorum (pure {{lookupKeyVar}}) envPageSize) x5
  
  select{{keySpaceCaml}}{{tableNameCaml}}All :: PrepQuery R () Row{{keySpaceCaml}}{{tableNameCaml}}
  select{{keySpaceCaml}}{{tableNameCaml}}All = "SELECT {{columns}} FROM {{tableName}}"
  
  read{{keySpaceCaml}}{{tableNameCaml}}All :: Env -> ConduitM () [Row{{keySpaceCaml}}{{tableNameCaml}}] IO ()
  read{{keySpaceCaml}}{{tableNameCaml}}All Env {..} =
    transPipe (runClient env{{keySpaceCaml}}) $
      paginateC select{{keySpaceCaml}}{{tableNameCaml}}All (paramsP LocalQuorum () envPageSize) x5
  
  export{{keySpaceCaml}}{{tableNameCaml}}Full :: Env -> FilePath -> IO ()
  export{{keySpaceCaml}}{{tableNameCaml}}Full env@Env {..} path = do
    putStrLn $ "Exporting " <> "{{keySpace}}.{{tableName}}"  <> " to " <> path
    withBinaryFile path WriteMode $ \handle ->
      runConduit $
        read{{keySpaceCaml}}{{tableNameCaml}}All env
          .| sinkJsonLines handle
  
  insert{{keySpaceCaml}}{{tableNameCaml}} :: PrepQuery W Row{{keySpaceCaml}}{{tableNameCaml}} ()
  insert{{keySpaceCaml}}{{tableNameCaml}} =
      "INSERT INTO {{tableName}} ({{columns}}) VALUES {{placeHolders}}"
  
  import{{keySpaceCaml}}{{tableNameCaml}} :: Env -> FilePath -> IO ()
  import{{keySpaceCaml}}{{tableNameCaml}} Env{..} path = do
    exists <- doesFileExist path
    if exists then do
      putStrLn $ "Importing " <> path <> " to " <> "{{keySpace}}.{{tableName}}"
      withBinaryFile path ReadMode $ \handle -> do
        runConduit $
          sourceJsonLines handle
            .| transPipe (runClient env{{keySpaceCaml}}) (sinkTableRows insert{{keySpaceCaml}}{{tableNameCaml}})
    else do
       putStrLn $ "Skipping because not found: " <> path
       pure ()
  {{/chunkTable}}
  
  importAllTables :: Env -> IO ()
  importAllTables env@Env {..} = do
  {{#chunkTable}}
    import{{keySpaceCaml}}{{tableNameCaml}} env (envTargetPath </> "{{keySpace}}.{{tableName}}")
  {{/chunkTable}}
  
  exportAllTablesFull :: Env -> IO ()
  exportAllTablesFull env@Env {..} = do
  {{#chunkTable}}
    export{{keySpaceCaml}}{{tableNameCaml}}Full env (envTargetPath </> "{{keySpace}}.{{tableName}}")
  {{/chunkTable}}
  |]
  
  _printAllTables :: [CreateTable] -> IO ()
  _printAllTables createTables =
    for_ createTables $ \CreateTable {..} ->
      putStrLn $ T.unpack (ctKeyspace <> "." <> ctTablename)
  
  main :: IO ()
  main = do
    Arguments schemaFile mbOutFile <-
      OA.execParser
        ( OA.info
            (OA.helper <*> argParser)
            ( OA.header "gen-table-types"
                <> OA.progDesc "Generate Haskell types from Cassandra schema dump"
                <> OA.fullDesc
            )
        )
    contents <- T.readFile schemaFile
    createTables <- case parse parser "" contents of
      Left e -> putStr (errorBundlePretty e) >> exitFailure
      Right x -> pure x
  
    let mkChunk' = mkChunk createTables
  
        mkChunkUsers ks tn =
          mkChunk createTables ks tn "[UserId]" "uids" "user in ?"
  
        mkChunkTeam ks tn =
          mkChunk createTables ks tn "TeamId" "tid" "team = ?"
  
    withOutputHandle mbOutFile $ \out ->
      case compileMustacheText "" moduleTemplate of
        Left bundle -> putStrLn (errorBundlePretty bundle)
        Right template -> do
          (TIO.hPutStr out . renderMustache template)
            ( object
                [ "chunkTable"
                    .= [ -- brig.activation_keys,
                         -- brig.blacklist
                         -- brig.budget
                         -- brig.clients
                         --   PRIMARY KEY (user, client)
                         mkChunkUsers "brig" "clients",
                         -- brig.codes
                         -- brig.connection
                         --   PRIMARY KEY (left, right)
                         --   FUTUREWORK: is this symmetric? At first glance it looks like yes
                         mkChunk' "brig" "connection" "[UserId]" "uids" "left in ?",
                         -- brig.excluded_phones
                         --   PRIMARY KEY (mapped_id)
                         --   FUTUREWORK: is mapped_id a user id?
                         -- brig.invitation
                         --   mkChunk' "brig" "invitation" "[UserId]" "uids" "inviter in ?",
                         -- brig.invitation_info
                         --   mkChunk' "brig" "invitation_info" "[UserId]" "uids" "inviter in ?",
                         --   mkChunk createTables "brig" "invitee_info" "[UserId]" "uids" "invitee in ? OR inviter in ?" "([UserId], [UserId])" "(uids, uids)",
                         -- brig.login_codes
                         --   PRIMARY KEY (user)
                         --   FUTUREWORK: do we need this?
                         mkChunkUsers "brig" "login_codes",
                         -- brig.meta
                         -- brig.password_reset
                         --   PRIMARY KEY(key)
                         --   FUTUREWORK: do we need this? can we do better than a full table scan?
                         mkChunk' "brig" "password_reset" "[PasswordResetKey]" "reset_keys" "key in ?",
                         -- brig.prekeys
                         --   PRIMARY KEY (user, client, key)
                         --   FUTUREWORK: do new need this ?
                         mkChunkUsers "brig" "prekeys",
                         -- brig.properties
                         --   PRIMARY KEY (user, key)
                         mkChunkUsers "brig" "properties",
                         -- brig.provider
                         -- brig.provider_keys,
                         -- brig.rich_info
                         --   PRIMARY KEY (user)
                         mkChunkUsers "brig" "rich_info",
                         -- brig.service
                         -- brig.service_prefix
                         -- brig.service_tag
                         -- brig.service_team
                         -- brig.service_user
                         -- brig.service_whitelist
                         -- brig.service_whitelist_rev
                         -- brig.team_invitation
                         -- brig.team_invitation_email
                         -- brig.team_invitation_info
                         -- brig.unique_claims
                         -- brig.user
                         --   PRIMARY KEY (id)
                         mkChunk' "brig" "user" "[UserId]" "uids" "id in ?",
                         -- brig.user_cookies
                         -- brig.user_handle
                         --   PRIMARY KEY (handle :: text)
                         --   FUTUREWORK: can we do better than a full table scan?
                         mkChunk' "brig" "user_handle" "[Handle]" "handles" "handle in ?",
                         -- brig.user_keys:
                         --   PRIMARY KEY (key :: text) (Brig.Data.UserKey)
                         --   FUTUREWORK: do we need it?  can we do better than a full table scan?
                         mkChunk' "brig" "user_keys" "[Int32]" "keys" "key in ?",
                         -- brig.vcodes
                         -- galley.billing_team_member
                         --   PRIMARY KEY (team, user)
                         mkChunkTeam "galley" "billing_team_member",
                         -- galley.clients
                         --   PRIMARY KEY (user)
                         mkChunkUsers "galley" "clients",
                         -- galley.conversation
                         --   PRIMARY KEY (conv)
                         mkChunk' "galley" "conversation" "[ConvId]" "cids" "conv in ?",
                         -- galley.conversation_codes
                         -- galley.custom_backend
                         -- galley.data_migration
                         -- galley.id_mapping
                         -- galley.legalhold_pending_prekeys
                         -- galley.legalhold_pending_prekeys,
                         -- galley.legalhold_service,
                         -- galley.member
                         --   PRIMARY KEY (conv, user)
                         mkChunk' "galley" "member" "[ConvId]" "cids" "conv in ?",
                         -- galley.meta
                         -- galley.service
                         -- galley.team
                         --   PRIMARY KEY (team)
                         mkChunkTeam "galley" "team",
                         -- galley.team_conv
                         --   PRIMARY KEY (team, conv)
                         mkChunkTeam "galley" "team_conv",
                         -- galley.team_features
                         --   PRIMARY KEY (team_id)
                         mkChunk' "galley" "team_features" "TeamId" "tid" "team_id = ?",
                         -- galley.team_member
                         --   PRIMARY KEY (team, user)
                         mkChunkTeam "galley" "team_member",
                         -- galley.team_notifications
                         --   PRIMARY KEY (team, id)
                         mkChunkTeam "galley" "team_notifications",
                         -- galley.user
                         --   PRIMARY KEY (user, conv)
                         mkChunkUsers "galley" "user",
                         -- galley.user_team
                         --   PRIMARY KEY (user, team)
                         mkChunkUsers "galley" "user_team",
                         -- gundeck.met
                         --   PRIMARY KEY (user, id)
                         mkChunkUsers "gundeck" "notifications",
                         -- gundeck.push
                         -- gundeck.user_push
                         -- spar.authreq
                         -- spar.authresp
                         -- spar.bind_cookie
                         -- spar.default_idp
                         -- spar.idp
                         -- spar.idp_raw_metadata
                         -- spar.issuer_idp
                         -- spar.meta
                         -- spar.scim_external
                         mkChunkTeam "spar" "scim_external",
                         -- spar.scim_user_times
                         --   PRIMARY KEY (uid)
                         mkChunk' "spar" "scim_user_times" "[UserId]" "uids" "uid in ?",
                         -- spar.team_idp
                         -- spar.team_provisioning_by_team
                         -- spar.team_provisioning_by_token
                         -- spar.user
                         mkChunk' "spar" "user" "[Text]" "issuer" "issuer in ?"
                         -- PRIMARY KEY (issuer, sso_id)
                         -- spar.verdict
                       ]
                ]
            )
  
    for_ mbOutFile $ \filename ->
      system $ "ormolu --mode inplace " <> filename
    where
      withOutputHandle :: Maybe FilePath -> (Handle -> IO ()) -> IO ()
      withOutputHandle mbOutFile fn = case mbOutFile of
        Nothing -> fn stdout
        Just filename -> withFile filename WriteMode fn
  
  -- | FUTUREWORK: can be implemented more easily with `stack path --project-root`
  projectFile :: FilePath -> IO FilePath
  projectFile relativeFilename =
    (</> relativeFilename) <$> findProjectRoot
    where
      findFileUpwards :: FilePath -> FilePath -> IO (Maybe FilePath)
      findFileUpwards startDir =
        findFile (fmap (nthParent startDir) [0 .. 10])
        where
          nthParent :: FilePath -> Int -> FilePath
          nthParent dir i =
            foldl' (</>) dir (replicate i "..")
  
      findProjectRoot :: IO FilePath
      findProjectRoot = do
        curDir <- getCurrentDirectory
        Just p <- findFileUpwards curDir "stack-deps.nix"
        pure $ takeDirectory p
  
  debug :: IO ()
  debug = do
    cassandraSchema <- projectFile "cassandra-schema.cql"
    withArgs [cassandraSchema] main
  
  debugwrite :: IO ()
  debugwrite = do
    cassandraSchema <- projectFile "cassandra-schema.cql"
    outputFile <- projectFile "tools/db/move-team/src/Schema.hs"
    withArgs [cassandraSchema, "--output=" <> outputFile] main
  
  data Chunk = Chunk
    { lookupKeyType :: Text,
      lookupKeyVar :: Text,
      lookupKeyWhere :: Text,
      keySpace :: Text,
      tableName :: Text,
      keySpaceCaml :: Text,
      tableNameCaml :: Text,
      columns :: Text,
      columnsQuoted :: Text,
      typeOfRow :: Text,
      placeHolders :: Text
    }
    deriving (Generic)
  
  instance ToJSON Chunk
  
  mkChunk :: [CreateTable] -> Text -> Text -> Text -> Text -> Text -> Chunk
  mkChunk createTables ks tn lookupKeyType lookupKeyVar lookupKeyWhere =
    let (CreateTable _ks _tn cols) = findCreateTable ks tn createTables
        keySpace = ks
        tableName = tn
        keySpaceCaml = caml ks
        tableNameCaml = caml tn
        columns = T.intercalate ", " (fmap colName cols)
        columnsQuoted = T.intercalate ", " (fmap (quote . colName) cols)
        typeOfRow = T.intercalate ", " (fmap (("Maybe " <>) . toHaskellType . colType) cols)
        placeHolders = "(" <> T.intercalate ", " (replicate (length cols) "?") <> ")"
     in Chunk {..}
    where
      caml :: Text -> Text
      caml = T.pack . go 0 . T.unpack
        where
          go _ ('_' : c : cs) = toUpper c : go 1 cs
          go 0 (c : cs) = toUpper c : go 1 cs
          go _ (c : cs) = c : go 1 cs
          go _ "" = ""
  
      quote :: Text -> Text
      quote str = "\"" <> str <> "\""
  