package com.jscape.inet.scp;

import com.jscape.inet.scp.protocol.marshaling.CommandFormat;
import com.jscape.inet.scp.protocol.marshaling.MessageCodec;
import com.jscape.inet.scp.protocol.marshaling.Messages;
import com.jscape.inet.scp.protocol.messages.*;
import com.jscape.inet.ssh.protocol.v2.connection.SessionChannelConnection;
import com.jscape.inet.ssh.protocol.v2.connection.SessionConnection;
import com.jscape.util.FileTools;
import com.jscape.util.IOTools;
import com.jscape.util.net.connection.Connector;
import com.jscape.util.net.connection.SyncRawConnection;
import com.jscape.util.service.AbstractService;

import java.io.*;
import java.util.logging.Level;
import java.util.logging.Logger;

import static com.jscape.inet.scp.FileService.OperationException.wrap;
import static com.jscape.util.Parameters.assertGreater;
import static com.jscape.util.Parameters.assertNotNull;

/**
 * @author Alex Usun
 */
public class ScpFileService
        extends AbstractService<Void>
        implements FileService {

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

    private static final String DEFAULT_DIRECTORY_MODE = "0755";
    private static final String DEFAULT_FILE_MODE = "0644";

    private static final int BUFFER_FACTOR = 3;
    private static final int STREAM_BUFFER_SIZE = 8 * 1024;

    private final Connector<SessionConnection> connector;
    private final int transferBufferSizeBytes;
    private final Logger logger;
    private SessionConnection sessionConnection;
    private volatile SessionChannelConnection channelConnection;
    private InputStream in;
    private OutputStream out;

    public ScpFileService(Connector<SessionConnection> connector, int transferBufferSizeBytes, Logger logger) {
        assertNotNull(connector);
        this.connector = connector;

        assertGreater(transferBufferSizeBytes, 0, "Bad transfer buffer size value.");
        this.transferBufferSizeBytes = transferBufferSizeBytes;

        assertNotNull(logger);
        this.logger = logger;
    }

    @Override
    public boolean available() {
        return !this.sessionConnection.closed();
    }

    @Override
    public Long readFile(String sourcePath, OutputStream destination, boolean fileAttributesPreservingRequired, TransferListener listener)
            throws OperationException {
        assertNotNull(listener);

        try {
            openChannelConnection();

            send(new ReadFileCommand(sourcePath, fileAttributesPreservingRequired));
            sendOk();

            Long modificationTime = receiveModificationTime(fileAttributesPreservingRequired);

            FileMessage fileMessage = receive();
            sendOk();

            listener.onTransferStarted(fileMessage.fileName, fileMessage.fileLength);
            transfer(this.in, destination, fileMessage.fileLength, listener);

            receiveOk();
            sendOk();

            return modificationTime;
        } catch (Exception e) {
            throw wrap(e);
        } finally {
            closeChannelConnection();
        }
    }

    @Override
    public void readFiles(String wildcard, File directory, boolean fileAttributesPreservingRequired, TransferListener listener)
            throws OperationException {
        assertNotNull(listener);

        try {
            openChannelConnection();

            send(new ReadFilesCommand(wildcard, fileAttributesPreservingRequired));
            sendOk();

            readEntries(directory, false, listener);
        } catch (Exception e) {
            throw wrap(e);
        } finally {
            closeChannelConnection();
        }
    }

    @Override
    public void readDirectory(String sourcePath, File destination, boolean fileAttributesPreservingRequired, TransferListener listener)
            throws OperationException {
        assertNotNull(listener);

        try {
            openChannelConnection();
            send(new ReadDirectoryCommand(sourcePath, fileAttributesPreservingRequired));
            sendOk();

            Long modificationTime = receiveModificationTime(fileAttributesPreservingRequired);
            DirectoryMessage message = receive();
            File directory = readDirectory(message, destination, listener);
            setModificationTime(directory, modificationTime);
        } catch (Exception e) {
            throw wrap(e);
        } finally {
            closeChannelConnection();
        }
    }

    @Override
    public void writeFile(InputStream source,
                          String destinationPath,
                          String destinationName,
                          long lengthBytes,
                          String mode,
                          Long modificationDate,
                          TransferListener listener)
            throws OperationException {
        try {
            openChannelConnection();

            destinationPath = destinationPath.trim();
            send(new WriteFileCommand(destinationPath.isEmpty() ? "." : destinationPath, modificationDate != null));
            receiveOk();

            if (modificationDate != null) {
                exchange(new TimeMessage(modificationDate, modificationDate));
            }

            exchange(new FileMessage(destinationName, lengthBytes, mode != null ? mode : DEFAULT_FILE_MODE));
            listener.onTransferStarted(destinationName, lengthBytes);
            transfer(source, this.out, lengthBytes, listener);
            exchange(new OkMessage());
        } catch (Exception e) {
            throw wrap(e);
        } finally {
            closeChannelConnection();
        }
    }

    @Override
    public void writeDirectory(File source, String destinationPath, String destinationName, String mode, boolean fileAttributesPreservingRequired, TransferListener listener)
            throws OperationException {
        try {
            openChannelConnection();
            send(new WriteDirectoryCommand(destinationPath, fileAttributesPreservingRequired));
            receiveOk();

            if (fileAttributesPreservingRequired) {
                exchange(new TimeMessage(source.lastModified(), source.lastModified()));
            }

            exchange(new DirectoryMessage(destinationName, mode != null ? mode : DEFAULT_DIRECTORY_MODE));
            writeEntries(source, fileAttributesPreservingRequired, listener);
        } catch (Exception e) {
            throw wrap(e);
        } finally {
            closeChannelConnection();
        }
    }

    @Override
    public void cancelTransferOperation() {
        closeChannelConnection();
    }

    public SessionConnection sessionConnection() {
        return this.sessionConnection;
    }

    @Override
    public String toString() {
        return "ScpFileService {"
               + "connector=" + this.connector
               + ", transferBufferSizeBytes=" + this.transferBufferSizeBytes
               + '}';
    }

    @Override
    protected void actualStart()
            throws Exception {
        this.sessionConnection = this.connector.connect();
        logServiceStarted();
    }

    @Override
    protected void actualStop() {
        this.sessionConnection.close();
        logServiceStopped();
    }

    private void openChannelConnection()
            throws Exception {
        this.channelConnection = this.sessionConnection.openSessionChannel(true);
        this.in = new BufferedInputStream(new SyncRawConnection.ConnectionInputStream(this.channelConnection));
        this.out = new BufferedOutputStream(new SyncRawConnection.ConnectionOutputStream(this.channelConnection));
    }

    private void closeChannelConnection() {
        if (this.channelConnection != null) {
            this.channelConnection.close();
            this.channelConnection = null;
        }
    }

    private void send(Command command)
            throws Exception {
        logOutgoingCommand(command);
        String commandLine = CommandFormat.format(command);
        this.channelConnection.exec(commandLine, true);
    }

    private void send(Message message)
            throws Exception {
        logOutgoingMessage(message);
        MESSAGE_CODEC.write(message, this.out);
    }

    private void sendOk()
            throws Exception {
        send(new OkMessage());
    }

    @SuppressWarnings("unchecked")
    private <M extends Message> M receive()
            throws Exception {
        Message message = MESSAGE_CODEC.read(this.in);
        logIncomingMessage(message);

        assertNoErrors(message);

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

    private void receiveOk()
            throws Exception {
        this.<OkMessage>receive();
    }

    private Long receiveModificationTime(boolean fileAttributesPreservingRequired)
            throws Exception {
        if (!fileAttributesPreservingRequired) {
            return null;
        }

        TimeMessage message = receive();
        sendOk();

        return message.modificationDate;
    }

    private void assertNoErrors(Message message)
            throws ScpException {
        if (message instanceof WarningMessage) {
            throw new ScpException(((WarningMessage) message).message);
        }

        if (message instanceof FatalErrorMessage) {
            throw new ScpException(((FatalErrorMessage) message).message);
        }
    }

    private void exchange(Message message)
            throws Exception {
        send(message);
        receiveOk();
    }

    private File readFile(FileMessage message, File parent, TransferListener listener)
            throws Exception {
        OutputStream out = null;

        try {
            File file = new File(parent, message.fileName);
            out = applyBuffering(new FileOutputStream(file));

            sendOk();

            listener.onTransferStarted(message.fileName, message.fileLength);
            transfer(this.in, out, message.fileLength, listener);

            receiveOk();
            sendOk();

            return file;
        } finally {
            IOTools.close(out);
        }
    }

    private File readDirectory(DirectoryMessage message, File parent, TransferListener listener)
            throws Exception {
        File directory = new File(parent, message.directoryName);
        FileTools.createDirectory(directory);

        sendOk();

        readEntries(directory, true, listener);

        return directory;
    }

    private void readEntries(File parent, boolean transferExpected, TransferListener listener)
            throws Exception {
        Long modificationTime = null;

        try {
            do {
                Message message = receive();

                if (message instanceof TimeMessage) {
                    modificationTime = ((TimeMessage) message).modificationDate;
                    sendOk();
                    transferExpected = true;
                } else if (message instanceof FileMessage) {
                    transferExpected = true;
                    File file = readFile((FileMessage) message, parent, listener);
                    setModificationTime(file, modificationTime);
                    transferExpected = false;
                } else if (message instanceof DirectoryMessage) {
                    transferExpected = true;
                    File directory = readDirectory((DirectoryMessage) message, parent, listener);
                    setModificationTime(directory, modificationTime);
                    transferExpected = false;
                } else if (message instanceof EndDirectoryMessage) {
                    sendOk();
                    break;
                }
            } while (true);
        } catch (EOFException e) {
            process(e, transferExpected);
        }
    }

    private void process(EOFException e, boolean transferExpected)
            throws EOFException {
        if (transferExpected) {
            throw e;
        }
    }

    private void writeEntries(File parent, boolean fileAttributesPreservingRequired, TransferListener listener)
            throws Exception {
        File[] entries = parent.listFiles();

        if (entries == null) {
            return;
        }

        for (File entry : entries) {
            if (fileAttributesPreservingRequired) {
                exchange(new TimeMessage(entry.lastModified(), entry.lastModified()));
            }

            if (entry.isDirectory()) {
                writeDirectory(entry, fileAttributesPreservingRequired, listener);
            } else {
                writeFile(entry, listener);
            }
        }

        exchange(new EndDirectoryMessage());
    }

    private void writeDirectory(File directory, boolean preserveFileAttributes, TransferListener listener)
            throws Exception {
        exchange(new DirectoryMessage(directory.getName(), DEFAULT_DIRECTORY_MODE));
        writeEntries(directory, preserveFileAttributes, listener);
    }

    private void writeFile(File file, TransferListener listener)
            throws Exception {
        InputStream in = applyBuffering(new FileInputStream(file));

        try {
            exchange(new FileMessage(file.getName(), file.length(), DEFAULT_FILE_MODE));
            listener.onTransferStarted(file.getName(), file.length());
            transfer(in, this.out, file.length(), listener);
            exchange(new OkMessage());
        } finally {
            IOTools.close(in);
        }
    }

    private void transfer(InputStream source, OutputStream destination, long length, TransferListener listener)
            throws Exception {
        byte[] buffer = new byte[this.transferBufferSizeBytes];
        long bytesToRead = Math.min(buffer.length, length);
        long transferredBytes = 0;
        int read;

        while (transferredBytes < length && !eof(read = source.read(buffer, 0, (int) bytesToRead))) {
            destination.write(buffer, 0, read);
            transferredBytes += read;
            listener.onTransferProgress(transferredBytes);
            bytesToRead = Math.min(buffer.length, length - transferredBytes);
        }

        destination.flush();
        listener.onTransferCompleted(transferredBytes);
    }

    private void setModificationTime(File file, Long modificationTime) {
        if (modificationTime != null) {
            file.setLastModified(modificationTime);
        }
    }

    private boolean eof(int read) {
        return read == -1;
    }

    private InputStream applyBuffering(InputStream in) {
        return new BufferedInputStream(in, Math.max(this.transferBufferSizeBytes * BUFFER_FACTOR, STREAM_BUFFER_SIZE));
    }

    private OutputStream applyBuffering(OutputStream out) {
        return new BufferedOutputStream(out, Math.max(this.transferBufferSizeBytes * BUFFER_FACTOR, STREAM_BUFFER_SIZE));
    }

    private void logServiceStarted() {
        if (this.logger.isLoggable(Level.FINE)) {
            this.logger.log(Level.FINE, String.format("SCP connection created: [%s <-> %s].",
                                                      this.sessionConnection.localAddress(),
                                                      this.sessionConnection.remoteAddress()));
        }
    }

    private void logServiceStopped() {
        if (this.logger.isLoggable(Level.FINE)) {
            this.logger.log(Level.FINE, String.format("SCP connection closed: [%s <-> %s].",
                                                      this.sessionConnection.localAddress(),
                                                      this.sessionConnection.remoteAddress()));
        }
    }

    private void logOutgoingCommand(Command command) {
        if (this.logger.isLoggable(Level.FINER)) {
            this.logger.log(Level.FINER,
                            String.format("Outgoing SCP command at [%s <-> %s]: %s.",
                                          this.sessionConnection.localAddress(),
                                          this.sessionConnection.remoteAddress(),
                                          command));
        }
    }

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

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