package com.jscape.inet.proxy;

import com.jscape.inet.proxy.socks.protocol.v4.marshaling.MessageCodec;
import com.jscape.inet.proxy.socks.protocol.v4.marshaling.Messages;
import com.jscape.inet.proxy.socks.protocol.v4.messages.CommandResponse;
import com.jscape.inet.proxy.socks.protocol.v4.messages.CommandResult;
import com.jscape.inet.proxy.socks.protocol.v4.messages.ConnectRequest;
import com.jscape.inet.proxy.socks.protocol.v4.messages.Message;
import com.jscape.util.net.TransportAddress;
import com.jscape.util.net.connection.Connection;
import com.jscape.util.net.connection.SyncRawConnection;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Inet4Address;
import java.util.logging.Level;
import java.util.logging.Logger;

import static com.jscape.util.Parameters.assertNotNull;

/**
 * @author Alex Usun
 */
public class Socks4Proxy
        implements Proxy {

    private static final MessageCodec MESSAGE_CODEC = Messages.init(new MessageCodec());

    private final Logger logger;

    public Socks4Proxy(Logger logger) {
        assertNotNull(logger);
        this.logger = logger;
    }

    @Override
    public SyncRawConnection openConnection(SyncRawConnection connection, TransportAddress address, String username, String password)
            throws IOException {
        InputStream in = new SyncRawConnection.ConnectionInputStream(connection);
        OutputStream out = new SyncRawConnection.ConnectionOutputStream(connection);

        send(new ConnectRequest(address.getPort(), (Inet4Address) address.inetAddress(), username), out, connection);
        CommandResponse response = receive(CommandResponse.class, in, connection);
        
        if (response.result != CommandResult.REQUEST_GRANTED.code) {
            throw new OperationException("SOCKS4 proxy error: " + response.result);
        }
        
        return connection;
    }

    private void send(Message message, OutputStream out, Connection connection)
            throws IOException {
        logOutgoingMessage(connection, message);
        MESSAGE_CODEC.write(message, out);
    }

    @SuppressWarnings("unchecked")
    private <M extends Message> M receive(Class<M> messageClass, InputStream in, Connection connection)
            throws IOException {
        Message message = MESSAGE_CODEC.read(messageClass, in);
        logIncomingMessage(connection, message);

        try {
            return (M) message;
        } catch (Exception e) {
            throw new IOException("Unexpected incoming message: " + message);
        }
    }

    private void logOutgoingMessage(Connection connection, Message message) {
        if (this.logger.isLoggable(Level.FINER)) {
            this.logger.log(Level.FINER, String.format("Outgoing SOCKS4 proxy message at [%s <-> %s]: %s.",
                                                       connection.localAddress(),
                                                       connection.remoteAddress(),
                                                       message));
        }
    }

    private void logIncomingMessage(Connection connection, Message message) {
        if (this.logger.isLoggable(Level.FINER)) {
            this.logger.log(Level.FINER,
                            String.format("Incoming SOCKS4 proxy message at [%s <-> %s]: %s.",
                                          connection.localAddress(),
                                          connection.remoteAddress(),
                                          message));
        }
    }
}
