{-# OPTIONS_GHC -Wno-unused-top-binds #-}

-- 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 Spar.DataMigration.V2_UserV2 (migration) where

import Cassandra
import qualified Conduit as C
import qualified Data.ByteString.UTF8 as UTF8
import Data.Conduit
import qualified Data.Conduit.Combinators as CC
import Data.Conduit.Internal (zipSources)
import qualified Data.Conduit.List as CL
import Data.Id
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import Data.Time (UTCTime)
import Imports
import qualified SAML2.WebSSO as SAML
import SAML2.WebSSO.Types (Issuer (..))
import Spar.Data (NormalizedUNameID (unNormalizedUNameID), normalizeQualifiedNameId)
import Spar.Data.Instances ()
import Spar.DataMigration.Types (logger)
import Spar.DataMigration.Types hiding (logger)
import qualified System.Logger as Log
import URI.ByteString (serializeURIRef')
import Wire.IdPConfigStore.Orphans ()

-- row in spar.user
-- primary key are the first two entries
type OldRow = (SAML.Issuer, CqlSafe SAML.NameID, UserId)

-- row in spar.user_v2
-- primary key are the first two entries
type NewRow = (SAML.Issuer, NormalizedUNameID, SAML.NameID, UserId)

-- Roughly equivalent to a mapping NewRow->[OldRow]. If there are multiple
-- entries per key this means there is a collision of tow old rows to a new row
type MigrationMapInverse = Map (SAML.Issuer, NormalizedUNameID) [(SAML.NameID, UserId)]

type EnvIO = ReaderT Env IO

data List2 a = List2 a a [a]
  deriving (Show)

-- Resolve collisions in MigrationMapInverse
-- Left =~ collision could not be resolved
-- Right =~ collision could be resolved to 1 "winner"
type CollisionResolver =
  (SAML.Issuer, NormalizedUNameID) ->
  List2 (SAML.NameID, UserId) ->
  EnvIO (Either (List2 (SAML.NameID, UserId)) (SAML.NameID, UserId))

-- | Use this if you want to paginate without crashing
newtype CqlSafe a = CqlSafe {unCqlSafe :: Either String a}

instance (Cql a) => Cql (CqlSafe a) where
  ctype = Tagged $ untag (ctype @a)
  toCql _ = error "CqlSafe is not meant for serialization"
  fromCql val =
    case fromCql @a val of
      Left err -> Right $ CqlSafe (Left err)
      Right x -> Right $ CqlSafe (Right x)

migration :: Migration
migration =
  Migration
    { version = MigrationVersion 2,
      text = "Migrate to spar.user_v2",
      action = \env ->
        flip runReaderT env $
          case dryRun env of
            DryRun -> performDryRun
            NoDryRun -> performMigration
    }

performDryRun :: EnvIO ()
performDryRun = do
  migMapInv <- collectMapping
  runConduit $
    filterResolved combinedResolver migMapInv
      .| C.sinkNull

performMigration :: EnvIO ()
performMigration = do
  migMapInv <- collectMapping
  runConduit $
    filterResolved combinedResolver migMapInv
      .| CL.mapM_ insertNewRow
  where
    insertNewRow :: NewRow -> EnvIO ()
    insertNewRow newRow = do
      runSpar $ write insert (params LocalQuorum newRow)
      where
        insert :: PrepQuery W NewRow ()
        insert = "INSERT INTO user_v2 (issuer, normalized_uname_id, sso_id, uid) VALUES (?, ?, ?, ?)"

logInfo :: String -> EnvIO ()
logInfo msg = do
  logger_ <- asks logger
  Log.info logger_ (Log.msg msg)

logError :: String -> EnvIO ()
logError msg = do
  logger_ <- asks logger
  Log.err logger_ (Log.msg msg)

runSpar :: Client a -> EnvIO a
runSpar cl = do
  cass <- asks sparCassandra
  runClient cass cl

runBrig :: Client a -> EnvIO a
runBrig cl = do
  cass <- asks brigCassandra
  runClient cass cl

collectMapping :: EnvIO MigrationMapInverse
collectMapping = do
  runConduit $
    zipSources (CL.sourceList [(1 :: Int32) ..]) readOldRows
      .| migrateRows
      .| CC.foldl collect Map.empty
  where
    readOldRows :: ConduitT () [OldRow] EnvIO ()
    readOldRows = do
      pSize <- lift $ asks pageSize
      transPipe runSpar $
        paginateC select (paramsP LocalQuorum () pSize) x5
      where
        select :: PrepQuery R () OldRow
        select = "SELECT issuer, sso_id, uid FROM user"

    migrateRows :: ConduitT (Int32, [OldRow]) NewRow EnvIO ()
    migrateRows = do
      mb <- await
      case mb of
        Just (page, rows) -> do
          lift $ logInfo $ "page " <> show page
          for_ rows $ \row@(_, _, uid) -> do
            case migrateRow row of
              Left err -> do
                lift $ logError $ "Failed to parse row for user " <> show uid <> " with error " <> err
              Right newRow -> yield newRow
          migrateRows
        Nothing -> pure ()

    collect :: MigrationMapInverse -> NewRow -> MigrationMapInverse
    collect migMapInv (issuer, normNid, nid, uid') = Map.insertWith (<>) (issuer, normNid) [(nid, uid')] migMapInv

    migrateRow :: OldRow -> Either String NewRow
    migrateRow (issuer, nidSafe, uid) =
      unCqlSafe nidSafe <&> \nid -> (issuer, normalizeQualifiedNameId nid, nid, uid)

filterResolved :: CollisionResolver -> MigrationMapInverse -> ConduitT () NewRow EnvIO ()
filterResolved resolver migMapInv =
  C.yieldMany (Map.assocs migMapInv) .| go
  where
    go :: ConduitT ((SAML.Issuer, NormalizedUNameID), [(SAML.NameID, UserId)]) NewRow EnvIO ()
    go = do
      mbAssoc <- await
      for_ mbAssoc $ \(new@(issuer, nid), olds) -> do
        let yieldOld (nameId, uid) = yield (issuer, nid, nameId, uid)
        let issuerURI = UTF8.toString . serializeURIRef' . _fromIssuer $ issuer
        case olds of
          [] -> pure ()
          [old] -> yieldOld old
          (old1 : old2 : rest) ->
            lift (resolver new (List2 old1 old2 rest)) >>= \case
              Left _ ->
                lift $ logError $ unwords ["Couldnt resolve collisision of", issuerURI, T.unpack (unNormalizedUNameID nid), show olds]
              Right old -> do
                lift $ logInfo $ unwords ["Resolved collision", issuerURI, T.unpack (unNormalizedUNameID nid), show (fmap snd olds), "to", show (snd old)]
                yieldOld old
        go

combineResolver :: CollisionResolver -> CollisionResolver -> CollisionResolver
combineResolver resolver1 resolver2 pair olds =
  resolver1 pair olds >>= \case
    Left olds' -> resolver2 pair olds'
    Right old -> pure (Right old)

resolveViaActivated :: CollisionResolver
resolveViaActivated _ input@(List2 old1 old2 rest) = do
  olds' <- filterM (isActivated . snd) (old1 : old2 : rest)
  case olds' of
    [] -> pure (Left input)
    [old] -> pure (Right old)
    (old1' : old2' : rest') -> pure (Left (List2 old1' old2' rest'))
  where
    isActivated :: UserId -> EnvIO Bool
    isActivated u =
      runBrig $
        (== Just (Identity True))
          <$> retry x1 (query1 activatedSelect (params LocalQuorum (Identity u)))

    activatedSelect :: PrepQuery R (Identity UserId) (Identity Bool)
    activatedSelect = "SELECT activated FROM user WHERE id = ?"

