#
# This file is licensed under the Affero General Public License (AGPL) version 3.
#
# Copyright (C) 2023 New Vector, Ltd
#
# 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.
#
# See the GNU Affero General Public License for more details:
# <https://www.gnu.org/licenses/agpl-3.0.html>.
#
# Originally licensed under the Apache License, Version 2.0:
# <http://www.apache.org/licenses/LICENSE-2.0>.
#
# [This file includes modifications made by New Vector Limited]
#
#
from typing import Any, Dict
from unittest.mock import Mock

from twisted.internet.testing import MemoryReactor

from synapse.api.constants import EventTypes
from synapse.rest import admin
from synapse.rest.client import login, room
from synapse.server import HomeServer
from synapse.types import JsonDict, create_requester
from synapse.util import Clock
from synapse.visibility import filter_events_for_client

from tests import unittest
from tests.unittest import override_config

one_hour_ms = 3600000
one_day_ms = one_hour_ms * 24


class RetentionTestCase(unittest.HomeserverTestCase):
    servlets = [
        admin.register_servlets,
        login.register_servlets,
        room.register_servlets,
    ]

    def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
        config = self.default_config()

        # merge this default retention config with anything that was specified in
        # @override_config
        retention_config = {
            "enabled": True,
            "default_policy": {
                "min_lifetime": one_day_ms,
                "max_lifetime": one_day_ms * 3,
            },
            "allowed_lifetime_min": one_day_ms,
            "allowed_lifetime_max": one_day_ms * 3,
        }
        retention_config.update(config.get("retention", {}))
        config["retention"] = retention_config

        self.hs = self.setup_test_homeserver(config=config)

        return self.hs

    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
        self.user_id = self.register_user("user", "password")
        self.token = self.login("user", "password")

        self.store = self.hs.get_datastores().main
        self.serializer = self.hs.get_event_client_serializer()
        self.clock = self.hs.get_clock()

    def test_retention_event_purged_with_state_event(self) -> None:
        """Tests that expired events are correctly purged when the room's retention policy
        is defined by a state event.
        """
        room_id = self.helper.create_room_as(self.user_id, tok=self.token)

        # Set the room's retention period to 2 days.
        lifetime = one_day_ms * 2
        self.helper.send_state(
            room_id=room_id,
            event_type=EventTypes.Retention,
            body={"max_lifetime": lifetime},
            tok=self.token,
        )

        self._test_retention_event_purged(room_id, one_day_ms * 1.5)

    def test_retention_event_purged_with_state_event_outside_allowed(self) -> None:
        """Tests that the server configuration can override the policy for a room when
        running the purge jobs.
        """
        room_id = self.helper.create_room_as(self.user_id, tok=self.token)

        # Set a max_lifetime higher than the maximum allowed value.
        self.helper.send_state(
            room_id=room_id,
            event_type=EventTypes.Retention,
            body={"max_lifetime": one_day_ms * 4},
            tok=self.token,
        )

        # Check that the event is purged after waiting for the maximum allowed duration
        # instead of the one specified in the room's policy.
        self._test_retention_event_purged(room_id, one_day_ms * 1.5)

        # Set a max_lifetime lower than the minimum allowed value.
        self.helper.send_state(
            room_id=room_id,
            event_type=EventTypes.Retention,
            body={"max_lifetime": one_hour_ms},
            tok=self.token,
        )

        # Check that the event is purged after waiting for the minimum allowed duration
        # instead of the one specified in the room's policy.
        self._test_retention_event_purged(room_id, one_day_ms * 0.5)

    def test_retention_event_purged_without_state_event(self) -> None:
        """Tests that expired events are correctly purged when the room's retention policy
        is defined by the server's configuration's default retention policy.
        """
        room_id = self.helper.create_room_as(self.user_id, tok=self.token)

        self._test_retention_event_purged(room_id, one_day_ms * 2)

    @override_config({"retention": {"purge_jobs": [{"interval": "5d"}]}})
    def test_visibility(self) -> None:
        """Tests that synapse.visibility.filter_events_for_client correctly filters out
        outdated events, even if the purge job hasn't got to them yet.

        We do this by setting a very long time between purge jobs.
        """
        store = self.hs.get_datastores().main
        storage_controllers = self.hs.get_storage_controllers()
        room_id = self.helper.create_room_as(self.user_id, tok=self.token)

        # Send a first event, which should be filtered out at the end of the test.
        resp = self.helper.send(room_id=room_id, body="1", tok=self.token)
        first_event_id = resp.get("event_id")
        assert isinstance(first_event_id, str)

        # Advance the time by 2 days. We're using the default retention policy, therefore
        # after this the first event will still be valid.
        self.reactor.advance(one_day_ms * 2 / 1000)

        # Send another event, which shouldn't get filtered out.
        resp = self.helper.send(room_id=room_id, body="2", tok=self.token)
        valid_event_id = resp.get("event_id")
        assert isinstance(valid_event_id, str)

        # Advance the time by another 2 days. After this, the first event should be
        # outdated but not the second one.
        self.reactor.advance(one_day_ms * 2 / 1000)

        # Fetch the events, and run filter_events_for_client on them
        events = self.get_success(
            store.get_events_as_list([first_event_id, valid_event_id])
        )
        self.assertEqual(2, len(events), "events retrieved from database")
        filtered_events = self.get_success(
            filter_events_for_client(
                storage_controllers,
                self.user_id,
                events,
            )
        )

        # We should only get one event back.
        self.assertEqual(len(filtered_events), 1, filtered_events)
        # That event should be the second, not outdated event.
        self.assertEqual(filtered_events[0].event_id, valid_event_id, filtered_events)

    def _test_retention_event_purged(self, room_id: str, increment: float) -> None:
        """Run the following test scenario to test the message retention policy support:

        1. Send event 1
        2. Increment time by `increment`
        3. Send event 2
        4. Increment time by `increment`
        5. Check that event 1 has been purged
        6. Check that event 2 has not been purged
        7. Check that state events that were sent before event 1 aren't purged.
        The main reason for sending a second event is because currently Synapse won't
        purge the latest message in a room because it would otherwise result in a lack of
        forward extremities for this room. It's also a good thing to ensure the purge jobs
        aren't too greedy and purge messages they shouldn't.

        Args:
            room_id: The ID of the room to test retention in.
            increment: The number of milliseconds to advance the clock each time. Must be
                defined so that events in the room aren't purged if they are `increment`
                old but are purged if they are `increment * 2` old.
        """
        # Get the create event to, later, check that we can still access it.
        message_handler = self.hs.get_message_handler()
        create_event = self.get_success(
            message_handler.get_room_data(
                create_requester(self.user_id), room_id, EventTypes.Create, state_key=""
            )
        )

        # Send a first event to the room. This is the event we'll want to be purged at the
        # end of the test.
        resp = self.helper.send(room_id=room_id, body="1", tok=self.token)

        expired_event_id = resp.get("event_id")
        assert expired_event_id is not None

        # Check that we can retrieve the event.
        expired_event = self.get_event(expired_event_id)
        self.assertEqual(
            expired_event.get("content", {}).get("body"), "1", expired_event
        )

        # Advance the time.
        self.reactor.advance(increment / 1000)

        # Send another event. We need this because the purge job won't purge the most
        # recent event in the room.
        resp = self.helper.send(room_id=room_id, body="2", tok=self.token)

        valid_event_id = resp.get("event_id")
        assert valid_event_id is not None

        # Advance the time again. Now our first event should have expired but our second
        # one should still be kept.
        self.reactor.advance(increment / 1000)

        # Check that the first event has been purged from the database, i.e. that we
        # can't retrieve it anymore, because it has expired.
        self.get_event(expired_event_id, expect_none=True)

        # Check that the event that hasn't expired can still be retrieved.
        valid_event = self.get_event(valid_event_id)
        self.assertEqual(valid_event.get("content", {}).get("body"), "2", valid_event)

        # Check that we can still access state events that were sent before the event that
        # has been purged.
        self.get_event(room_id, bool(create_event))

    def get_event(self, event_id: str, expect_none: bool = False) -> JsonDict:
        event = self.get_success(self.store.get_event(event_id, allow_none=True))

        if expect_none:
            self.assertIsNone(event)
            return {}

        assert event is not None

        time_now = self.clock.time_msec()
        serialized = self.get_success(self.serializer.serialize_event(event, time_now))

        return serialized


