package com.atlassian.bitbucket.mesh.request;

import com.atlassian.bitbucket.mesh.rpc.MetadataKeys;
import io.grpc.*;

import javax.annotation.Nonnull;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * A {@link RequestDescriptor} for a gRPC {@link ServerCall}, with callbacks to track the number of bytes and
 * messages sent and received, as well as the {@link Status} returned when the call is closed.
 */
class GrpcRequestDescriptor implements RequestDescriptor {

    private final String action;
    private final AtomicLong bytesReceived;
    private final AtomicLong bytesSent;
    private final String agent;
    private final ServerCall<?, ?> call;
    private final String correlationId;
    private final AtomicInteger messagesReceived;
    private final AtomicInteger messagesSent;
    private final String remoteAddress;
    private final String user;

    private volatile Integer result;

    GrpcRequestDescriptor(ServerCall<?, ?> call, Metadata headers) {
        this.call = call;

        action = formatAction(call.getMethodDescriptor());
        agent = headers.get(MetadataKeys.KEY_REMOTE_CLIENT);
        bytesReceived = new AtomicLong();
        bytesSent = new AtomicLong();
        correlationId = headers.get(MetadataKeys.KEY_REMOTE_REQUEST);
        messagesReceived = new AtomicInteger();
        messagesSent = new AtomicInteger();
        remoteAddress = formatAddress(call.getAttributes());
        user = headers.get(MetadataKeys.KEY_REMOTE_USER);
    }

    @Nonnull
    @Override
    public String getAction() {
        return action;
    }

    @Override
    public String getAgent() {
        return agent;
    }

    @Override
    public long getBytesReceived() {
        return bytesReceived.get();
    }

    @Override
    public long getBytesSent() {
        return bytesSent.get();
    }

    @Override
    public String getCorrelationId() {
        return correlationId;
    }

    @Override
    public int getMessagesReceived() {
        return messagesReceived.get();
    }

    @Override
    public int getMessagesSent() {
        return messagesSent.get();
    }

    @Nonnull
    @Override
    public String getProtocol() {
        return "grpc";
    }

    @Nonnull
    @Override
    public Object getRawRequest() {
        return call;
    }

    @Override
    public String getRemoteAddress() {
        return remoteAddress;
    }

    @Override
    public Integer getResult() {
        return result;
    }

    @Override
    public String getUser() {
        return user;
    }

    void onBytesReceived(long bytes) {
        bytesReceived.addAndGet(bytes);
    }

    void onBytesSent(long bytes) {
        bytesSent.addAndGet(bytes);
    }

    void onClose(Status status) {
        result = status.getCode().value();
    }

    void onMessageReceived() {
        messagesReceived.incrementAndGet();
    }

    void onMessageSent() {
        messagesSent.incrementAndGet();
    }

    private static String formatAction(MethodDescriptor<?, ?> descriptor) {
        String fullMethodName = descriptor.getFullMethodName();

        int split = fullMethodName.lastIndexOf('.');
        String action = split == -1 ? fullMethodName : fullMethodName.substring(split + 1);

        return "\"" + action + "\"";
    }

    private static String formatAddress(Attributes attributes) {
        SocketAddress address = attributes.get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR);
        if (address instanceof InetSocketAddress) {
            return ((InetSocketAddress) address).getAddress().getHostAddress();
        }

        return address == null ? null : address.toString();
    }
}
