/*_############################################################################
  _## 
  _##  SNMP4J-UNIX-TRANSPORT - UnixSocketStreamTransportMappingTest.java  
  _## 
  _##  Copyright (C) 2021-2022  Frank Fock (SNMP4J.org)
  _##  
  _##  Licensed under the Apache License, Version 2.0 (the "License");
  _##  you may not use this file except in compliance with the License.
  _##  You may obtain a copy of the License at
  _##  
  _##      http://www.apache.org/licenses/LICENSE-2.0
  _##  
  _##  Unless required by applicable law or agreed to in writing, software
  _##  distributed under the License is distributed on an "AS IS" BASIS,
  _##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  _##  See the License for the specific language governing permissions and
  _##  limitations under the License.
  _##  
  _##########################################################################*/
package org.snmp4j.transport.unix;

import org.junit.Test;
import org.snmp4j.PDU;
import org.snmp4j.TransportMapping;
import org.snmp4j.TransportStateReference;
import org.snmp4j.asn1.BEROutputStream;
import org.snmp4j.log.ConsoleLogAdapter;
import org.snmp4j.log.ConsoleLogFactory;
import org.snmp4j.log.LogAdapter;
import org.snmp4j.log.LogFactory;
import org.snmp4j.mp.SnmpConstants;
import org.snmp4j.smi.*;
import org.snmp4j.transport.TransportListener;
import org.snmp4j.util.ThreadPool;
import org.snmp4j.util.WorkerTask;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import static org.junit.Assert.*;

public class UnixSocketStreamTransportMappingTest {

    public static final int MAX_MESSAGE_FOR_LOAD_TEST = 100;
    public static final String UNIXDOMAIN_TEST_SERVER_PREFIX = "snmp4j-unixdomain-test-server";

    static {
        LogFactory.setLogFactory(new ConsoleLogFactory());
        ConsoleLogAdapter.setWarnEnabled(true);
    }

    private static final LogAdapter LOGGER = LogFactory.getLogger(UnixSocketStreamTransportMappingTest.class);

    @Test
    public void sendMessage() throws IOException, InterruptedException {

        UnixSocketStreamTransportMapping serverTransportMapping = new UnixSocketStreamTransportMapping(
                new UnixDomainAddress(getSocketFile(UNIXDOMAIN_TEST_SERVER_PREFIX)));
        final List<OctetString> bytesReceivedList = new ArrayList<OctetString>();
        final Object sync = new Object();
        serverTransportMapping.addTransportListener(new TransportListener() {
            public <A extends Address> void processMessage(TransportMapping<? super A> sourceTransport,
                                                           A incomingAddress, ByteBuffer wholeMessage,
                                                           TransportStateReference tmStateReference) {
                OctetString bytesReceived = new OctetString(wholeMessage.array());
                LOGGER.info("Received from "+incomingAddress+": "+bytesReceived.toHexString());
                bytesReceivedList.add(bytesReceived);
                synchronized (sync) {
                    sync.notifyAll();
                }
            }
        });
        serverTransportMapping.setServerEnabled(true);
        serverTransportMapping.listen();
        UnixDomainAddress serverAddress = serverTransportMapping.getListenAddress();
        UnixSocketStreamTransportMapping clientTransportMapping = new UnixSocketStreamTransportMapping(null);
        //clientTransportMapping.listen();
        TransportStateReference transportStateReference =
                new TransportStateReference(clientTransportMapping, null, null,
                        null, null, false, new Object());
        PDU v2cPDU = new PDU();
        v2cPDU.add(new VariableBinding(new OID(SnmpConstants.sysDescr), new OctetString("hello World")));
        BEROutputStream berOutputStream = new BEROutputStream(ByteBuffer.allocate(v2cPDU.getBERLength()));
        v2cPDU.encodeBER(berOutputStream);
        byte[] bytes2Send = berOutputStream.getBuffer().array();
        synchronized (sync) {
            clientTransportMapping.sendMessage(serverAddress, bytes2Send, transportStateReference, 1000, 0);
            sync.wait(2000);
        }
        assertEquals(1, bytesReceivedList.size());
        assertEquals(new OctetString(bytes2Send).toHexString(), bytesReceivedList.get(0).toHexString());
        clientTransportMapping.close();
        serverTransportMapping.close();
    }

    private Path getSocketFile(String prefix) throws IOException {
        Path path = Files.createTempFile(prefix, "");
        // make sure path does not exist when bind is called:
        Files.deleteIfExists(path);
        return path;
    }

