/*
 * 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 com.facebook.presto.connector.jmx;

import com.facebook.presto.spi.ConnectorColumnHandle;
import com.facebook.presto.spi.ConnectorRecordSetProvider;
import com.facebook.presto.spi.ConnectorSplit;
import com.facebook.presto.spi.InMemoryRecordSet;
import com.facebook.presto.spi.NodeManager;
import com.facebook.presto.spi.RecordSet;
import com.facebook.presto.spi.type.Type;
import com.facebook.presto.util.IterableTransformer;
import com.google.common.base.Function;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import io.airlift.slice.Slice;

import java.util.Arrays;
import javax.inject.Inject;
import javax.management.Attribute;
import javax.management.JMException;
import javax.management.MBeanServer;
import javax.management.ObjectName;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import static com.facebook.presto.util.Types.checkType;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

public class JmxRecordSetProvider
        implements ConnectorRecordSetProvider
{
    private final MBeanServer mbeanServer;
    private final String nodeId;

    @Inject
    public JmxRecordSetProvider(MBeanServer mbeanServer, NodeManager nodeManager)
    {
        this.mbeanServer = checkNotNull(mbeanServer, "mbeanServer is null");
        this.nodeId = checkNotNull(nodeManager, "nodeManager is null").getCurrentNode().getNodeIdentifier();
    }

    @Override
    public RecordSet getRecordSet(ConnectorSplit split, List<? extends ConnectorColumnHandle> columns)
    {
        JmxTableHandle tableHandle = checkType(split, JmxSplit.class, "split").getTableHandle();

        checkNotNull(columns, "columns is null");
        checkArgument(!columns.isEmpty(), "must provide at least one column");

        ImmutableMap.Builder<String, Type> builder = ImmutableMap.builder();
        for (ConnectorColumnHandle column : columns) {
            JmxColumnHandle jmxColumnHandle = checkType(column, JmxColumnHandle.class, "column");
            builder.put(jmxColumnHandle.getColumnName(), jmxColumnHandle.getColumnType());
        }
        ImmutableMap<String, Type> columnTypes = builder.build();

        List<List<Object>> rows;
        try {
            Map<String, Object> attributes = getAttributes(columnTypes.keySet(), tableHandle);
            List<Object> row = new ArrayList<>();
            // NOTE: data must be produced in the order of the columns parameter.  This code relies on the
            // fact that columnTypes is an ImmutableMap which is an order preserving LinkedHashMap under
            // the covers.
            for (Entry<String, Type> entry : columnTypes.entrySet()) {
                if (entry.getKey().equals("node")) {
                    row.add(nodeId);
                }
                else {
                    Object value = attributes.get(entry.getKey());
                    if (value == null) {
                        row.add(null);
                    }
                    else {
                        Class<?> javaType = entry.getValue().getJavaType();
                        if (javaType == boolean.class) {
                            if (value instanceof Boolean) {
                                row.add(value);
                            }
                            else {
                                // mbeans can lie about types
                                row.add(null);
                            }
                        }
                        else if (javaType == long.class) {
                            if (value instanceof Number) {
                                row.add(((Number) value).longValue());
                            }
                            else {
                                // mbeans can lie about types
                                row.add(null);
                            }
                        }
                        else if (javaType == double.class) {
                            if (value instanceof Number) {
                                row.add(((Number) value).doubleValue());
                            }
                            else {
                                // mbeans can lie about types
                                row.add(null);
                            }
                        }
                        else if (javaType == Slice.class) {
                            if (value.getClass().isArray()) {
                                // return a string representation of the array
                                if (value.getClass().getComponentType() == String.class) {
                                    row.add(Arrays.toString((String[]) value));
                                }
                                else if (value.getClass().getComponentType() == boolean.class) {
                                    row.add(Arrays.toString((boolean[]) value));
                                }
                                else if (value.getClass().getComponentType() == byte.class) {
                                    row.add(Arrays.toString((byte[]) value));
                                }
                                else if (value.getClass().getComponentType() == char.class) {
                                    row.add(Arrays.toString((char[]) value));
                                }
                                else if (value.getClass().getComponentType() == double.class) {
                                    row.add(Arrays.toString((double[]) value));
                                }
                                else if (value.getClass().getComponentType() == float.class) {
                                    row.add(Arrays.toString((float[]) value));
                                }
                                else if (value.getClass().getComponentType() == int.class) {
                                    row.add(Arrays.toString((int[]) value));
                                }
                                else if (value.getClass().getComponentType() == long.class) {
                                    row.add(Arrays.toString((long[]) value));
                                }
                                else if (value.getClass().getComponentType() == short.class) {
                                    row.add(Arrays.toString((short[]) value));
                                }
                            }
                            else {
                                row.add(value.toString());
                            }
                        }
                    }
                }
            }
            rows = ImmutableList.of(row);
        }
        catch (JMException e) {
            rows = ImmutableList.of();
        }

        return new InMemoryRecordSet(columnTypes.values(), rows);
    }

    private Map<String, Object> getAttributes(Set<String> uniqueColumnNames, JmxTableHandle tableHandle)
            throws JMException
    {
        ObjectName objectName = new ObjectName(tableHandle.getObjectName());

        String[] columnNamesArray = uniqueColumnNames.toArray(new String[uniqueColumnNames.size()]);

        return IterableTransformer.on(mbeanServer.getAttributes(objectName, columnNamesArray).asList())
                .uniqueIndex(attributeNameGetter())
                .transformValues(attributeValueGetter())
                .map();
    }

    private static Function<Attribute, String> attributeNameGetter()
    {
        return new Function<Attribute, String>()
        {
            @Override
            public String apply(Attribute attribute)
            {
                return attribute.getName();
            }
        };
    }

    private static Function<Attribute, Object> attributeValueGetter()
    {
        return new Function<Attribute, Object>()
        {
            @Override
            public Object apply(Attribute attribute)
            {
                return attribute.getValue();
            }
        };
    }
}

@@ -994,48 +994,8 @@
 er;%0A
-import com.google.common.base.Function;%0A
 impo
@@ -7933,25 +7933,25 @@
 ueIndex(
-a
+A
 ttribute
 NameGett
@@ -7942,28 +7942,25 @@
 ttribute
-NameGetter()
+::getName
 )%0A      
@@ -7990,654 +7990,51 @@
 ues(
-attributeValueGetter())%0A                .map();%0A    %7D%0A%0A    private static Function%3CAttribute, String%3E attributeNameGetter()%0A    %7B%0A        return new Function%3CAttribute, String%3E()%0A        %7B%0A            @Override%0A            public String apply(Attribute attribute)%0A            %7B%0A                return attribute.getName();%0A            %7D%0A        %7D;%0A    %7D%0A%0A    private static Function%3CAttribute, Object%3E attributeValueGetter()%0A    %7B%0A        return new Function%3CAttribute, Object%3E()%0A        %7B%0A            @Override%0A            public Object apply(Attribute attribute)%0A            %7B%0A                return attribute.getValue();%0A            %7D%0A        %7D
+Attribute::getValue)%0A                .map()
 ;%0A  

Remove getter functions from JmxRecordSetProvider