package com.jscape.inet.ssh.protocol.v2.marshaling;

import com.jscape.inet.ssh.protocol.messages.IdentificationString;
import com.jscape.inet.ssh.protocol.messages.Message;
import com.jscape.util.ArrayTools;
import com.jscape.util.IOTools;
import com.jscape.util.MessageFormat;
import com.jscape.util.io.DataCodec;
import com.jscape.util.io.primitive.CommonLineReader;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.StringReader;
import java.nio.charset.StandardCharsets;

/**
 * The SSH protocol identification string codec.
 *
 * @author Alex Usun
 */
public class IdentificationStringCodec
        implements DataCodec<Message> {

    private static final int MAX_LINE_LENGTH = 255;

    private static final String EOL = "\r\n";

    private static final byte[] IDENTIFICATION_STRING_PREFIX = {'S', 'S', 'H', '-'};
    private static final String TEMPLATE = "SSH-{0}-{1}{2}\r\n";
    private static final String KEX_TEMPLATE = "SSH-{0}-{1}{2}";
    private static final char VERSION_BOUND = '-';
    private static final char COMMENTS_BOUND = ' ';

    /**
     * {@inheritDoc}
     */
    public Message read(InputStream in)
            throws IOException {
        CommonLineReader reader = new CommonLineReader(in, MAX_LINE_LENGTH);
        StringBuilder bannerBuffer = new StringBuilder();
        byte[] line;
        int i = 0;
        while (!ArrayTools.starts(line = CommonLineReader.readLine(reader), IDENTIFICATION_STRING_PREFIX)) {
            if (i++ > 0) {
                bannerBuffer.append(EOL);
            }

            bannerBuffer.append(new String(line, StandardCharsets.UTF_8));
        }

        reader.consumeLF();

        String banner = bannerBuffer.length() > 0 ? bannerBuffer.toString() : null;

        StringReader identificationReader = new StringReader(new String(line, StandardCharsets.US_ASCII));
        String protocolVersion = parseProtocolVersion(identificationReader);
        String softwareVersion = parseSoftwareVersion(identificationReader);
        String comments = parseComments(identificationReader);

        try {
            IdentificationString message = new IdentificationString(banner, protocolVersion, softwareVersion, comments);
            message.rawData = line;
            
            return message;
        } catch (Exception e) {
            throw IOTools.ioException(e);
        }
    }

    /**
     * {@inheritDoc}
     */
    public void write(Message value, OutputStream out)
            throws IOException {
        IdentificationString message = (IdentificationString) value;
        writeBanner(message.banner, out);
        writeIdentificationString(message, out);

        out.flush();
    }

    private void writeBanner(String banner, OutputStream out)
            throws IOException {
        if (banner != null && !banner.isEmpty()) {
            String bannerToWrite = banner.endsWith(EOL) ? banner : banner + EOL;
            out.write(bannerToWrite.getBytes(StandardCharsets.UTF_8));
        }
    }

    private void writeIdentificationString(IdentificationString value, OutputStream out)
            throws IOException {
        byte[] data = format(value);
        out.write(data);
        
        value.rawData = formatForDHKeyExchange(value);
    }

    private byte[] format(IdentificationString string) {
        String data = MessageFormat.format(TEMPLATE,
                                    string.protocolVersion,
                                    string.softwareVersion,
                                    string.comments != null && string.comments.length() > 0 ? COMMENTS_BOUND + string.comments : "");
        
        return data.getBytes(StandardCharsets.US_ASCII);
    }
    
    private byte[] formatForDHKeyExchange(IdentificationString string) {
        String data = MessageFormat.format(KEX_TEMPLATE,
                                           string.protocolVersion,
                                           string.softwareVersion,
                                           string.comments != null && string.comments.length() > 0 ? COMMENTS_BOUND + string.comments : "");
        
        return data.getBytes(StandardCharsets.US_ASCII);
    }

    private String parseProtocolVersion(StringReader reader)
            throws IOException {
        readUntil(VERSION_BOUND, reader);
        String version = readUntil(VERSION_BOUND, reader);

        if (version == null || version.length() == 0) {
            throw new IOException("Invalid identification string: protocol version not found.");
        }

        return version;
    }

    private String parseSoftwareVersion(StringReader reader)
            throws IOException {
        String version = readUntil(COMMENTS_BOUND, reader);

//        if (version == null || version.length() == 0) {
//            throw new IOException("Invalid identification string: software version not found.");
//        }

        return version != null ? version : "";
    }

    private String parseComments(StringReader reader)
            throws IOException {
        StringBuilder buffer = new StringBuilder();
        int ch;
        while ((ch = reader.read()) != IOTools.EOF) {
            buffer.append((char) (ch & 0xFFFF));
        }

        return buffer.length() > 0 ? buffer.toString() : null;
    }

    private String readUntil(char bound, StringReader reader)
            throws IOException {
        StringBuilder buffer = new StringBuilder();
        int ch;
        while ((ch = reader.read()) != IOTools.EOF && ch != bound) {
            buffer.append((char) (ch & 0xFFFF));
        }

        return ch != -1 || buffer.length() > 0 ? buffer.toString() : null;
    }

}
