#
# 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]
#
#
import logging
from collections import defaultdict
from typing import Any, Dict, List, Optional, Set, Tuple

from twisted.internet.address import IPv4Address
from twisted.internet.protocol import Protocol, connectionDone
from twisted.python.failure import Failure
from twisted.test.proto_helpers import MemoryReactor
from twisted.web.resource import Resource

from synapse.app.generic_worker import GenericWorkerServer
from synapse.config.workers import InstanceTcpLocationConfig, InstanceUnixLocationConfig
from synapse.http.site import SynapseRequest, SynapseSite
from synapse.replication.http import ReplicationRestResource
from synapse.replication.tcp.client import ReplicationDataHandler
from synapse.replication.tcp.handler import ReplicationCommandHandler
from synapse.replication.tcp.protocol import (
    ClientReplicationStreamProtocol,
    ServerReplicationStreamProtocol,
)
from synapse.replication.tcp.resource import ReplicationStreamProtocolFactory
from synapse.server import HomeServer
from synapse.util import Clock

from tests import unittest
from tests.server import FakeTransport
from tests.utils import USE_POSTGRES_FOR_TESTS

try:
    import hiredis
except ImportError:
    hiredis = None  # type: ignore

logger = logging.getLogger(__name__)


class BaseStreamTestCase(unittest.HomeserverTestCase):
    """Base class for tests of the replication streams"""

    # hiredis is an optional dependency so we don't want to require it for running
    # the tests.
    if not hiredis:
        skip = "Requires hiredis"

    if not USE_POSTGRES_FOR_TESTS:
        # Redis replication only takes place on Postgres
        skip = "Requires Postgres"

    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
        # build a replication server
        server_factory = ReplicationStreamProtocolFactory(hs)
        self.streamer = hs.get_replication_streamer()
        self.server: ServerReplicationStreamProtocol = server_factory.buildProtocol(
            IPv4Address("TCP", "127.0.0.1", 0)
        )

        # Make a new HomeServer object for the worker
        self.reactor.lookups["testserv"] = "1.2.3.4"
        self.worker_hs = self.setup_test_homeserver(
            homeserver_to_use=GenericWorkerServer,
            config=self._get_worker_hs_config(),
            reactor=self.reactor,
            federation_http_client=None,
        )

        # Since we use sqlite in memory databases we need to make sure the
        # databases objects are the same.
        self.worker_hs.get_datastores().main.db_pool = hs.get_datastores().main.db_pool

        # Normally we'd pass in the handler to `setup_test_homeserver`, which would
        # eventually hit "Install @cache_in_self attributes" in tests/utils.py.
        # Unfortunately our handler wants a reference to the homeserver. That leaves
        # us with a chicken-and-egg problem.
        # We can workaround this: create the homeserver first, create the handler
        # and bodge it in after the fact. The bodging requires us to know the
        # dirty details of how `cache_in_self` works. We politely ask mypy to
        # ignore our dirty dealings.
        self.test_handler = self._build_replication_data_handler()
        self.worker_hs._replication_data_handler = self.test_handler  # type: ignore[attr-defined]

        repl_handler = ReplicationCommandHandler(self.worker_hs)
        self.client = ClientReplicationStreamProtocol(
            self.worker_hs,
            "client",
            "test",
            clock,
            repl_handler,
        )

        self._client_transport: Optional[FakeTransport] = None
        self._server_transport: Optional[FakeTransport] = None

    def create_resource_dict(self) -> Dict[str, Resource]:
        d = super().create_resource_dict()
        d["/_synapse/replication"] = ReplicationRestResource(self.hs)
        return d

    def _get_worker_hs_config(self) -> dict:
        config = self.default_config()
        config["worker_app"] = "synapse.app.generic_worker"
        config["instance_map"] = {"main": {"host": "testserv", "port": 8765}}
        return config

    def _build_replication_data_handler(self) -> "TestReplicationDataHandler":
        return TestReplicationDataHandler(self.worker_hs)

    def reconnect(self) -> None:
        if self._client_transport:
            self.client.close()

        if self._server_transport:
            self.server.close()

        self._client_transport = FakeTransport(self.server, self.reactor)
        self.client.makeConnection(self._client_transport)

        self._server_transport = FakeTransport(self.client, self.reactor)
        self.server.makeConnection(self._server_transport)

    def disconnect(self) -> None:
        if self._client_transport:
            self._client_transport = None
            self.client.close()

        if self._server_transport:
            self._server_transport = None
            self.server.close()

    def replicate(self) -> None:
        """Tell the master side of replication that something has happened, and then
        wait for the replication to occur.
        """
        self.streamer.on_notifier_poke()
        self.pump(0.1)

    def handle_http_replication_attempt(self) -> SynapseRequest:
        """Asserts that a connection attempt was made to the master HS on the
        HTTP replication port, then proxies it to the master HS object to be
        handled.

        Returns:
            The request object received by master HS.
        """

        # We should have an outbound connection attempt.
        clients = self.reactor.tcpClients
        self.assertEqual(len(clients), 1)
        (host, port, client_factory, _timeout, _bindAddress) = clients.pop(0)
        self.assertEqual(host, "1.2.3.4")
        self.assertEqual(port, 8765)

        # Set up client side protocol
        client_address = IPv4Address("TCP", "127.0.0.1", 1234)
        client_protocol = client_factory.buildProtocol(("127.0.0.1", 1234))

        # Set up the server side protocol
        server_address = IPv4Address("TCP", host, port)
        channel = self.site.buildProtocol((host, port))

        # hook into the channel's request factory so that we can keep a record
        # of the requests
        requests: List[SynapseRequest] = []
        real_request_factory = channel.requestFactory

        def request_factory(*args: Any, **kwargs: Any) -> SynapseRequest:
            request = real_request_factory(*args, **kwargs)
            requests.append(request)
            return request

        channel.requestFactory = request_factory

        # Connect client to server and vice versa.
        client_to_server_transport = FakeTransport(
            channel, self.reactor, client_protocol, server_address, client_address
        )
        client_protocol.makeConnection(client_to_server_transport)

        server_to_client_transport = FakeTransport(
            client_protocol, self.reactor, channel, client_address, server_address
        )
        channel.makeConnection(server_to_client_transport)

        # The request will now be processed by `self.site` and the response
        # streamed back.
        self.reactor.advance(0)

        # We tear down the connection so it doesn't get reused without our
        # knowledge.
        server_to_client_transport.loseConnection()
        client_to_server_transport.loseConnection()

        # there should have been exactly one request
        self.assertEqual(len(requests), 1)

        return requests[0]

    def assert_request_is_get_repl_stream_updates(
        self, request: SynapseRequest, stream_name: str
    ) -> None:
        """Asserts that the given request is a HTTP replication request for
        fetching updates for given stream.
        """

        path: bytes = request.path  # type: ignore
        self.assertRegex(
            path,
            rb"^/_synapse/replication/get_repl_stream_updates/%s/[^/]+$"
            % (stream_name.encode("ascii"),),
        )

        self.assertEqual(request.method, b"GET")