class RetentionNoDefaultPolicyTestCase(unittest.HomeserverTestCase):
    servlets = [
        admin.register_servlets,
        login.register_servlets,
        room.register_servlets,
    ]

    def default_config(self) -> Dict[str, Any]:
        config = super().default_config()

        retention_config = {
            "enabled": True,
        }

        # Update this config with what's in the default config so that
        # override_config works as expected.
        retention_config.update(config.get("retention", {}))
        config["retention"] = retention_config

        return config

    def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
        mock_federation_client = Mock(spec=["backfill"])

        self.hs = self.setup_test_homeserver(
            federation_client=mock_federation_client,
        )
        return self.hs

    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
        self.user_id = self.register_user("user", "password")
        self.token = self.login("user", "password")

    def test_no_default_policy(self) -> None:
        """Tests that an event doesn't get expired if there is neither a default retention
        policy nor a policy specific to the room.
        """
        room_id = self.helper.create_room_as(self.user_id, tok=self.token)

        self._test_retention(room_id)

    def test_state_policy(self) -> None:
        """Tests that an event gets correctly expired if there is no default retention
        policy but there's a policy specific to the room.
        """
        room_id = self.helper.create_room_as(self.user_id, tok=self.token)

        # Set the maximum lifetime to 35 days so that the first event gets expired but not
        # the second one.
        self.helper.send_state(
            room_id=room_id,
            event_type=EventTypes.Retention,
            body={"max_lifetime": one_day_ms * 35},
            tok=self.token,
        )

        self._test_retention(room_id, expected_code_for_first_event=404)

    @unittest.override_config({"retention": {"enabled": False}})
    def test_visibility_when_disabled(self) -> None:
        """Retention policies should be ignored when the retention feature is disabled."""
        room_id = self.helper.create_room_as(self.user_id, tok=self.token)

        self.helper.send_state(
            room_id=room_id,
            event_type=EventTypes.Retention,
            body={"max_lifetime": one_day_ms},
            tok=self.token,
        )

        resp = self.helper.send(room_id=room_id, body="test", tok=self.token)

        self.reactor.advance(one_day_ms * 2 / 1000)

        self.get_event(room_id, resp["event_id"])

    def _test_retention(
        self, room_id: str, expected_code_for_first_event: int = 200
    ) -> None:
        # Send a first event to the room. This is the event we'll want to be purged at the
        # end of the test.
        resp = self.helper.send(room_id=room_id, body="1", tok=self.token)

        first_event_id = resp.get("event_id")
        assert first_event_id is not None

        # Check that we can retrieve the event.
        expired_event = self.get_event(room_id, first_event_id)
        self.assertEqual(
            expired_event.get("content", {}).get("body"), "1", expired_event
        )

        # Advance the time by a month.
        self.reactor.advance(one_day_ms * 30 / 1000)

        # Send another event. We need this because the purge job won't purge the most
        # recent event in the room.
        resp = self.helper.send(room_id=room_id, body="2", tok=self.token)

        second_event_id = resp.get("event_id")
        assert second_event_id is not None

        # Advance the time by another month.
        self.reactor.advance(one_day_ms * 30 / 1000)

        # Check if the event has been purged from the database.
        first_event = self.get_event(
            room_id, first_event_id, expected_code=expected_code_for_first_event
        )

        if expected_code_for_first_event == 200:
            self.assertEqual(
                first_event.get("content", {}).get("body"), "1", first_event
            )

        # Check that the event that hasn't been purged can still be retrieved.
        second_event = self.get_event(room_id, second_event_id)
        self.assertEqual(second_event.get("content", {}).get("body"), "2", second_event)

    def get_event(
        self, room_id: str, event_id: str, expected_code: int = 200
    ) -> JsonDict:
        url = "/_matrix/client/r0/rooms/%s/event/%s" % (room_id, event_id)

        channel = self.make_request("GET", url, access_token=self.token)

        self.assertEqual(channel.code, expected_code, channel.result)

        return channel.json_body
