/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.apache.activemq.artemis.cli.commands.tools.xml;

import javax.xml.stream.XMLStreamConstants;
import javax.xml.stream.XMLStreamReader;
import javax.xml.transform.stax.StAXSource;
import javax.xml.validation.Validator;
import java.io.FileInputStream;
import java.io.InputStream;
import java.lang.invoke.MethodHandles;
import java.net.URL;
import java.nio.ByteBuffer;
import java.security.AccessController;
import java.security.PrivilegedAction;
import java.util.Comparator;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.TreeSet;

import org.apache.activemq.artemis.api.core.Message;
import org.apache.activemq.artemis.api.core.QueueConfiguration;
import org.apache.activemq.artemis.api.core.RoutingType;
import org.apache.activemq.artemis.api.core.SimpleString;
import org.apache.activemq.artemis.api.core.TransportConfiguration;
import org.apache.activemq.artemis.api.core.client.ActiveMQClient;
import org.apache.activemq.artemis.api.core.client.ClientMessage;
import org.apache.activemq.artemis.api.core.client.ClientProducer;
import org.apache.activemq.artemis.api.core.client.ClientRequestor;
import org.apache.activemq.artemis.api.core.client.ClientSession;
import org.apache.activemq.artemis.api.core.client.ClientSessionFactory;
import org.apache.activemq.artemis.api.core.client.ServerLocator;
import org.apache.activemq.artemis.api.core.management.ManagementHelper;
import org.apache.activemq.artemis.api.core.management.ResourceNames;
import org.apache.activemq.artemis.cli.commands.ActionContext;
import org.apache.activemq.artemis.cli.commands.messages.BasicConnectionAbstract;
import org.apache.activemq.artemis.core.filter.impl.FilterImpl;
import org.apache.activemq.artemis.core.protocol.core.impl.PacketImpl;
import org.apache.activemq.artemis.core.remoting.impl.netty.NettyConnectorFactory;
import org.apache.activemq.artemis.core.remoting.impl.netty.TransportConstants;
import org.apache.activemq.artemis.utils.ClassloadingUtil;
import org.apache.activemq.artemis.utils.ListUtil;
import org.apache.activemq.artemis.utils.XmlProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import picocli.CommandLine.Command;
import picocli.CommandLine.Option;

/**
 * Read XML output from <code>org.apache.activemq.artemis.core.persistence.impl.journal.XmlDataExporter</code>, create a core session, and
 * send the messages to a running instance of ActiveMQ Artemis.  It uses the StAX <code>javax.xml.stream.XMLStreamReader</code>
 * for speed and simplicity.
 */
@Command(name = "imp", description = "Import all message-data using an XML that could be interpreted by any system.")
public final class XmlDataImporter extends BasicConnectionAbstract {

   private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

   private XMLStreamReader reader;

   private XMLMessageImporter messageReader;

   // this session is really only needed if the "session" variable does not auto-commit sends
   ClientSession managementSession;

   boolean localSession = false;

   final Map<String, String> addressMap = new HashMap<>();

   final Map<String, Long> queueIDs = new HashMap<>();

   HashMap<String, String> oldPrefixTranslation = new HashMap<>();

   private ClientSession session;
   private ClientProducer producer;

   @Deprecated(forRemoval = true)
   @Option(names = "--host", description = "The host used to import the data. Default: null.", hidden = true)
   public String host = null;

   @Deprecated(forRemoval = true)
   @Option(names = "--port", description = "The port used to import the data. Default: 61616.", hidden = true)
   public int port = 61616;

   @Option(names = "--transaction", description = "Import every message using a single transction. If anything goes wrong during the process the entire import will be aborted. Default: false.", hidden = true)
   public boolean transactional;

   @Option(names = "--commit-interval", description = "How often to commit.", hidden = true)
   public int commitInterval = 1000;

   @Option(names = "--input", description = "The input file name. Default: exp.dmp.", required = true)
   public String input = "exp.dmp";

   @Option(names = "--sort", description = "Sort the messages from the input (used for older versions that won't sort messages).")
   public boolean sort = false;

   @Option(names = "--legacy-prefixes", description = "Do not remove prefixes from legacy imports.")
   public boolean legacyPrefixes = false;

   TreeSet<XMLMessageImporter.MessageInfo> messages;

   @Override
   public Object execute(ActionContext context) throws Exception {
      process(input, host, port);
      return null;
   }

   public void process(String inputFileName, String host, int port) throws Exception {
      try (FileInputStream inputFile = new FileInputStream(inputFileName)) {
         this.process(inputFile, host, port);
      }
   }

