package tc.oc.api.queue;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Type;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nullable;
import javax.inject.Inject;

import com.google.common.base.Charsets;
import com.google.common.reflect.TypeToken;
import com.google.common.util.concurrent.SettableFuture;
import com.google.gson.Gson;
import com.rabbitmq.client.AMQP;
import com.rabbitmq.client.DefaultConsumer;
import com.rabbitmq.client.Envelope;
import tc.oc.api.connectable.Connectable;
import tc.oc.api.message.Message;
import tc.oc.api.message.MessageHandler;
import tc.oc.api.message.MessageListener;
import tc.oc.api.message.MessageQueue;
import tc.oc.api.message.MessageRegistry;
import tc.oc.api.message.NoSuchMessageException;
import tc.oc.api.serialization.Pretty;
import tc.oc.commons.core.exception.ExceptionHandler;
import tc.oc.commons.core.logging.Loggers;
import tc.oc.commons.core.reflect.Methods;
import tc.oc.commons.core.reflect.Types;
import tc.oc.commons.core.util.CachingTypeMap;
import tc.oc.minecraft.suspend.Suspendable;

/**
 * An AMQP queue.
 *
 * Declaration and binding (note that Queue can be instantiated without an AMQP connection):
 *
 *     final Queue MY_QUEUE = new Queue("my_queue", true, false, false, null);
 *     MY_QUEUE.declare(Api.get().queueClient());
 *
 * Binding:
 *
 *     MY_QUEUE.bind(Exchange.DIRECT); // queue name is default routing key
 *     MY_QUEUE.bind(Exchange.FANOUT);
 *     MY_QUEUE.bind(Exchange.TOPIC, "fun_stuff");
 *
 * Subscription with {@link MessageHandler}:
 *
 *     MY_QUEUE.subscribe(ServerReconfigure.class, new MessageHandler<ServerReconfigure>() {
 *           @Override
 *           public void handleDelivery(ServerReconfigure message, Delivery delivery) {
 *               // ...
 *           }
 *     }, false, syncExecutor);
 *
 * Subscription with {@link MessageListener}:
 *
 *     class MyWorker implements MessageListener {
 *         @HandleMessage
 *         public void handleHealthReport(ServerHealthReport message, Delivery delivery) {
 *             // ...
 *         }
 *
 *         @HandleMessage
 *         public void handlePlayerReport(ServerPlayerReport message, Delivery delivery) {
 *             // ...
 *         }
 *     }
 *
 *     MY_QUEUE.subscribe(new MyWorker(), syncExecutor);
 */
public class Queue implements MessageQueue, Connectable, Suspendable {

    private static class RegisteredHandler<T extends Message> {
        final @Nullable MessageListener listener;
        final MessageHandler<T> handler;
        final @Nullable Executor executor;

        private RegisteredHandler(@Nullable MessageListener listener, MessageHandler<T> handler, @Nullable Executor executor) {
            this.listener = listener;
            this.handler = handler;
            this.executor = executor;
        }
    }

    protected Logger logger;
    @Inject protected MessageRegistry messageRegistry;
    @Inject protected Gson gson;
    @Inject @Pretty protected Gson prettyGson;
    @Inject protected QueueClient client;
    @Inject protected Exchange.Topic topic;
    @Inject protected ExceptionHandler exceptionHandler;

    protected final Consume consume;

    @Nullable String consumerTag;
    private MultiDispatcher dispatcher;

    private final Set<RegisteredHandler<?>> handlers = new HashSet<>();
    private final CachingTypeMap<Message, RegisteredHandler<?>> handlersByType = CachingTypeMap.create();
    private volatile boolean suspended;

    public Consume consume() { return consume; }
    public String name() { return consume.name(); }
    public QueueClient client() { return client; }

    public Queue(Consume consume) {
        this.consume = consume;
    }

    @Inject void init(Loggers loggers) {
        logger = loggers.get(getClass());
    }