class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
    """Base class for tests running multiple workers.

    Enables Redis, providing a fake Redis server.

    Automatically handle HTTP replication requests from workers to master,
    unlike `BaseStreamTestCase`.
    """

    if not hiredis:
        skip = "Requires hiredis"

    if not USE_POSTGRES_FOR_TESTS:
        # Redis replication only takes place on Postgres
        skip = "Requires Postgres"

    def default_config(self) -> Dict[str, Any]:
        """
        Overrides the default config to enable Redis.
        Even if the test only uses make_worker_hs, the main process needs Redis
        enabled otherwise it won't create a Fake Redis server to listen on the
        Redis port and accept fake TCP connections.
        """
        base = super().default_config()
        base["redis"] = {"enabled": True}
        base["instance_map"] = {"main": {"host": "testserv", "port": 8765}}
        return base

    def setUp(self) -> None:
        super().setUp()

        # build a replication server
        self.streamer = self.hs.get_replication_streamer()

        # Fake in memory Redis server that servers can connect to.
        self._redis_server = FakeRedisPubSubServer()

        # We may have an attempt to connect to redis for the external cache already.
        self.connect_any_redis_attempts()

        store = self.hs.get_datastores().main
        self.database_pool = store.db_pool

        self.reactor.lookups["testserv"] = "1.2.3.4"
        self.reactor.lookups["localhost"] = "127.0.0.1"

        # A map from a HS instance to the associated HTTP Site to use for
        # handling inbound HTTP requests to that instance.
        self._hs_to_site = {self.hs: self.site}

        # Handle attempts to connect to fake redis server.
        self.reactor.add_tcp_client_callback(
            "localhost",
            6379,
            self.connect_any_redis_attempts,
        )

        self.hs.get_replication_command_handler().start_replication(self.hs)

        # When we see a connection attempt to the master replication listener we
        # automatically set up the connection. This is so that tests don't
        # manually have to go and explicitly set it up each time (plus sometimes
        # it is impossible to write the handling explicitly in the tests).
        #
        # Register the master replication listener:
        self.reactor.add_tcp_client_callback(
            "1.2.3.4",
            8765,
            lambda: self._handle_http_replication_attempt(self.hs, 8765),
        )

    def create_test_resource(self) -> ReplicationRestResource:
        """Overrides `HomeserverTestCase.create_test_resource`."""
        # We override this so that it automatically registers all the HTTP
        # replication servlets, without having to explicitly do that in all
        # subclassses.

        resource = ReplicationRestResource(self.hs)

        for servlet in self.servlets:
            servlet(self.hs, resource)

        return resource

    def make_worker_hs(
        self, worker_app: str, extra_config: Optional[dict] = None, **kwargs: Any
    ) -> HomeServer:
        """Make a new worker HS instance, correctly connecting replication
        stream to the master HS.

        Args:
            worker_app: Type of worker, e.g. `synapse.app.generic_worker`.
            extra_config: Any extra config to use for this instances.
            **kwargs: Options that get passed to `self.setup_test_homeserver`,
                useful to e.g. pass some mocks for things like `federation_http_client`

        Returns:
            The new worker HomeServer instance.
        """

        config = self._get_worker_hs_config()
        config["worker_app"] = worker_app
        config.update(extra_config or {})

        worker_hs = self.setup_test_homeserver(
            homeserver_to_use=GenericWorkerServer,
            config=config,
            reactor=self.reactor,
            **kwargs,
        )

        # If the instance is in the `instance_map` config then workers may try
        # and send HTTP requests to it, so we register it with
        # `_handle_http_replication_attempt` like we do with the master HS.
        instance_name = worker_hs.get_instance_name()
        instance_loc = worker_hs.config.worker.instance_map.get(instance_name)
        if instance_loc and isinstance(instance_loc, InstanceTcpLocationConfig):
            # Ensure the host is one that has a fake DNS entry.
            if instance_loc.host not in self.reactor.lookups:
                raise Exception(
                    "Host does not have an IP for instance_map[%r].host = %r"
                    % (
                        instance_name,
                        instance_loc.host,
                    )
                )

            # Copy the port into a new, non-Optional variable so mypy knows we're
            # not going to reset `instance_loc` to `None` under its feet. See
            # https://mypy.readthedocs.io/en/latest/common_issues.html#narrowing-and-inner-functions
            port = instance_loc.port

            self.reactor.add_tcp_client_callback(
                self.reactor.lookups[instance_loc.host],
                instance_loc.port,
                lambda: self._handle_http_replication_attempt(worker_hs, port),
            )
        elif instance_loc and isinstance(instance_loc, InstanceUnixLocationConfig):
            raise Exception(
                "Unix sockets are not supported for unit tests at this time."
            )

        store = worker_hs.get_datastores().main
        store.db_pool._db_pool = self.database_pool._db_pool

        # Set up a resource for the worker
        resource = ReplicationRestResource(worker_hs)

        for servlet in self.servlets:
            servlet(worker_hs, resource)

        self._hs_to_site[worker_hs] = SynapseSite(
            logger_name="synapse.access.http.fake",
            site_tag="{}-{}".format(
                worker_hs.config.server.server_name, worker_hs.get_instance_name()
            ),
            config=worker_hs.config.server.listeners[0],
            resource=resource,
            server_version_string="1",
            max_request_body_size=8192,
            reactor=self.reactor,
            hs=worker_hs,
        )

        worker_hs.get_replication_command_handler().start_replication(worker_hs)

        return worker_hs

    def _get_worker_hs_config(self) -> dict:
        config = self.default_config()
        return config

    def replicate(self) -> None:
        """Tell the master side of replication that something has happened, and then
        wait for the replication to occur.
        """
        self.streamer.on_notifier_poke()
        self.pump()

    def _handle_http_replication_attempt(self, hs: HomeServer, repl_port: int) -> None:
        """Handles a connection attempt to the given HS replication HTTP
        listener on the given port.
        """

        # We should have at least one outbound connection attempt, where the
        # last is one to the HTTP repication IP/port.
        clients = self.reactor.tcpClients
        self.assertGreaterEqual(len(clients), 1)
        (host, port, client_factory, _timeout, _bindAddress) = clients.pop()
        self.assertEqual(host, "1.2.3.4")
        self.assertEqual(port, repl_port)

        # Set up client side protocol
        client_address = IPv4Address("TCP", "127.0.0.1", 1234)
        client_protocol = client_factory.buildProtocol(("127.0.0.1", 1234))

        # Set up the server side protocol
        server_address = IPv4Address("TCP", host, port)
        channel = self._hs_to_site[hs].buildProtocol((host, port))

        # Connect client to server and vice versa.
        client_to_server_transport = FakeTransport(
            channel, self.reactor, client_protocol, server_address, client_address
        )
        client_protocol.makeConnection(client_to_server_transport)

        server_to_client_transport = FakeTransport(
            client_protocol, self.reactor, channel, client_address, server_address
        )
        channel.makeConnection(server_to_client_transport)

        # Note: at this point we've wired everything up, but we need to return
        # before the data starts flowing over the connections as this is called
        # inside `connecTCP` before the connection has been passed back to the
        # code that requested the TCP connection.

    def connect_any_redis_attempts(self) -> None:
        """If redis is enabled we need to deal with workers connecting to a
        redis server. We don't want to use a real Redis server so we use a
        fake one.
        """
        clients = self.reactor.tcpClients
        while clients:
            (host, port, client_factory, _timeout, _bindAddress) = clients.pop(0)
            self.assertEqual(host, "localhost")
            self.assertEqual(port, 6379)

            client_address = IPv4Address("TCP", "127.0.0.1", 6379)
            client_protocol = client_factory.buildProtocol(client_address)

            server_address = IPv4Address("TCP", host, port)
            server_protocol = self._redis_server.buildProtocol(server_address)

            client_to_server_transport = FakeTransport(
                server_protocol, self.reactor, client_protocol
            )
            client_protocol.makeConnection(client_to_server_transport)

            server_to_client_transport = FakeTransport(
                client_protocol, self.reactor, server_protocol
            )
            server_protocol.makeConnection(server_to_client_transport)