   /**
    * This is the normal constructor for programmatic access to the
    * <code>org.apache.activemq.artemis.core.persistence.impl.journal.XmlDataImporter</code> if the session passed
    * in uses auto-commit for sends.
    * <br>
    * If the session needs to be transactional then use the constructor which takes 2 sessions.
    *
    * @param inputStream the stream from which to read the XML for import
    * @param session     used for sending messages, must use auto-commit for sends
    */
   public void process(InputStream inputStream, ClientSession session) throws Exception {
      this.process(inputStream, session, null);
   }

   /**
    * This is the constructor to use if you wish to import all messages transactionally.
    * <br>
    * Pass in a session which doesn't use auto-commit for sends, and one that does (for management
    * operations necessary during import).
    *
    * @param inputStream       the stream from which to read the XML for import
    * @param session           used for sending messages, doesn't need to auto-commit sends
    * @param managementSession used for management queries, must use auto-commit for sends
    */
   public void process(InputStream inputStream,
                       ClientSession session,
                       ClientSession managementSession) throws Exception {
      Objects.requireNonNull(inputStream);
      reader = XmlProvider.createXMLStreamReader(inputStream);
      messageReader = new XMLMessageImporter(reader, session);
      messageReader.setOldPrefixTranslation(oldPrefixTranslation);

      this.session = session;
      this.producer = session.createProducer();
      if (managementSession != null) {
         this.managementSession = managementSession;
      } else {
         this.managementSession = session;
      }

      processXml();
   }

   public void process(InputStream inputStream, String host, int port) throws Exception {
      ServerLocator serverLocator;
      if (host != null) {
         Map<String, Object> connectionParams = Map.of(
            TransportConstants.HOST_PROP_NAME, host,
            TransportConstants.PORT_PROP_NAME, Integer.toString(port)
         );
         serverLocator = ActiveMQClient.createServerLocatorWithoutHA(new TransportConfiguration(NettyConnectorFactory.class.getName(), connectionParams));
      } else {
         serverLocator = ActiveMQClient.createServerLocator(brokerURL);
      }
      ClientSessionFactory sf = serverLocator.createSessionFactory();

      ClientSession session = null;
      ClientSession managementSession = null;

      try {

         if (user != null || password != null) {
            session = sf.createSession(user, password, false, false, true, false, 0);
            managementSession = sf.createSession(user, password, false, true, true, false, 0);
         } else {
            session = sf.createSession(false, false, true);
            managementSession = sf.createSession(false, true, true);
         }
         localSession = true;

         producer = session.createProducer();

         process(inputStream, session, managementSession);
      } finally {
         try {
            session.commit();
            session.close();
         } catch (Throwable ignored) {
         }
         try {
            managementSession.commit();
            managementSession.close();
         } catch (Throwable ignored) {
         }


      }
   }

   public void validate(String fileName) throws Exception {
      try (FileInputStream file = new FileInputStream(fileName)) {
         validate(file);
      }
   }

   public void validate(InputStream inputStream) throws Exception {
      XMLStreamReader reader = XmlProvider.createXMLStreamReader(inputStream);
      Validator validator = XmlProvider.newValidator(XmlDataImporter.findResource("schema/artemis-import-export.xsd"));
      validator.validate(new StAXSource(reader));
      reader.close();
   }

   private static URL findResource(final String resourceName) {
      return AccessController.doPrivileged(new PrivilegedAction<URL>() {
         @Override
         public URL run() {
            return ClassloadingUtil.findResource(resourceName);
         }
      });
   }

   private void processXml() throws Exception {
      if (sort) {
         messages = new TreeSet<XMLMessageImporter.MessageInfo>(new Comparator<XMLMessageImporter.MessageInfo>() {
            @Override
            public int compare(XMLMessageImporter.MessageInfo o1, XMLMessageImporter.MessageInfo o2) {
               return Long.compare(o1.id, o2.id);
            }
         });
      }
      while (reader.hasNext()) {
         if (logger.isDebugEnabled()) {
            logger.debug("EVENT:[{}][{}] ", reader.getLocation().getLineNumber(), reader.getLocation().getColumnNumber());
         }

         if (reader.getEventType() == XMLStreamConstants.START_ELEMENT) {
            if (XmlDataConstants.OLD_BINDING.equals(reader.getLocalName())) {
               oldBinding(); // export from 1.x
            } else if (XmlDataConstants.QUEUE_BINDINGS_CHILD.equals(reader.getLocalName())) {
               bindQueue();
            } else if (XmlDataConstants.ADDRESS_BINDINGS_CHILD.equals(reader.getLocalName())) {
               bindAddress();
            } else if (XmlDataConstants.MESSAGES_CHILD.equals(reader.getLocalName())) {
               processMessage();
            }
         }
         reader.next();
      }

      if (sort) {
         long messageNr = 0;
         for (XMLMessageImporter.MessageInfo msgtmp : messages) {
            sendMessage(msgtmp.queues, msgtmp.message);
            messageNr++;
            if (messageNr % commitInterval == 0) {
               session.commit();
            }
         }
      }

      session.commit();
   }