-- Resolve to the user with the latest access token.
resolveViaAccessToken :: CollisionResolver
resolveViaAccessToken _ input@(List2 old1 old2 rest) = do
  tps <- for (old1 : old2 : rest) $ \old@(_nid, uid) -> do
    mt <- latestCookieExpiry uid
    pure (old, mt)
  case strictMaximumBy (compareCookieTime `on` snd) tps of
    Nothing -> pure (Left input)
    Just (old, _) -> pure (Right old)
  where
    latestCookieExpiry :: UserId -> EnvIO (Maybe UTCTime)
    latestCookieExpiry uid =
      runBrig $
        runIdentity <$$> query1 select (params LocalQuorum (Identity uid))
      where
        select :: PrepQuery R (Identity UserId) (Identity UTCTime)
        select = "SELECT expires FROM user_cookies WHERE user = ? ORDER BY expires DESC LIMIT 1"

    -- Returns the element (if it exists) that is stricly greater than any other element.
    strictMaximumBy :: (a -> a -> Ordering) -> [a] -> Maybe a
    strictMaximumBy _ [] = Nothing
    strictMaximumBy _ [y] = Just y
    strictMaximumBy ordering ys = case sortBy (flip ordering) ys of
      [] -> Nothing
      [y] -> Just y
      (y1 : y2 : _) -> case y1 `ordering` y2 of
        LT -> Nothing
        EQ -> Nothing
        GT -> Just y1

    compareCookieTime :: Maybe UTCTime -> Maybe UTCTime -> Ordering
    compareCookieTime Nothing Nothing = EQ
    compareCookieTime Nothing (Just _t) = LT
    compareCookieTime (Just _t) Nothing = GT
    compareCookieTime (Just t1) (Just t2) = t1 `compare` t2

combinedResolver :: CollisionResolver
combinedResolver = resolveViaActivated `combineResolver` resolveViaAccessToken
