package com.jscape.inet.proxy.socks.protocol.v5.marshaling;

import com.jscape.inet.proxy.socks.protocol.v5.messages.Message;
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.Arrays;
import java.util.HashMap;
import java.util.Map;

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

/**
 * @author Alex Usun
 */
public class MessageCodec {

    private final Map<Class, Entry> readingEntries;
    private final Map<Class, Entry> writingEntries;

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

    public MessageCodec set(Entry... entries) {
        for (Entry entry : entries) {
            this.readingEntries.put(entry.readingMessageClass, entry);

            for (Class messageClass : entry.writingMessageClasses) {
                this.writingEntries.put(messageClass, entry);
            }
        }

        return this;
    }

    public Message read(Class messageClass, InputStream in)
            throws IOException {
        Entry entry = entryFor(messageClass);
        return entry.codec.read(in);
    }

    public void write(Message message, OutputStream out)
            throws IOException {
        Entry entry = entryFor(message);
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        entry.codec.write(message, buffer);
        out.write(buffer.buffer(), 0, buffer.position());
        out.flush();
    }

    private Entry entryFor(Class messageClass)
            throws IOException {
        Entry entry = this.readingEntries.get(messageClass);

        if (entry == null) {
            throw new IOException(String.format("Unsupported message class: %s.", messageClass));
        }

        return entry;
    }

    private Entry entryFor(Message message) {
        Entry entry = this.writingEntries.get(message.getClass());
    
        if (entry == null) {
            throw new IllegalArgumentException(String.format("Unsupported message: %s.", message.getClass()));
        }

        return entry;
    }

    public static class Entry {

        public final Class readingMessageClass;
        public final DataCodec<Message> codec;
        public final Class[] writingMessageClasses;

        public Entry(Class readingMessageClass, DataCodec<Message> codec, Class... writingMessageClasses) {
            assertNotNull(readingMessageClass);
            this.readingMessageClass = readingMessageClass;

            assertNotNull(codec);
            this.codec = codec;

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

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