   long messageNr = 0;

   private void processMessage() throws Exception {
      messageNr++;

      if (messageNr % commitInterval == 0) {
         getActionContext().err.println("Processed " + messageNr + " messages");
         session.commit();
      }
      XMLMessageImporter.MessageInfo info = messageReader.readMessage(false);
      if (sort) {
         messages.add(info);
      } else {
         sendMessage(info.queues, info.message);
      }
   }

   private void sendMessage(List<String> queues, Message message) throws Exception {
      final String destination = addressMap.get(queues.get(0));
      final ByteBuffer buffer = ByteBuffer.allocate(queues.size() * 8);

      final boolean debugLog = logger.isDebugEnabled();
      final StringBuilder debugLogMessage = debugLog ? new StringBuilder() : null;
      if (debugLog) {
         debugLogMessage.append("Sending ").append(message).append(" to address: ").append(destination).append("; routed to queues: ");
      }

      for (String queue : queues) {
         long queueID = -1;

         if (queueIDs.containsKey(queue)) {
            queueID = queueIDs.get(queue);
         } else {
            // Get the ID of the queues involved so the message can be routed properly.  This is done because we cannot
            // send directly to a queue, we have to send to an address instead but not all the queues related to the
            // address may need the message
            try (ClientRequestor requestor = new ClientRequestor(managementSession, "activemq.management")) {
               ClientMessage managementMessage = managementSession.createMessage(false);
               ManagementHelper.putAttribute(managementMessage, ResourceNames.QUEUE + queue, "ID");
               managementSession.start();
               if (debugLog) {
                  logger.debug("Requesting ID for: {}", queue);
               }
               ClientMessage reply = requestor.request(managementMessage);
               if (ManagementHelper.hasOperationSucceeded(reply)) {
                  Number idObject = (Number) ManagementHelper.getResult(reply);
                  queueID = idObject.longValue();
               } else {
                  if (debugLog) {
                     logger.debug("Failed to get ID for {}, reply: {}", queue, ManagementHelper.getResult(reply, String.class));
                  }
               }
            }

            if (debugLog) {
               logger.debug("ID for {} is: {}", queue, queueID);
            }
            if (queueID != -1) {
               queueIDs.put(queue, queueID);  // store it so we don't have to look it up every time
            }
         }

         if (queueID != -1) {
            buffer.putLong(queueID);
         }
         if (debugLog) {
            debugLogMessage.append(queue).append(", ");
         }
      }

      if (debugLog) {
         debugLogMessage.delete(debugLogMessage.length() - 2, debugLogMessage.length()); // take off the trailing comma+space
         logger.debug(debugLogMessage.toString());
      }

      message.putBytesProperty(Message.HDR_ROUTE_TO_IDS, buffer.array());
      producer.send(SimpleString.toSimpleString(destination), message);
   }

