-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2026 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 Wire.IdPSubsystem.InterpreterSpec (spec) where

import Control.Lens hiding (elements)
import Data.ByteString.Lazy.UTF8 qualified as BSUTF8
import Data.Default (def)
import Data.Handle
import Data.HavePendingInvitations
import Data.Id
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as Map
import Data.Qualified
import Data.Set qualified as Set
import Imports
import Polysemy
import Polysemy.Error
import Polysemy.State
import SAML2.WebSSO qualified as SAML
import SAML2.WebSSO.Types (IdPMetadata (_edIssuer))
import System.Logger.Message qualified as Log
import Test.Hspec
import Test.Hspec.QuickCheck
import Test.QuickCheck
import Wire.API.Team.Member
import Wire.API.User
import Wire.API.User.IdentityProvider
import Wire.BrigAPIAccess
import Wire.GalleyAPIAccess
import Wire.IdPConfigStore
import Wire.IdPConfigStore.Mem (idPToMem)
import Wire.IdPSubsystem
import Wire.IdPSubsystem.Interpreter
import Wire.MockInterpreters.GalleyAPIAccess
import Wire.Sem.Logger qualified as Log
import Wire.Sem.Logger.TinyLog

type AllEffects =
  [ IdPSubsystem,
    Error IdPSubsystemError,
    Log.Logger (Log.Msg -> Log.Msg),
    BrigAPIAccess,
    GalleyAPIAccess,
    IdPConfigStore,
    State [RecordedLog]
  ]

runAllEffects ::
  Bool ->
  [IdP] ->
  Map TeamId [TeamMember] ->
  ([UserId] -> [Handle] -> [EmailAddress] -> HavePendingInvitations -> [User]) ->
  Sem AllEffects a ->
  (Either IdPSubsystemError a, [RecordedLog])
runAllEffects enableDiscovery idps teams brigAPIMockFn action = swap $ run $ runState [] $ do
  (_idpState, result) <- idPToMem $ do
    forM_ idps insertConfig
    -- Run the action
    miniGalleyAPIAccess teams def
      . brigAPIAccessMock brigAPIMockFn
      . pureRecordLogs
      . runError
      . interpretIdPSubsystem enableDiscovery
      $ action
  pure result

-- Unfortunately, the implementation of `GetUsersByVariousKeys` is too complex
-- to cover it by a simple state. Thus, we resort to handwritten mock
-- functions. The integration tests provide certainty that
-- `GetUsersByVariousKeys` works for us as expected.
brigAPIAccessMock ::
  ([UserId] -> [Handle] -> [EmailAddress] -> HavePendingInvitations -> [User]) ->
  InterpreterFor BrigAPIAccess r
brigAPIAccessMock mockFn = interpret $ \case
  GetUsersByVariousKeys uids handles emails havePending -> pure $ mockFn uids handles emails havePending
  _ -> error "Not implemented in brigAPIAccessMock"

brigAPIAccessMockFn ::
  EmailAddress ->
  User ->
  [UserId] ->
  [Handle] ->
  [EmailAddress] ->
  HavePendingInvitations ->
  [User]
brigAPIAccessMockFn expectedEmail resUser [] [] emails NoPendingInvitations
  | emails == [expectedEmail] =
      [resUser]
brigAPIAccessMockFn expectedEmail _resUser [] [] emails NoPendingInvitations
  | emails /= [expectedEmail] =
      []
brigAPIAccessMockFn expectedEmail resUser uids handles emails havePending =
  error $
    "Exepected call "
      <> show (expectedEmail, resUser, uids, handles, emails, havePending)

