/*
 * 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.hive;

import com.facebook.presto.hive.shaded.org.apache.thrift.transport.TSocket;
import com.facebook.presto.hive.shaded.org.apache.thrift.transport.TTransport;
import com.facebook.presto.hive.shaded.org.apache.thrift.transport.TTransportException;
import com.google.common.net.HostAndPort;
import com.google.common.primitives.Ints;
import com.google.inject.Inject;
import io.airlift.units.Duration;

import javax.annotation.Nullable;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.net.Socket;
import java.net.SocketAddress;

import static com.google.common.base.Preconditions.checkNotNull;

public class HiveMetastoreClientFactory
{
    private final HostAndPort socksProxy;
    private final Duration timeout;

    public HiveMetastoreClientFactory(@Nullable HostAndPort socksProxy, Duration timeout)
    {
        this.socksProxy = socksProxy;
        this.timeout = checkNotNull(timeout, "timeout is null");
    }

    @Inject
    public HiveMetastoreClientFactory(HiveClientConfig config)
    {
        this(config.getMetastoreSocksProxy(), config.getMetastoreTimeout());
    }

    @SuppressWarnings("SocketOpenedButNotSafelyClosed")
    public HiveMetastoreClient create(String host, int port)
            throws TTransportException
    {
        TTransport transport;
        if (socksProxy == null) {
            transport = new TTransportWrapper(new TSocket(host, port, (int) timeout.toMillis()), host);
            transport.open();
        }
        else {
            SocketAddress address = InetSocketAddress.createUnresolved(socksProxy.getHostText(), socksProxy.getPort());
            Socket socks = new Socket(new Proxy(Proxy.Type.SOCKS, address));
            try {
                socks.connect(InetSocketAddress.createUnresolved(host, port), (int) timeout.toMillis());
                socks.setSoTimeout(Ints.checkedCast(timeout.toMillis()));
            }
            catch (IOException e) {
                throw rewriteException(new TTransportException(e), host);
            }
            try {
                transport = new TTransportWrapper(new TSocket(socks), host);
            }
            catch (TTransportException e) {
                throw rewriteException(e, host);
            }
        }

        return new HiveMetastoreClient(transport);
    }

    private static TTransportException rewriteException(TTransportException e, String host)
    {
        return new TTransportException(e.getType(), String.format("%s: %s", host, e.getMessage()), e.getCause());
    }

    private static class TTransportWrapper
            extends TTransport
    {
        private final TTransport transport;
        private final String host;

        TTransportWrapper(TTransport transport, String host)
        {
            this.transport = transport;
            this.host = host;
        }

        @Override
        public boolean isOpen()
        {
            return transport.isOpen();
        }

        @Override
        public boolean peek()
        {
            return transport.peek();
        }

        @Override
        public byte[] getBuffer()
        {
            return transport.getBuffer();
        }

        @Override
        public int getBufferPosition()
        {
            return transport.getBufferPosition();
        }

        @Override
        public int getBytesRemainingInBuffer()
        {
            return transport.getBytesRemainingInBuffer();
        }

        @Override
        public void consumeBuffer(int len)
        {
            transport.consumeBuffer(len);
        }

        @Override
        public void close()
        {
            transport.close();
        }

        @Override
        public void open()
                throws TTransportException
        {
            try {
                transport.open();
            }
            catch (TTransportException e) {
                throw rewriteException(e, host);
            }
        }

        @Override
        public int readAll(byte[] bytes, int off, int len)
                throws TTransportException
        {
            try {
                return transport.readAll(bytes, off, len);
            }
            catch (TTransportException e) {
                throw rewriteException(e, host);
            }
        }

        @Override
        public int read(byte[] bytes, int off, int len)
                throws TTransportException
        {
            try {
                return transport.read(bytes, off, len);
            }
            catch (TTransportException e) {
                throw rewriteException(e, host);
            }
        }

        @Override
        public void write(byte[] bytes)
                throws TTransportException
        {
            try {
                transport.write(bytes);
            }
            catch (TTransportException e) {
                throw rewriteException(e, host);
            }
        }

        @Override
        public void write(byte[] bytes, int off, int len)
                throws TTransportException
        {
            try {
                transport.write(bytes, off, len);
            }
            catch (TTransportException e) {
                throw rewriteException(e, host);
            }
        }

        @Override
        public void flush()
                throws TTransportException
        {
            try {
                transport.flush();
            }
            catch (TTransportException e) {
                throw rewriteException(e, host);
            }
        }
    }
}

@@ -1166,74 +1166,8 @@
 s;%0A%0A
-import static com.google.common.base.Preconditions.checkNotNull;%0A%0A
 publ
@@ -1264,32 +1264,27 @@
 e final 
-Duration
+int
  timeout
 ;%0A%0A    p
@@ -1275,16 +1275,22 @@
  timeout
+Millis
 ;%0A%0A    p
@@ -1442,50 +1442,52 @@
 eout
- = checkNotNull(timeout, %22timeout is null%22
+Millis = Ints.checkedCast(timeout.toMillis()
 );%0A 
@@ -1661,64 +1661,8 @@
  %7D%0A%0A
-    @SuppressWarnings(%22SocketOpenedButNotSafelyClosed%22)%0A
     
@@ -1896,30 +1896,24 @@
 t, port,
- (int)
  timeout
 .toMilli
@@ -1904,27 +1904,22 @@
  timeout
-.to
 Millis
-()
 ), host)
@@ -1997,184 +1997,45 @@
 cket
-Address address = InetSocketAddress.createUnresolved(socksProxy.getHostText(), socksProxy.getPort());%0A            
+ socks = createSocks
 Socket
- 
+(
 socks
- = new Socket(new Proxy(Proxy.Type.SOCKS, address)
+Proxy
 );%0A 
@@ -2132,22 +2132,16 @@
 rt),
- (int)
  timeout
 .toM
@@ -2132,35 +2132,30 @@
 rt), timeout
-.to
 Millis
-()
 );%0A         
@@ -2184,44 +2184,21 @@
 out(
-Ints.checkedCast(
 timeout
-.to
 Millis
-())
 );%0A 
@@ -2611,32 +2611,279 @@
 nsport);%0A    %7D%0A%0A
+    private static Socket createSocksSocket(HostAndPort proxy)%0A    %7B%0A        SocketAddress address = InetSocketAddress.createUnresolved(proxy.getHostText(), proxy.getPort());%0A        return new Socket(new Proxy(Proxy.Type.SOCKS, address));%0A    %7D%0A%0A
     private stat

Fix warnings in HiveMetastoreClientFactory