    @Override
    public void connect() throws IOException {
        logger.fine("Declaring queue");
        client.getChannel().queueDeclare(consume.name(), consume.durable(), consume.exclusive(), consume.autoDelete(), consume.arguments());
        dispatcher = new MultiDispatcher();
        consumerTag = client.getChannel().basicConsume(consume.name(), false, "", false, true, Collections.emptyMap(), dispatcher);
    }

    @Override
    public void disconnect() throws IOException {
        if(dispatcher != null) {
            if(consumerTag != null) {
                client.getChannel().basicCancel(consumerTag);
                consumerTag = null;
            }
            try {
                dispatcher.awaitTermination(5L, TimeUnit.SECONDS);
            } catch(InterruptedException | ExecutionException | TimeoutException e) {
                throw new IOException("Failed to shutdown " + this, e);
            }
        }
    }

    @Override
    public void suspend() {
        suspended = true;
    }

    @Override
    public void resume() {
        suspended = false;
    }

    protected void bind(Exchange exchange, String routingKey) {
        try {
            logger.fine("Binding to exchange " + exchange.name() + " with routing key " + routingKey);

            client.getChannel().queueBind(name(), exchange.name(), routingKey, null);
        } catch(IOException e) {
            throw new IllegalStateException("Failed to bind to exchange " + exchange.name(), e);
        }
    }

    @Override
    public void bind(Class<? extends Message> type) {
        bind(topic, messageRegistry.typeName(type));
    }

    @Override
    public <T extends Message> void subscribe(TypeToken<T> messageType, MessageHandler<T> handler, @Nullable Executor executor) {
        subscribe(messageType, null, handler, executor);
    }

    private <T extends Message> void subscribe(TypeToken<T> messageType, @Nullable MessageListener listener, MessageHandler<T> handler, @Nullable Executor executor) {
        logger.fine("Subscribing handler " + handler);
        synchronized(handlers) {
            final RegisteredHandler<T> registered = new RegisteredHandler<>(listener, handler, executor);
            handlers.add(registered);
            handlersByType.put(messageType, registered);
            handlersByType.invalidate();
        }
    }

    private TypeToken<? extends Message> getMessageType(TypeToken decl, Method method) {
        if(method.getParameterTypes().length < 1 || method.getParameterTypes().length > 3) {
            throw new IllegalStateException("Message handler method must take 1 to 3 parameters");
        }

        final TypeToken<Message> base = new TypeToken<Message>(){};

        for(Type param : method.getGenericParameterTypes()) {
            final TypeToken paramToken = decl.resolveType(param);
            Types.assertFullySpecified(paramToken);
            if(base.isAssignableFrom(paramToken)) {
                messageRegistry.typeName(paramToken.getRawType()); // Verify message type is registered
                return paramToken;
            }
        }

        throw new IllegalStateException("Message handler has no message parameter");
    }

    @Override
    public void subscribe(final MessageListener listener, @Nullable Executor executor) {
        logger.fine("Subscribing listener " + listener);

        final TypeToken<? extends MessageListener> listenerType = TypeToken.of(listener.getClass());
        Methods.declaredMethodsInAncestors(listener.getClass()).forEach(method -> {
            final MessageListener.HandleMessage annot = method.getAnnotation(MessageListener.HandleMessage.class);
            if(annot != null) {
                method.setAccessible(true);
                final TypeToken<? extends Message> messageType = getMessageType(listenerType, method);

                logger.fine("  dispatching " + messageType.getRawType().getSimpleName() + " to method " + method.getName());

                MessageHandler handler = new MessageHandler() {
                    @Override
                    public void handleDelivery(Message message, TypeToken type, Metadata properties, Delivery delivery) {
                        try {
                            if(annot.protocolVersion() != -1 && annot.protocolVersion() != properties.protocolVersion()) {
                                return;
                            }

                            final Class<?>[] paramTypes = method.getParameterTypes();
                            Object[] params = new Object[paramTypes.length];
                            for(int i = 0; i < paramTypes.length; i++) {
                                if(paramTypes[i].isAssignableFrom(message.getClass())) {
                                    params[i] = message;
                                } else if(paramTypes[i].isAssignableFrom(Metadata.class)) {
                                    params[i] = properties;
                                } else if(paramTypes[i].isAssignableFrom(Delivery.class)) {
                                    params[i] = delivery;
                                }
                            }
                            method.invoke(listener, params);
                        } catch(IllegalAccessException e) {
                            throw new IllegalStateException(e);
                        } catch(InvocationTargetException e) {
                            if(e.getCause() instanceof RuntimeException) {
                                throw (RuntimeException) e.getCause();
                            } else {
                                throw new IllegalStateException(e);
                            }
                        }
                    }

                    @Override
                    public String toString() {
                        return listener + "." + method.getName();
                    }
                };

                subscribe(messageType, listener, handler, executor);
            }
        });
    }

