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

import com.jscape.inet.ssh.protocol.v2.marshaling.primitive.StringCodec;
import com.jscape.inet.ssh.protocol.v2.messages.PublicKeyAlgorithm;
import com.jscape.inet.ssh.protocol.v2.messages.Signature;
import com.jscape.util.io.ByteArrayInputStream;
import com.jscape.util.io.ByteArrayOutputStream;
import com.jscape.util.io.DataCodec;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.Map;

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

/**
 * @author Alex Usun
 */
public class SignatureCodec
        implements DataCodec<Signature> {

    private final Map<String, Entry> entries;

    public SignatureCodec(Entry... entries) {
        this.entries = new HashMap<>();
        set(entries);
    }

    public SignatureCodec set(Entry... entries) {
        for (Entry entry : entries) {
            this.entries.put(entry.identifier, entry);
        }

        return this;
    }

    @Override
    public Signature read(InputStream in)
            throws IOException {
        byte[] signatureData = StringCodec.readValue(in);
        InputStream signatureDataStream = new ByteArrayInputStream(signatureData);
        String identifier = StringCodec.readUsAsciiValue(signatureDataStream);
        Entry entry = entryFor(identifier);

        return entry.codec.read(identifier, signatureDataStream);
    }

    @Override
    public void write(Signature value, OutputStream out)
            throws IOException {
        Entry entry = entryFor(value.identifier);
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        StringCodec.writeUsAsciiValue(value.identifier, buffer);
        entry.codec.write(value, buffer);

        StringCodec.writeValue(buffer.buffer(), 0, buffer.position(), out);
    }

    @Override
    public String toString() {
        return "SignatureCodec {"
                + "entries=" + this.entries
                + '}';
    }

    private Entry entryFor(String formatIdentifier)
            throws IOException {
        Entry entry = this.entries.get(formatIdentifier);

        if (entry == null) {
            throw new IOException(String.format("Unsupported signature format identifier: %s.", formatIdentifier));
        }

        return entry;
    }

    public static class Entry {

        public final String identifier;
        public final EntryCodec codec;

        public Entry(String identifier, EntryCodec codec) {
            assertNotNull(identifier);
            this.identifier = identifier;

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

        public Entry(PublicKeyAlgorithm publicKeyAlgorithm, EntryCodec codec) {
            this(publicKeyAlgorithm.identifier, codec);
        }

        @Override
        public String toString() {
            return "Entry {"
                    + "identifier='" + this.identifier + '\''
                    + ", codec=" + this.codec
                    + '}';
        }
    }

    public interface EntryCodec {

        Signature read(String identifier, InputStream in)
                throws IOException;

        void write(Signature value, OutputStream out)
                throws IOException;
    }
}
