package jetbrains.jetpad.json.adHoc;

import jetbrains.jetpad.json.JsonArray;
import jetbrains.jetpad.json.JsonBoolean;
import jetbrains.jetpad.json.JsonNull;
import jetbrains.jetpad.json.JsonNumber;
import jetbrains.jetpad.json.JsonObject;
import jetbrains.jetpad.json.JsonString;
import jetbrains.jetpad.json.JsonValue;

import java.util.HashMap;
import java.util.Map;

public class JsonSerializationContext {
  private Map<Class, Byte> myRegistrations = new HashMap<Class, Byte>();
  private Map<Byte, Serializer> mySerializers = new HashMap<Byte, Serializer>();

  public JsonSerializationContext() {
    IntegerSerializer integerSerializer = new IntegerSerializer();
    register(Integer.class, integerSerializer);

    StringSerializer stringSerializer = new StringSerializer(integerSerializer);
    register(String.class, stringSerializer);

    register(JsonNull.class, new JsonNullSerializer());
    register(JsonBoolean.class, new JsonBooleanSerializer());

    register(JsonString.class, new JsonStringSerializer(stringSerializer));
    register(JsonNumber.class, new JsonNumberSerializer(integerSerializer, new DoubleSerializer()));

    register(JsonObject.class, new JsonObjectSerializer(this));
    register(JsonArray.class, new JsonArraySerializer(this));
  }

  <T> void register(Class<T> type, Serializer<T> serializer) {
    if (myRegistrations.containsKey(type)) throw new IllegalStateException("double register serializer for " + type);

    byte id = serializer.hasId() ? serializer.getId() : (byte) myRegistrations.size();
    if (mySerializers.containsKey(id)) throw new IllegalStateException("duplicate serializer id " + id);

    myRegistrations.put(type, id);
    mySerializers.put(id, serializer);
  }

  public JsonValue read(byte[] input) {
    Serializer serializer = getSerializer(input, 0);
    JsonValue value = (JsonValue) serializer.read(input, 1);
    if (serializer.getPosition() != input.length) throw new JsonSerializationException();
    return value;
  }

  public byte[] write(JsonValue value) {
    Byte type = myRegistrations.get(value.getClass());
    if (type == null) throw new JsonSerializationException("unregistered class " + value.getClass());
    return mySerializers.get(type).write(value);
  }

  Serializer getSerializer(byte[] input, int offset) {
    byte type = input[offset];
    if (!mySerializers.containsKey(type)) throw new JsonSerializationException("offset=" + offset + ", type=" + type);
    return mySerializers.get(type);
  }

  <T> Serializer<T> getSerializer(Class<T> type) {
    return mySerializers.get(myRegistrations.get(type));
  }
}

@@ -423,62 +423,74 @@
 Map%3C
-Class, Byte%3E myRegistrations = new HashMap%3CClass, Byte
+Byte, Serializer%3E mySerializersById = new HashMap%3CByte, Serializer
 %3E();
@@ -496,36 +496,37 @@
 ;%0A  private Map%3C
-Byte
+Class
 , Serializer%3E my
@@ -536,16 +536,23 @@
 ializers
+ByClass
  = new H
@@ -550,36 +550,37 @@
 s = new HashMap%3C
-Byte
+Class
 , Serializer%3E();
@@ -1364,28 +1364,33 @@
   if (my
-Registration
+SerializersByClas
 s.contai
@@ -1543,28 +1543,33 @@
 byte) my
-Registration
+SerializersByClas
 s.size()
@@ -1583,32 +1583,36 @@
 f (mySerializers
+ById
 .containsKey(id)
@@ -1686,28 +1686,33 @@
 %0A%0A    my
-Registration
+SerializersByClas
 s.put(ty
@@ -1715,18 +1715,26 @@
 t(type, 
-id
+serializer
 );%0A    m
@@ -1745,16 +1745,20 @@
 ializers
+ById
 .put(id,
@@ -2092,34 +2092,51 @@
     
-Byte type = myRegistration
+Serializer serializer = mySerializersByClas
 s.ge
@@ -2164,20 +2164,26 @@
     if (
-type
+serializer
  == null
@@ -2267,35 +2267,33 @@
 ());%0A    return 
-myS
+s
 erializers.get(t
@@ -2285,27 +2285,16 @@
 rializer
-s.get(type)
 .write(v
@@ -2369,20 +2369,18 @@
    byte 
-type
+id
  = input
@@ -2411,16 +2411,20 @@
 ializers
+ById
 .contain
@@ -2420,36 +2420,34 @@
 yId.containsKey(
-type
+id
 )) throw new Jso
@@ -2498,21 +2498,17 @@
  %22, 
-type
+id
 =%22 + 
-type
+id
 );%0A 
@@ -2530,25 +2530,27 @@
 ializers
+ById
 .get(
-type
+id
 );%0A  %7D%0A%0A
@@ -2628,27 +2628,14 @@
 zers
-.get(myRegistration
+ByClas
 s.ge
@@ -2641,17 +2641,16 @@
 et(type)
-)
 ;%0A  %7D%0A%7D%0A

optimize JsonSerializationContext.getSerializer(Class)