    @Override
    public void unsubscribe(MessageHandler<?> handler) {
        synchronized(handlers) {
            handlers.removeIf(registered -> registered.handler == handler);
            handlersByType.entries().removeIf(registered -> registered.getValue().handler == handler);
        }
    }

    @Override
    public void unsubscribe(MessageListener listener) {
        if(listener == null) return;
        synchronized(handlers) {
            handlers.removeIf(registered -> registered.listener == listener);
            handlersByType.entries().removeIf(registered -> registered.getValue().listener == listener);
        }
    }

    private class MultiDispatcher extends DefaultConsumer {

        private SettableFuture cancelled;

        MultiDispatcher() {
            super(client.getChannel());
        }

        void awaitTermination(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
            if(cancelled != null) {
                cancelled.get(timeout, unit);
            }
        }

        @Override
        public void handleConsumeOk(String consumerTag) {
            cancelled = SettableFuture.create();
        }

        @Override
        public void handleCancelOk(String consumerTag) {
            if(cancelled != null) {
                cancelled.set(true);
            }
        }

        @Override
        public void handleDelivery(String consumerTag, Envelope envelope, final AMQP.BasicProperties amqProperties, byte[] body) throws IOException {
            try {
                client.getChannel().basicAck(envelope.getDeliveryTag(), false);

                final TypeToken<? extends Message> type;
                try {
                    type = messageRegistry.resolve(amqProperties.getType(), Metadata.modelName(amqProperties));
                } catch(NoSuchMessageException e) {
                    // Probably a newer protocol
                    logger.warning("Skipping unknown message type: " + e.getMessage());
                    return;
                }

                final Collection<RegisteredHandler<?>> matchingHandlers;
                synchronized(handlers) {
                    matchingHandlers = handlersByType.allAssignableFrom(type);
                }
                if(matchingHandlers.isEmpty()) return;

                final String json = new String(body, Charsets.UTF_8);
                final Message message = gson.fromJson(json, type.getType());
                final Metadata properties = new Metadata(amqProperties);
                final Delivery delivery = new Delivery(client, consumerTag, envelope);

                if(logger.isLoggable(Level.FINE)) {
                    logger.fine("Received message " + properties.getType() + "\nMetadata: " + properties + "\n" + prettyGson.toJson(json));
                }

                for(final RegisteredHandler handler : matchingHandlers) {
                    if(suspended && handler.listener != null &&
                       !handler.listener.listenWhileSuspended()) continue;

                    logger.fine("Dispatching " + amqProperties.getType() + " to " + handler.handler.getClass());
                    if(handler.executor == null) {
                        exceptionHandler.run(() -> handler.handler.handleDelivery(message, type, properties, delivery));
                    } else {
                        handler.executor.execute(() -> {
                            synchronized(handlers) {
                                // Double check from the handler's executor that it is still registered.
                                // This makes it much less likely to dispatch a message to a handler
                                // after it unsubs. It should work perfectly if the handler unsubs on
                                // the same thread it handles messages on.
                                if(!handlers.contains(handler)) return;
                            }
                            exceptionHandler.run(() -> handler.handler.handleDelivery(message, type, properties, delivery));
                        });
                    }
                }
            } catch(Throwable t) {
                logger.log(Level.SEVERE, "Exception dispatching AMQP message", t);
                // Don't let any exceptions through to the AMQP driver or it will close the channel
            }
        }
    }
}
