import os

import mpitestutil as testutil
import mpiunittest as unittest

from mpi4py import MPI

try:
    import socket
except ImportError:
    socket = None


def ch4_ucx():
    return "ch4:ucx" in MPI.Get_library_version()


def ch4_ofi():
    return "ch4:ofi" in MPI.Get_library_version()


@unittest.skipMPI("mpich(<4.3.0)", not testutil.has_mpi_port())
@unittest.skipMPI("openmpi(<2.0.0)")
@unittest.skipMPI("openmpi(>=5.0.0,<5.0.4)")
@unittest.skipMPI("msmpi(<8.1.0)")
@unittest.skipMPI("mvapich(<3.0.0)")
@unittest.skipIf(MPI.COMM_WORLD.Get_size() < 2, "mpi-world-size<2")
class TestDPM(unittest.TestCase):
    #
    message = [
        None,
        True,
        False,
        -7,
        0,
        7,
        -(2**63) + 1,
        2**63 - 1,
        -2.17,
        0.0,
        +1.62,
        1 + 2j,
        2 - 3j,
        "mpi4py",
        (1, 2, 3),
        [1, 2, 3],
        {1: 2},
    ]

    def testNamePublishing(self):
        rank = MPI.COMM_WORLD.Get_rank()
        service = f"mpi4py-{rank}"
        port = MPI.Open_port()
        MPI.Publish_name(service, port)
        found = MPI.Lookup_name(service)
        self.assertEqual(port, found)
        MPI.Unpublish_name(service, port)
        MPI.Close_port(port)

    @unittest.skipMPI("mpich(==3.4.1)", ch4_ofi())
    @unittest.skipMPI("mvapich", ch4_ofi())
    @unittest.skipMPI("impi", MPI.COMM_WORLD.Get_size() > 2)
    def testAcceptConnect(self):
        comm_self = MPI.COMM_SELF
        comm_world = MPI.COMM_WORLD
        wrank = comm_world.Get_rank()
        group_world = comm_world.Get_group()
        group = group_world.Excl([0])
        group_world.Free()
        comm = comm_world.Create(group)
        group.Free()
        if wrank == 0:
            self.assertEqual(comm, MPI.COMM_NULL)
        else:
            self.assertNotEqual(comm, MPI.COMM_NULL)
            self.assertEqual(comm.size, comm_world.size - 1)
            self.assertEqual(comm.rank, comm_world.rank - 1)
        if wrank == 0:
            port = MPI.Open_port()
            comm_world.send(port, dest=1)
            intercomm = comm_self.Accept(port)
            self.assertEqual(intercomm.remote_size, comm_world.size - 1)
            self.assertEqual(intercomm.size, 1)
            self.assertEqual(intercomm.rank, 0)
            MPI.Close_port(port)
        else:
            if wrank == 1:
                port = comm_world.recv(source=0)
            else:
                port = None
            intercomm = comm.Connect(port, root=0)
            self.assertEqual(intercomm.remote_size, 1)
            self.assertEqual(intercomm.size, comm_world.size - 1)
            self.assertEqual(intercomm.rank, comm.rank)
            comm.Free()
        if wrank == 0:
            message = TestDPM.message
            root = MPI.ROOT
        else:
            message = None
            root = 0
        message = intercomm.bcast(message, root)
        if wrank == 0:
            self.assertIsNone(message)
        else:
            self.assertEqual(message, TestDPM.message)
        intercomm.Free()

    def testConnectAccept(self):
        comm_self = MPI.COMM_SELF
        comm_world = MPI.COMM_WORLD
        wrank = comm_world.Get_rank()
        group_world = comm_world.Get_group()
        group = group_world.Excl([0])
        group_world.Free()
        comm = comm_world.Create(group)
        group.Free()
        if wrank == 0:
            self.assertEqual(comm, MPI.COMM_NULL)
        else:
            self.assertNotEqual(comm, MPI.COMM_NULL)
            self.assertEqual(comm.size, comm_world.size - 1)
            self.assertEqual(comm.rank, comm_world.rank - 1)
        if wrank == 0:
            port = comm_world.recv(source=1)
            intercomm = comm_self.Connect(port)
            self.assertEqual(intercomm.remote_size, comm_world.size - 1)
            self.assertEqual(intercomm.size, 1)
            self.assertEqual(intercomm.rank, 0)
        else:
            if wrank == 1:
                port = MPI.Open_port()
                comm_world.send(port, dest=0)
            else:
                port = None
            intercomm = comm.Accept(port, root=0)
            if wrank == 1:
                MPI.Close_port(port)
            self.assertEqual(intercomm.remote_size, 1)
            self.assertEqual(intercomm.size, comm_world.size - 1)
            self.assertEqual(intercomm.rank, comm.rank)
            comm.Free()
        if wrank == 0:
            message = TestDPM.message
            root = MPI.ROOT
        else:
            message = None
            root = 0
        message = intercomm.bcast(message, root)
        if wrank == 0:
            self.assertIsNone(message)
        else:
            self.assertEqual(message, TestDPM.message)
        intercomm.Free()

    @unittest.skipIf(socket is None, "socket")
    @unittest.skipMPI("impi", testutil.github() and os.name == "nt")
    def testJoin(self):
        rank = MPI.COMM_WORLD.Get_rank()
        server = client = address = None
        host = socket.gethostname()
        addrinfo = socket.getaddrinfo(host, None, type=socket.SOCK_STREAM)
        addr_families = [info[0] for info in addrinfo]
        # if both INET and INET6 are available, don't assume the order
        # is the same on both server and client. Prefer INET if available.
        addr_family = None
        if socket.AF_INET in addr_families:
            addr_family = socket.AF_INET
        elif socket.AF_INET6 in addr_families:
            addr_family = socket.AF_INET6
        addr_family = MPI.COMM_WORLD.bcast(addr_family, root=0)
        supported = addr_family in addr_families
        supported = MPI.COMM_WORLD.allreduce(supported, op=MPI.LAND)
        if not supported:
            self.skipTest("socket-inet")
        # create server/client sockets
        if rank == 0:  # server
            server = socket.socket(addr_family, socket.SOCK_STREAM)
            server.bind((host, 0))
            server.listen(0)
        if rank == 1:  # client
            client = socket.socket(addr_family, socket.SOCK_STREAM)
        # communicate address
        if rank == 0:
            address = server.getsockname()
            MPI.COMM_WORLD.ssend(address, 1)
        if rank == 1:
            address = MPI.COMM_WORLD.recv(None, 0)
        MPI.COMM_WORLD.Barrier()
        # stablish client/server connection
        connected = False
        if rank == 0:  # server
            client = server.accept()[0]
            server.close()
        if rank == 1:  # client
            client.connect(address)
            connected = True
        connected = MPI.COMM_WORLD.bcast(connected, root=1)
        # test Comm.Join()
        MPI.COMM_WORLD.Barrier()
        if client:
            fd = client.fileno()
            intercomm = MPI.Comm.Join(fd)
            client.close()
            if intercomm != MPI.COMM_NULL:
                self.assertEqual(intercomm.remote_size, 1)
                self.assertEqual(intercomm.size, 1)
                self.assertEqual(intercomm.rank, 0)
                if rank == 0:
                    message = TestDPM.message
                    root = MPI.ROOT
                else:
                    message = None
                    root = 0
                message = intercomm.bcast(message, root)
                if rank == 0:
                    self.assertIsNone(message)
                else:
                    self.assertEqual(message, TestDPM.message)
                intercomm.Free()
        MPI.COMM_WORLD.Barrier()


if __name__ == "__main__":
    unittest.main()