class TestReplicationDataHandler(ReplicationDataHandler):
    """Drop-in for ReplicationDataHandler which just collects RDATA rows"""

    def __init__(self, hs: HomeServer):
        super().__init__(hs)

        # list of received (stream_name, token, row) tuples
        self.received_rdata_rows: List[Tuple[str, int, Any]] = []

    async def on_rdata(
        self, stream_name: str, instance_name: str, token: int, rows: list
    ) -> None:
        await super().on_rdata(stream_name, instance_name, token, rows)
        for r in rows:
            self.received_rdata_rows.append((stream_name, token, r))


class FakeRedisPubSubServer:
    """A fake Redis server for pub/sub."""

    def __init__(self) -> None:
        self._subscribers_by_channel: Dict[
            bytes, Set["FakeRedisPubSubProtocol"]
        ] = defaultdict(set)

    def add_subscriber(self, conn: "FakeRedisPubSubProtocol", channel: bytes) -> None:
        """A connection has called SUBSCRIBE"""
        self._subscribers_by_channel[channel].add(conn)

    def remove_subscriber(self, conn: "FakeRedisPubSubProtocol") -> None:
        """A connection has lost connection"""
        for subscribers in self._subscribers_by_channel.values():
            subscribers.discard(conn)

    def publish(
        self, conn: "FakeRedisPubSubProtocol", channel: bytes, msg: object
    ) -> int:
        """A connection want to publish a message to subscribers."""
        for sub in self._subscribers_by_channel[channel]:
            sub.send(["message", channel, msg])

        return len(self._subscribers_by_channel)

    def buildProtocol(self, addr: IPv4Address) -> "FakeRedisPubSubProtocol":
        return FakeRedisPubSubProtocol(self)