    @Test
    public void testSendMessageAfterReconnect() throws Exception {
        UnixSocketStreamTransportMapping serverTransportMapping = new UnixSocketStreamTransportMapping(
                new UnixDomainAddress(getSocketFile(UNIXDOMAIN_TEST_SERVER_PREFIX)));
        final List<OctetString> bytesReceivedList = new ArrayList<>();
        final Object sync = new Object();
        final Object sync2 = new Object();
        serverTransportMapping.addTransportListener(new TransportListener() {
            boolean second = false;
            public <A extends Address> void processMessage(TransportMapping<? super A> sourceTransport,
                                                           A incomingAddress, ByteBuffer wholeMessage,
                                                           TransportStateReference tmStateReference) {
                OctetString bytesReceived = new OctetString(wholeMessage.array());
                if (LOGGER.isDebugEnabled()) {
                    LOGGER.debug("Received from " + incomingAddress + ": " + bytesReceived.toHexString() +
                            " with session = " + tmStateReference.getSessionID());
                }
                bytesReceivedList.add(bytesReceived);
                synchronized (sync2) {
                    if (second) {
                        sync2.notifyAll();
                    }
                    else {
                        second = true;
                    }
                }
                synchronized (sync) {
                    sync.notifyAll();
                }
            }
        });
        serverTransportMapping.setServerEnabled(true);
        serverTransportMapping.listen();

        UnixDomainAddress serverAddress = serverTransportMapping.getListenAddress();
        UnixSocketStreamTransportMapping clientTransportMapping = new UnixSocketStreamTransportMapping();
        clientTransportMapping.listen();
        TransportStateReference transportStateReference =
                new TransportStateReference(clientTransportMapping, null, null,
                        null, null, false, sync);
        PDU v2cPDU = new PDU();
        v2cPDU.add(new VariableBinding(new OID(SnmpConstants.sysDescr), new OctetString("hello World")));
        BEROutputStream berOutputStream = new BEROutputStream(ByteBuffer.allocate(v2cPDU.getBERLength()));
        v2cPDU.encodeBER(berOutputStream);
        byte[] bytes2Send = berOutputStream.getBuffer().array();
        synchronized (sync) {
            clientTransportMapping.sendMessage(serverAddress, bytes2Send, transportStateReference, 1000, 0);
            sync.wait(5000);
        }
        assertEquals(1, bytesReceivedList.size());
        assertEquals(new OctetString(bytes2Send).toHexString(), bytesReceivedList.get(0).toHexString());
        serverTransportMapping.close();

        // Second message
        bytesReceivedList.clear();
        serverTransportMapping.listen();
        v2cPDU.add(new VariableBinding(SnmpConstants.sysUpTime, new Integer32(1234)));
        berOutputStream = new BEROutputStream(ByteBuffer.allocate(v2cPDU.getBERLength()));
        v2cPDU.encodeBER(berOutputStream);
        bytes2Send = berOutputStream.getBuffer().array();
        transportStateReference =
                new TransportStateReference(clientTransportMapping, null, null,
                        null, null, false, sync2);
        synchronized (sync2) {
            clientTransportMapping.sendMessage(serverAddress, bytes2Send, transportStateReference, 1000, 0);
            sync2.wait(2000);
        }
        assertEquals(1, bytesReceivedList.size());
        assertEquals(new OctetString(bytes2Send).toHexString(), bytesReceivedList.get(0).toHexString());
        clientTransportMapping.close();
        serverTransportMapping.close();
    }