spec :: Spec
spec = describe "IdPSubsystem.Interpreter" $ do
  describe "getSsoCodeByEmail" $ do
    prop "returns IdP for SCIM user with single IdP only if the domain fits" $ \(teamMember :: TeamMember) user userRef email teamId mbDomain -> do
      idp <- generate $ do
        -- Generate IdPs with and without fitting domain
        idp' <- arbitrary <&> SAML.idpExtraInfo . domain .~ mbDomain
        more <- arbitrary
        elements $ idp' : more

      let userWithEmail =
            user
              { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                userEmailUnvalidated = Just email,
                userTeam = Just teamId,
                userManagedBy = ManagedByScim
              }

          idpWithTeam = idp & SAML.idpExtraInfo . team .~ teamId

          teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

          teams = Map.singleton teamId [teamMember']

          (result, logs) =
            runAllEffects
              True
              [idpWithTeam]
              teams
              (brigAPIAccessMockFn email userWithEmail)
              (getSsoCodeByEmail mbDomain email)

          expectedResult =
            if (idp ^. SAML.idpExtraInfo . domain) == mbDomain
              then Right (Just idp._idpId)
              else Right Nothing

      result `shouldBe` expectedResult
      expectedSevereLogs logs mempty

    prop "finds IdP for SCIM user by domain" $ \(teamMember :: TeamMember) user idp userRef email teamId dom otherIdPs ->
      not
        ( any
            ( \otherIdP ->
                Just dom == otherIdP._idpExtraInfo._domain
                  || idp._idpId == otherIdP._idpId
                  -- In reality, this constraint is enforced by `validateNewIdP`, called close to the endpoint.
                  -- Depending on the IdP API version, the issuer must be either unique per backend (V1) or per team (V2)
                  || idp._idpMetadata._edIssuer == otherIdP._idpMetadata._edIssuer
            )
            otherIdPs
        )
        ==> let userWithEmail =
                  user
                    { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                      userEmailUnvalidated = Just email,
                      userTeam = Just teamId,
                      userManagedBy = ManagedByScim
                    }

                -- Set the target multi-ingress domain for one IdP
                patchedIdp = idp & SAML.idpExtraInfo . domain ?~ dom
                patchedIdpsHead = patchedIdp NE.<| otherIdPs
                idpsWithTeamHead = patchedIdpsHead <&> SAML.idpExtraInfo . team .~ teamId

                teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

                teams = Map.singleton teamId [teamMember']

                (resultHead, logsHead) =
                  runAllEffects
                    True
                    (NE.toList idpsWithTeamHead)
                    teams
                    (brigAPIAccessMockFn email userWithEmail)
                    (getSsoCodeByEmail (Just dom) email)
             in do
                  resultHead `shouldBe` Right (Just idp._idpId)
                  expectedSevereLogs logsHead mempty

                  -- Same game, but this time the IdP to find is at the end of the list.
                  -- This ensures that the IdP is really looked up and not just taken from
                  -- the head of the list.
                  let patchedIdpsTail = NE.reverse $ patchedIdp NE.<| (NE.reverse otherIdPs)
                      idpsWithTeamTail = patchedIdpsTail <&> SAML.idpExtraInfo . team .~ teamId

                      (resultTail, logsTail) =
                        runAllEffects
                          True
                          (NE.toList idpsWithTeamTail)
                          teams
                          (brigAPIAccessMockFn email userWithEmail)
                          (getSsoCodeByEmail (Just dom) email)

                  resultTail `shouldBe` Right (Just idp._idpId)
                  expectedSevereLogs logsTail mempty

    prop "returns any IdP if there are multiple" $ \(teamMember :: TeamMember) user userRef email teamId mbDomain idps ->
      -- This should not happen, because the IdP management API allows to
      -- create only one IdP per domain. However, better not have undefined
      -- behaviour - Just in case...

      allUniqueAndAtLeast 2 idps ==>
        let userWithEmail =
              user
                { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                  userEmailUnvalidated = Just email,
                  userTeam = Just teamId,
                  userManagedBy = ManagedByScim
                }

            idpsWithTeam =
              idps
                <&> SAML.idpExtraInfo . domain .~ mbDomain
                <&> SAML.idpExtraInfo . team .~ teamId

            teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

            teams = Map.singleton teamId [teamMember']

            (result, logs) =
              runAllEffects
                True
                (NE.toList idpsWithTeam)
                teams
                (brigAPIAccessMockFn email userWithEmail)
                (getSsoCodeByEmail mbDomain email)

            expectedIdPIds :: [SAML.IdPId] = NE.toList $ ((SAML._idpId) <$> idpsWithTeam)
         in do
              result
                `shouldSatisfy` ( \case
                                    Right (Just x) -> x `elem` expectedIdPIds
                                    e -> error $ "Unexpected result " <> show e
                                )

              expectedSevereLogs logs [(Log.Warn, "Found more than one IdP config for domain")]

    prop "returns Nothing for an unknown email" $ \(teamMember :: TeamMember) user idp userRef email anotherEmail teamId mbDomain -> do
      let userWithEmail =
            user
              { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                userEmailUnvalidated = Just email,
                userTeam = Just teamId,
                userManagedBy = ManagedByScim
              }

          idpWithTeam = idp & SAML.idpExtraInfo . team .~ teamId

          teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

          teams = Map.singleton teamId [teamMember']

          (result, logs) =
            runAllEffects
              True
              [idpWithTeam]
              teams
              (brigAPIAccessMockFn email userWithEmail)
              (getSsoCodeByEmail mbDomain anotherEmail)

      result `shouldBe` Right Nothing
      expectedSevereLogs logs mempty

    prop "returns Nothing for SCIM user with no IdP" $ \(teamMember :: TeamMember) user userRef email teamId -> do
      let userWithEmail =
            user
              { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                userEmailUnvalidated = Just email,
                userTeam = Just teamId,
                userManagedBy = ManagedByScim
              }

          teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

          teams = Map.singleton teamId [teamMember']

          (result, logs) =
            runAllEffects
              True
              mempty
              teams
              (brigAPIAccessMockFn email userWithEmail)
              (getSsoCodeByEmail Nothing email)

      result `shouldBe` Right Nothing
      expectedSevereLogs logs mempty

    prop "returns Nothing when the user does not belong to a team" $ \user idp userRef email -> do
      let userWithEmail =
            user
              { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                userEmailUnvalidated = Just email,
                userTeam = Nothing,
                userManagedBy = ManagedByScim
              }

          (result, logs) =
            runAllEffects
              True
              [idp]
              mempty
              (brigAPIAccessMockFn email userWithEmail)
              (getSsoCodeByEmail Nothing email)

      result `shouldBe` Right Nothing
      expectedSevereLogs logs mempty

    prop "returns Nothing for non SSO user" $ \(teamMember :: TeamMember) user idp email teamId -> do
      userManagedBy <- generate (arbitrary :: Gen ManagedBy)
      userIdentity <- generate (oneof [pure Nothing, pure $ Just (EmailIdentity email)])

      let userWithEmail =
            user
              { userIdentity = userIdentity,
                userEmailUnvalidated = Just email,
                userTeam = Just teamId,
                userManagedBy = userManagedBy
              }

          idpWithTeam = idp & SAML.idpExtraInfo . team .~ teamId

          teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

          teams = Map.singleton teamId [teamMember']

          (result, logs) =
            runAllEffects
              True
              [idpWithTeam]
              teams
              (brigAPIAccessMockFn email userWithEmail)
              (getSsoCodeByEmail Nothing email)

      result `shouldBe` Right Nothing
      expectedSevereLogs logs mempty

    prop "returns Nothing when the feature is disabled" $ \(teamMember :: TeamMember) user idp userRef email teamId -> do
      let userWithEmail =
            user
              { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                userEmailUnvalidated = Just email,
                userTeam = Just teamId,
                userManagedBy = ManagedByScim
              }

          idpWithTeam = idp & SAML.idpExtraInfo . team .~ teamId

          teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

          teams = Map.singleton teamId [teamMember']

          (result, logs) =
            runAllEffects
              False
              [idpWithTeam]
              teams
              (brigAPIAccessMockFn email userWithEmail)
              (getSsoCodeByEmail Nothing email)

      result `shouldBe` Right Nothing
      expectedSevereLogs logs mempty

    prop "returns Nothing if there are multiple, but none for this domain" \(teamMember :: TeamMember) user userRef email teamId mbDomain mbAnotherDomain idps ->
      allUniqueAndAtLeast 2 idps && mbAnotherDomain /= mbDomain ==>
        let userWithEmail =
              user
                { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                  userEmailUnvalidated = Just email,
                  userTeam = Just teamId,
                  userManagedBy = ManagedByScim
                }

            idpsWithTeam =
              idps
                <&> SAML.idpExtraInfo . domain .~ mbDomain
                <&> SAML.idpExtraInfo . team .~ teamId

            teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

            teams = Map.singleton teamId [teamMember']

            (result, logs) =
              runAllEffects
                True
                (NE.toList idpsWithTeam)
                teams
                (brigAPIAccessMockFn email userWithEmail)
                (getSsoCodeByEmail mbAnotherDomain email)
         in do
              result `shouldBe` Right Nothing
              expectedSevereLogs logs mempty

    prop "getting multiple users for an email leads to exception" $ \(teamMember :: TeamMember) user anotherUser userRef email teamId mbDomain -> do
      idp' <- generate $ do
        idp'' <- arbitrary <&> SAML.idpExtraInfo . domain .~ mbDomain
        more <- arbitrary
        elements $ idp'' : more

      let userWithEmail =
            user
              { userIdentity = Just (SSOIdentity (UserSSOId userRef) (Just email)),
                userEmailUnvalidated = Just email,
                userTeam = Just teamId,
                userManagedBy = ManagedByScim
              }

          idpWithTeam = idp' & SAML.idpExtraInfo . team .~ teamId

          teamMember' = teamMember & Wire.API.Team.Member.userId .~ (qUnqualified userWithEmail.userQualifiedId)

          teams = Map.singleton teamId [teamMember']

          brigAPIAccessMockFnMultipleUser :: [UserId] -> [Handle] -> [EmailAddress] -> HavePendingInvitations -> [User]
          brigAPIAccessMockFnMultipleUser _userIds _handles _emails _pendingInvitations = [userWithEmail, anotherUser]

          (result, logs) =
            runAllEffects
              True
              [idpWithTeam]
              teams
              brigAPIAccessMockFnMultipleUser
              (getSsoCodeByEmail mbDomain email)

      result `shouldBe` Left InconsistentUsers
      expectedSevereLogs logs [(Log.Warn, "Multiple users found for email address in getSsoCodeByEmail")]

-- | Similar to `RecordedLog`, but only carries the prefix of the rendered log message
type ExpectedLog = (Log.Level, LByteString)

expectedSevereLogs :: [RecordedLog] -> [ExpectedLog] -> Expectation
expectedSevereLogs recordedLogs expectedLogs =
  do
    let errorLogs = filter (\(lvl, _msg) -> lvl > Log.Info) recordedLogs
    length errorLogs `shouldBe` length expectedLogs
    forM_ (zip errorLogs expectedLogs) $ \((lvl, msg), (expectedLvl, expectedPrefix)) -> do
      lvl `shouldBe` expectedLvl
      BSUTF8.toString msg `shouldStartWith` BSUTF8.toString expectedPrefix

-- | Keep IdPs unique by ID and issuer so the in-memory store does not collapse them.
allUniqueAndAtLeast :: Int -> NE.NonEmpty (SAML.IdPConfig extra) -> Bool
allUniqueAndAtLeast n xs =
  length xs >= n
    && let ids = Set.fromList (SAML._idpId <$> NE.toList xs)
           issuers = Set.fromList ((^. SAML.idpMetadata . SAML.edIssuer) <$> NE.toList xs)
        in Set.size ids == length xs
             && Set.size issuers == length xs
