package com.jscape.inet.proxy;

import com.jscape.inet.proxy.socks.protocol.v5.marshaling.MessageCodec;
import com.jscape.inet.proxy.socks.protocol.v5.marshaling.Messages;
import com.jscape.inet.proxy.socks.protocol.v5.messages.*;
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.util.logging.Level;
import java.util.logging.Logger;

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

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

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

    private final Logger logger;

    public Socks5Proxy(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);

        int authenticationMethodCode = selectAuthenticationMethod(connection, in, out);
        if (authenticationMethodCode == AuthenticationMethodType.PASSWORD.code) {
            authenticate(username, password, in, out, connection);
        }

        send(new ConnectRequest(address.inetAddress(), address.getPort()), out, connection);
        CommandResponse response = receive(CommandResponse.class, in, connection);

        if (response.status != CommandStatus.SUCCEEDED.code) {
            throw new OperationException("SOCKS5 proxy error: " + response.status);
        }
        
        return connection;
    }

    private int selectAuthenticationMethod(SyncRawConnection connection, InputStream in, OutputStream out)
            throws IOException {
        send(new AuthenticationMethodRequest(AuthenticationMethodType.NONE, AuthenticationMethodType.PASSWORD), out, connection);
        AuthenticationMethodResponse response = receive(AuthenticationMethodResponse.class, in, connection);

        if (response.method != AuthenticationMethodType.NONE.code && response.method != AuthenticationMethodType.PASSWORD.code) {
            throw new OperationException(String.format("SOCKS5 proxy error: unsupported authentication method (%s).", response.method));
        }

        return response.method;
    }

    private void authenticate(String username, String password, InputStream in, OutputStream out, SyncRawConnection connection)
            throws IOException {
        send(new PasswordAuthenticationRequest(username, password), out, connection);
        PasswordAuthenticationResponse response = receive(PasswordAuthenticationResponse.class, in, connection);

        if (response.statusCode != PasswordAuthenticationStatus.SUCCESS.code) {
            throw new OperationException("SOCKS5 proxy error: authentication failed.");
        }
    }

    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 SOCKS5 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 SOCKS5 proxy message at [%s <-> %s]: %s.",
                                          connection.localAddress(),
                                          connection.remoteAddress(),
                                          message));
        }
    }
}