   private void oldBinding() throws Exception {
      String queueName = "";
      String address = "";
      String filter = "";

      for (int i = 0; i < reader.getAttributeCount(); i++) {
         String attributeName = reader.getAttributeLocalName(i);
         switch (attributeName) {
            case XmlDataConstants.OLD_ADDRESS:
               address = reader.getAttributeValue(i);
               break;
            case XmlDataConstants.OLD_QUEUE:
               queueName = reader.getAttributeValue(i);
               break;
            case XmlDataConstants.OLD_FILTER:
               filter = reader.getAttributeValue(i);
               break;
         }
      }

      if (queueName == null || address == null || filter == null) {
         // not expected to happen unless someone manually changed the format
         throw new IllegalStateException("invalid format, missing queue, address or filter");
      }

      RoutingType routingType = RoutingType.MULTICAST;

      if (address.startsWith(PacketImpl.OLD_QUEUE_PREFIX.toString())) {
         routingType = RoutingType.ANYCAST;
         if (!legacyPrefixes) {
            String newaddress = address.substring(PacketImpl.OLD_QUEUE_PREFIX.length());
            address = newaddress;
         }
      } else if (address.startsWith(PacketImpl.OLD_TOPIC_PREFIX.toString())) {
         routingType = RoutingType.MULTICAST;
         if (!legacyPrefixes) {
            String newaddress = address.substring(PacketImpl.OLD_TOPIC_PREFIX.length());
            address = newaddress;
         }
      }

      if (queueName.startsWith(PacketImpl.OLD_QUEUE_PREFIX.toString())) {
         if (!legacyPrefixes) {
            String newQueueName = queueName.substring(PacketImpl.OLD_QUEUE_PREFIX.length());
            oldPrefixTranslation.put(queueName, newQueueName);
            queueName = newQueueName;
         }
      } else if (queueName.startsWith(PacketImpl.OLD_TOPIC_PREFIX.toString())) {
         if (!legacyPrefixes) {
            String newQueueName = queueName.substring(PacketImpl.OLD_TOPIC_PREFIX.length());
            oldPrefixTranslation.put(queueName, newQueueName);
            queueName = newQueueName;
         }
      }


      ClientSession.AddressQuery addressQuery = session.addressQuery(SimpleString.toSimpleString(address));

      if (!addressQuery.isExists()) {
         session.createAddress(SimpleString.toSimpleString(address), routingType, true);
      }

      if (!filter.equals(FilterImpl.GENERIC_IGNORED_FILTER)) {
         ClientSession.QueueQuery queueQuery = session.queueQuery(new SimpleString(queueName));

         if (!queueQuery.isExists()) {
            session.createQueue(new QueueConfiguration(queueName).setAddress(address).setRoutingType(routingType).setFilterString(filter));
            if (logger.isDebugEnabled()) {
               logger.debug("Binding queue(name={}, address={}, filter={})", queueName, address, filter);
            }
         } else {
            logger.debug("Binding {} already exists so won't re-bind.", queueName);
         }
      }

      addressMap.put(queueName, address);
   }


   private void bindQueue() throws Exception {
      String queueName = "";
      String address = "";
      String filter = "";
      String routingType = "";

      for (int i = 0; i < reader.getAttributeCount(); i++) {
         String attributeName = reader.getAttributeLocalName(i);
         switch (attributeName) {
            case XmlDataConstants.QUEUE_BINDING_ADDRESS:
               address = reader.getAttributeValue(i);
               break;
            case XmlDataConstants.QUEUE_BINDING_NAME:
               queueName = reader.getAttributeValue(i);
               break;
            case XmlDataConstants.QUEUE_BINDING_FILTER_STRING:
               filter = reader.getAttributeValue(i);
               break;
            case XmlDataConstants.QUEUE_BINDING_ROUTING_TYPE:
               routingType = reader.getAttributeValue(i);
               break;
         }
      }

      ClientSession.QueueQuery queueQuery = session.queueQuery(new SimpleString(queueName));

      if (!queueQuery.isExists()) {
         session.createQueue(new QueueConfiguration(queueName).setAddress(address).setRoutingType(RoutingType.valueOf(routingType)).setFilterString(filter));
         if (logger.isDebugEnabled()) {
            logger.debug("Binding queue(name={}, address={}, filter={})", queueName, address, filter);
         }
      } else {
         logger.debug("Binding {} already exists so won't re-bind.", queueName);
      }

      addressMap.put(queueName, address);
   }

   private void bindAddress() throws Exception {
      String addressName = "";
      String routingTypes = "";

      for (int i = 0; i < reader.getAttributeCount(); i++) {
         String attributeName = reader.getAttributeLocalName(i);
         switch (attributeName) {
            case XmlDataConstants.ADDRESS_BINDING_NAME:
               addressName = reader.getAttributeValue(i);
               break;
            case XmlDataConstants.ADDRESS_BINDING_ROUTING_TYPE:
               routingTypes = reader.getAttributeValue(i);
               break;
         }
      }

      ClientSession.AddressQuery addressQuery = session.addressQuery(new SimpleString(addressName));

      if (!addressQuery.isExists()) {
         EnumSet<RoutingType> set = EnumSet.noneOf(RoutingType.class);
         for (String routingType : ListUtil.toList(routingTypes)) {
            set.add(RoutingType.valueOf(routingType));
         }
         session.createAddress(SimpleString.toSimpleString(addressName), set, false);
         logger.debug("Binding address(name={}, routingTypes={})", addressName, routingTypes);
      } else {
         logger.debug("Binding {} already exists so won't re-bind.", addressName);
      }
   }


}