class FakeRedisPubSubProtocol(Protocol):
    """A connection from a client talking to the fake Redis server."""

    transport: Optional[FakeTransport] = None

    def __init__(self, server: FakeRedisPubSubServer):
        self._server = server
        self._reader = hiredis.Reader()

    def dataReceived(self, data: bytes) -> None:
        self._reader.feed(data)

        # We might get multiple messages in one packet.
        while True:
            msg = self._reader.gets()

            if msg is False:
                # No more messages.
                return

            if not isinstance(msg, list):
                # Inbound commands should always be a list
                raise Exception("Expected redis list")

            self.handle_command(msg[0], *msg[1:])

    def handle_command(self, command: bytes, *args: bytes) -> None:
        """Received a Redis command from the client."""

        # We currently only support pub/sub.
        if command == b"PUBLISH":
            channel, message = args
            num_subscribers = self._server.publish(self, channel, message)
            self.send(num_subscribers)
        elif command == b"SUBSCRIBE":
            for idx, channel in enumerate(args):
                num_channels = idx + 1
                self._server.add_subscriber(self, channel)
                self.send(["subscribe", channel, num_channels])

        # Since we use SET/GET to cache things we can safely no-op them.
        elif command == b"SET":
            self.send("OK")
        elif command == b"GET":
            self.send(None)

        # Connection keep-alives.
        elif command == b"PING":
            self.send("PONG")

        else:
            raise Exception(f"Unknown command: {command!r}")

    def send(self, msg: object) -> None:
        """Send a message back to the client."""
        assert self.transport is not None

        raw = self.encode(msg).encode("utf-8")

        self.transport.write(raw)
        self.transport.flush()

    def encode(self, obj: object) -> str:
        """Encode an object to its Redis format.

        Supports: strings/bytes, integers and list/tuples.
        """

        if isinstance(obj, bytes):
            # We assume bytes are just unicode strings.
            obj = obj.decode("utf-8")

        if obj is None:
            return "$-1\r\n"
        if isinstance(obj, str):
            return f"${len(obj)}\r\n{obj}\r\n"
        if isinstance(obj, int):
            return f":{obj}\r\n"
        if isinstance(obj, (list, tuple)):
            items = "".join(self.encode(a) for a in obj)
            return f"*{len(obj)}\r\n{items}"

        raise Exception("Unrecognized type for encoding redis: %r: %r", type(obj), obj)

    def connectionLost(self, reason: Failure = connectionDone) -> None:
        self._server.remove_subscriber(self)