-- 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 Galley.API.Teams.Export (getTeamMembersCSV) where

import Control.Concurrent
import Control.Concurrent.Async qualified as Async
import Control.Lens (view, (^.))
import Control.Monad.Codensity
import Data.ByteString (toStrict)
import Data.ByteString.Builder
import Data.Csv
import Data.Handle
import Data.IORef (atomicModifyIORef, newIORef)
import Data.Id
import Data.Map qualified as Map
import Data.Qualified (Local, tUnqualified)
import Imports hiding (atomicModifyIORef, newEmptyMVar, newIORef, putMVar, readMVar, takeMVar, threadDelay, tryPutMVar)
import Polysemy
import Polysemy.Async
import Polysemy.Resource
import Wire.API.Error
import Wire.API.Error.Galley
import Wire.API.Routes.LowLevelStream (LowLevelStreamingBody)
import Wire.API.Team.Export
import Wire.API.Team.Member
import Wire.API.User (ScimUserInfo (suiCreatedOn), User (..))
import Wire.BrigAPIAccess
import Wire.Sem.Concurrency
import Wire.Sem.Concurrency.IO
import Wire.Sem.Paging qualified as E
import Wire.Sem.Paging.Cassandra (InternalPaging)
import Wire.SparAPIAccess qualified as Spar
import Wire.TeamMemberStore (TeamMemberStore, listTeamMembers)
import Wire.TeamSubsystem (TeamSubsystem)
import Wire.TeamSubsystem qualified as TeamSubsystem

-- | Cache of inviter handles.
--
-- This is used to make sure that inviters are only looked up once in brig,
-- even if they appear as inviters of several users in the team.
type InviterCache = IORef (Map UserId (MVar (Maybe Handle)))

lookupInviter ::
  (Member Resource r, Member BrigAPIAccess r, Member (Final IO) r) =>
  InviterCache ->
  UserId ->
  Sem r (Maybe Handle)
lookupInviter cache uid = flip onException ensureCache $ do
  empty <- embedFinal newEmptyMVar
  (cached, var) <-
    embedFinal $ atomicModifyIORef cache $ \m -> case Map.lookup uid m of
      Nothing -> (Map.insert uid empty m, (False, empty))
      Just v -> (m, (True, v))
  -- the cache did not contain this user, so write it in the corresponding MVar
  unless cached $ do
    u <- listToMaybe <$> getUsers [uid]
    embedFinal $ putMVar var (u >>= userHandle)
  -- at this point, we know that the MVar contains a value or some other thread
  -- is about to write one, so it is safe to just read from the MVar with a
  -- blocking call
  embedFinal $ readMVar var
  where
    -- this is run in case of errors to guarantee that other threads will never
    -- deadlock while reading the cache
    ensureCache = embedFinal $ do
      m <- readIORef cache
      for_ (Map.lookup uid m) $ \var ->
        tryPutMVar var Nothing

getUserRecord ::
  ( Member BrigAPIAccess r,
    Member Spar.SparAPIAccess r,
    Member (ErrorS TeamMemberNotFound) r,
    Member (Final IO) r,
    Member Resource r
  ) =>
  InviterCache ->
  TeamMember ->
  Sem r TeamExportUser
getUserRecord cache member = do
  let uid = member ^. userId
  export <- getUserExportData uid >>= noteS @TeamMemberNotFound
  mCreatedOn <- do
    let mFromInvitation = snd <$> member ^. invitation
    case mFromInvitation of
      Just ts -> pure $ Just ts
      Nothing -> suiCreatedOn <$> Spar.lookupScimUserInfo uid
  -- look up inviter handle from the cache
  let mInviterId = fst <$> member ^. invitation
  invitedBy <- join <$> traverse (lookupInviter cache) mInviterId
  pure
    export
      { tExportInvitedBy = invitedBy,
        tExportRole = permissionsRole . view permissions $ member,
        tExportCreatedOn = mCreatedOn
      }

-- | Export team info as a CSV, and stream it to the client.
--
-- We paginate through the team member list, then spawn a thread for each user
-- (out of a thread pool) in order to fetch information for that user from brig
-- and spar. Inviter IDs are resolved to handles via a brig request, then
-- stored in a cache so that they can be reused by subsequent requests.
getTeamMembersCSV ::
  forall r.
  ( Member BrigAPIAccess r,
    Member (ErrorS 'AccessDenied) r,
    Member (TeamMemberStore InternalPaging) r,
    Member (Final IO) r,
    Member Spar.SparAPIAccess r,
    Member TeamSubsystem r
  ) =>
  Local UserId ->
  TeamId ->
  Sem r LowLevelStreamingBody
getTeamMembersCSV lusr tid = do
  TeamSubsystem.internalGetTeamMember (tUnqualified lusr) tid >>= \case
    Nothing -> throwS @'AccessDenied
    Just member -> unless (member `hasPermission` DownloadTeamMembersCsv) $ throwS @'AccessDenied

  chan <- embedFinal newChan
  cache <- embedFinal $ newIORef mempty

  let encodeRow r = encodeDefaultOrderedByNameWith customEncodeOptions [r]
  let produceTeamExportUsers = do
        embedFinal $ writeChan chan (Just headerLine)
        E.withChunks (\mps -> listTeamMembers @InternalPaging tid mps maxBound) $
          \members -> unsafePooledForConcurrentlyN_ 8 members $ \member -> do
            mRecord <-
              runErrorS @TeamMemberNotFound $
                getUserRecord cache member
            let mRow = encodeRow <$> mRecord
            when (isJust mRow) $
              embedFinal $
                writeChan chan mRow

  -- In case an exception is thrown inside the producer thread, the response
  -- will not contain a correct error message, but rather be an http error such
  -- as 'InvalidChunkHeaders'. The exception however still reaches the
  -- middleware and is being tracked in logging and metrics.
  let producerThread =
        produceTeamExportUsers
          `finally` embedFinal (writeChan chan Nothing)

  asyncToIOFinal . resourceToIOFinal . unsafelyPerformConcurrency @_ @Unsafe $ do
    -- Here we should really capture the Wai continuation and run the finaliser
    -- after that. Unfortunately, this is not really possible with Servant,
    -- because the continuation is not exposed by the Handler monad. The best
    -- we can do is return a Codensity value with the correct finaliser, but
    -- that still leaves a short window between when the resource is acquired
    -- and when the finaliser is installed where the resource might be leaked.
    -- I don't have a good solution for that.
    bracketOnError
      (async producerThread)
      cancel
      $ \producer -> do
        pure $ do
          void $ Codensity $ \k -> do
            r <- k ()
            Async.cancel producer
            pure r
          pure $ \write flush -> do
            let go = do
                  readChan chan >>= \case
                    Nothing -> write "" >> flush
                    Just line -> write (byteString (toStrict line)) >> flush >> go
            go

headerLine :: LByteString
headerLine = encodeDefaultOrderedByNameWith (customEncodeOptions {encIncludeHeader = True}) ([] :: [TeamExportUser])

customEncodeOptions :: EncodeOptions
customEncodeOptions =
  EncodeOptions
    { encDelimiter = fromIntegral (ord ','),
      encUseCrLf = True, -- to be compatible with Mac and Windows
      encIncludeHeader = False, -- (so we can flush when the header is on the wire)
      encQuoting = QuoteAll
    }