    @Test
    public void testSendMessageAfterDisconnect() throws Exception {
        UnixSocketStreamTransportMapping serverTransportMapping = new UnixSocketStreamTransportMapping(
                new UnixDomainAddress(getSocketFile(UNIXDOMAIN_TEST_SERVER_PREFIX)));
        final List<OctetString> bytesReceivedList = new ArrayList<>();
        final Object sync = new Object();
        final Object sync2 = new Object();
        final Object sync3 = new Object();
        serverTransportMapping.addTransportListener(new TransportListener() {
            boolean second = false;
            public <A extends Address> void processMessage(TransportMapping<? super A> sourceTransport,
                                                           A incomingAddress, ByteBuffer wholeMessage,
                                                           TransportStateReference tmStateReference) {
                OctetString bytesReceived = new OctetString(wholeMessage.array());
                if (LOGGER.isDebugEnabled()) {
                    LOGGER.debug("Received from " + incomingAddress + ": " + bytesReceived.toHexString() +
                            " with session = " + tmStateReference.getSessionID());
                }
                bytesReceivedList.add(bytesReceived);
                synchronized (sync2) {
                    if (second) {
                        sync2.notifyAll();
                    }
                    else {
                        second = true;
                    }
                }
                synchronized (sync) {
                    sync.notifyAll();
                }
            }
        });
        serverTransportMapping.setServerEnabled(true);
        serverTransportMapping.listen();

        final ArrayList<Object> droppedMessage = new ArrayList<>();
        UnixDomainAddress serverAddress = serverTransportMapping.getListenAddress();
        UnixSocketStreamTransportMapping clientTransportMapping = new UnixSocketStreamTransportMapping() {
            @Override
            protected void handleDroppedMessageToSend(UnixDomainAddress address, byte[] message,
                                                      TransportStateReference transportStateReference,
                                                      long timeoutMillis, int maxRetries) {
                droppedMessage.add(message);
                synchronized (sync3) {
                    sync3.notifyAll();
                }
            }
        };
        clientTransportMapping.listen();
        TransportStateReference transportStateReference =
                new TransportStateReference(clientTransportMapping, null, null,
                        null, null, false, sync);
        PDU v2cPDU = new PDU();
        v2cPDU.add(new VariableBinding(new OID(SnmpConstants.sysDescr), new OctetString("hello World")));
        BEROutputStream berOutputStream = new BEROutputStream(ByteBuffer.allocate(v2cPDU.getBERLength()));
        v2cPDU.encodeBER(berOutputStream);
        byte[] bytes2Send = berOutputStream.getBuffer().array();
        synchronized (sync) {
            clientTransportMapping.sendMessage(serverAddress, bytes2Send, transportStateReference, 1000, 0);
            sync.wait(2000);
        }
        assertEquals(1, bytesReceivedList.size());
        assertEquals(new OctetString(bytes2Send).toHexString(), bytesReceivedList.get(0).toHexString());
        clientTransportMapping.close(serverAddress);
        assertEquals(0, clientTransportMapping.getSockets().size());

        // Second message
        bytesReceivedList.clear();
        v2cPDU.add(new VariableBinding(SnmpConstants.sysUpTime, new Integer32(1234)));
        berOutputStream = new BEROutputStream(ByteBuffer.allocate(v2cPDU.getBERLength()));
        v2cPDU.encodeBER(berOutputStream);
        bytes2Send = berOutputStream.getBuffer().array();
        transportStateReference =
                new TransportStateReference(clientTransportMapping, null, null,
                        null, null, false, sync2);
        synchronized (sync2) {
            clientTransportMapping.sendMessage(serverAddress, bytes2Send, transportStateReference, 1000, 0);
            sync2.wait(2000);
        }
        assertEquals(1, bytesReceivedList.size());
        assertEquals(new OctetString(bytes2Send).toHexString(), bytesReceivedList.get(0).toHexString());
        //clientTransportMapping.setOpenSocketOnSending(false);
        clientTransportMapping.suspendAddress(serverAddress);
        clientTransportMapping.close(serverAddress);
        assertEquals(0, clientTransportMapping.getSockets().size());

        // Third message
        bytesReceivedList.clear();
        v2cPDU.add(new VariableBinding(SnmpConstants.sysUpTime, new Integer32(1234)));
        berOutputStream = new BEROutputStream(ByteBuffer.allocate(v2cPDU.getBERLength()));
        v2cPDU.encodeBER(berOutputStream);
        bytes2Send = berOutputStream.getBuffer().array();
        transportStateReference =
                new TransportStateReference(clientTransportMapping, null, null,
                        null, null, false, sync3);
        synchronized (sync3) {
            clientTransportMapping.sendMessage(serverAddress, bytes2Send, transportStateReference, 1000, 0);
            sync3.wait(2000);
        }
        // no message sent:
        assertEquals(0, bytesReceivedList.size());
        // message dropped:
        assertEquals(1, droppedMessage.size());
        clientTransportMapping.resumeAddress(serverAddress);
        assertEquals(0, clientTransportMapping.getSuspendedAddresses().size());

        // Fourth message
        v2cPDU.add(new VariableBinding(SnmpConstants.sysUpTime, new Integer32(1234)));
        berOutputStream = new BEROutputStream(ByteBuffer.allocate(v2cPDU.getBERLength()));
        v2cPDU.encodeBER(berOutputStream);
        bytes2Send = berOutputStream.getBuffer().array();
        transportStateReference =
                new TransportStateReference(clientTransportMapping, null, null,
                        null, null, false, sync3);
        synchronized (sync2) {
            clientTransportMapping.sendMessage(serverAddress, bytes2Send, transportStateReference, 1000, 0);
            sync2.wait(2000);
        }
        assertEquals(1, bytesReceivedList.size());
        assertEquals(new OctetString(bytes2Send).toHexString(), bytesReceivedList.get(0).toHexString());

        serverTransportMapping.close();
        assertNull(serverTransportMapping.getListenWorkerTask());
        assertEquals(0, serverTransportMapping.getSockets().size());
    }

    @Test
    public void testReconnectUnderLoad() throws Exception {

        UnixSocketStreamTransportMapping serverTransportMapping = new UnixSocketStreamTransportMapping(
                new UnixDomainAddress(getSocketFile(UNIXDOMAIN_TEST_SERVER_PREFIX)));
        final List<OctetString> bytesReceivedList = Collections.synchronizedList(new ArrayList<>());
        final List<OctetString> bytesSentList = Collections.synchronizedList(new ArrayList<>());
        serverTransportMapping.addTransportListener(new TransportListener() {
            public <A extends Address> void processMessage(TransportMapping<? super A> sourceTransport,
                                                           A incomingAddress, ByteBuffer wholeMessage,
                                                           TransportStateReference tmStateReference) {
                OctetString bytesReceived = new OctetString(wholeMessage.array());
                if (LOGGER.isDebugEnabled()) {
                    LOGGER.debug("Received from " + incomingAddress + ": " + bytesReceived.toHexString() +
                            " with session = " + tmStateReference.getSessionID());
                }
                bytesReceivedList.add(bytesReceived);
            }
        });
        serverTransportMapping.setServerEnabled(true);
        serverTransportMapping.listen();

        UnixDomainAddress serverAddress = serverTransportMapping.getListenAddress();

        int numTransportMappings = 1;
        final Object[] sessionIDs = new Object[numTransportMappings];
        final UnixSocketStreamTransportMapping[] clientTransportMappings =
                new UnixSocketStreamTransportMapping[numTransportMappings];
        final TransportStateReference[] transportStateReferences = new TransportStateReference[numTransportMappings];
        for (int i=0; i<clientTransportMappings.length; i++) {
            clientTransportMappings[i] = new UnixSocketStreamTransportMapping();
            clientTransportMappings[i].listen();
            transportStateReferences[i] =
                    new TransportStateReference(clientTransportMappings[i], null, null,
                            null, null, false, sessionIDs[i]);
        }
        ThreadPool clientPool = ThreadPool.create("UnixDomainTM", numTransportMappings);
        for (TransportStateReference tmRef : transportStateReferences) {
            @SuppressWarnings("unchecked")
            TransportMapping<UnixDomainAddress> clientTransport = (TransportMapping<UnixDomainAddress>) tmRef.getTransport();
            clientPool.execute(new WorkerTask() {
                boolean stop = false;

                @Override
                public void shutdown() {

                }

                @Override
                public void terminate() {
                    stop = true;
                }

                @Override
                public void join() throws InterruptedException {

                }

                @Override
                public boolean join(long timeout) throws InterruptedException {
                    return false;
                }

                @Override
                public void interrupt() {

                }

                @Override
                public void run() {
                    for (int i = 0; (!stop) && i< MAX_MESSAGE_FOR_LOAD_TEST; i++) {
                        try {
                            PDU v2cPDU = new PDU();
                            v2cPDU.add(new VariableBinding(new OID(SnmpConstants.sysDescr), new OctetString("hello World"+i)));
                            BEROutputStream berOutputStream = new BEROutputStream(ByteBuffer.allocate(v2cPDU.getBERLength()));
                            v2cPDU.encodeBER(berOutputStream);
                            byte[] bytes2Send = berOutputStream.getBuffer().array();
                            bytesSentList.add(OctetString.fromByteArray(bytes2Send));
                            clientTransport.sendMessage(serverAddress, bytes2Send, tmRef, 100, 1);
                            Thread.sleep(10);
                        } catch (IOException | InterruptedException iox) {
                            iox.printStackTrace();
                        }
                    }
                }
            });
        }
        Thread.sleep(500);
        serverTransportMapping.close();
        assertEquals(0, serverTransportMapping.getSockets().size());
        serverTransportMapping.listen();
        assertNotNull(serverTransportMapping.serverThread);
        Thread.sleep(500);
        clientPool.stop();
        for (UnixSocketStreamTransportMapping ransportMapping : clientTransportMappings) {
            assertEquals(1, ransportMapping.getSockets().size());
        }
        assertTrue(bytesSentList.size() - bytesReceivedList.size() <= 1);
    }

}