package net.ranides.assira.generic;

import lombok.RequiredArgsConstructor;

import net.ranides.assira.collection.HashFunctions;
import net.ranides.assira.collection.maps.CustomMap;
import net.ranides.assira.reflection.*;
import net.ranides.assira.reflection.util.ReflectUtils;

import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.ObjectStreamClass;
import java.io.Serializable;
import java.util.Map;

/**
 * Subclasses of this class are serialized in non-standard way:
 * they save both their state, and bytecode.
 *
 * If you can't inherit from this class, you can save objects with bytecode, by serializing
 * object created by {@link #wrap(Serializable)}. Those wrapped objects will recreate both state
 * and code of the object during deserialization.
 */
public abstract class SerializableCode implements Serializable {

    private static final long serialVersionUID = 8L;

    protected boolean writeReplace = true;

    protected Object writeReplace() {
        return writeReplace ? SerializableCode.wrap(this) : this;
    }

    /**
     * Converts provided serializable object into special type, which will store both
     * state and bytecode inside stream, when serialized. During the deserialization "object"
     * will be recreated, including both state and code.
     *
     * @param object object
     * @return wrapper
     */
    public static Serializable wrap(Serializable object) {
        return new ResolveWrapper(new CodeContainer(object));
    }

    @RequiredArgsConstructor
    private static class ResolveWrapper implements Serializable {

        private static final long serialVersionUID = 8L;

        private final CodeContainer extern;

        private Object readResolve() {
            return extern.value;
        }

    }

    private static class CodeContainer implements Serializable {

        private static final IField WRITE_REPLACE = IClass.typeinfo(SerializableCode.class).field("writeReplace").get().accessible();

        private static final Map<byte[], IClass<?>> LOADED = new CustomMap<>(HashFunctions.BYTE_ARRAY);

        private static final long serialVersionUID = 8L;

        private byte[] bytecode;

        private transient Object value;

        public CodeContainer(Object value) {
            this.bytecode = ReflectUtils.getByteCode(value.getClass());
            this.value = value;
        }

        private void writeObject(java.io.ObjectOutputStream out) throws IOException {
            out.writeObject(bytecode);

            ObjectOutputStream substream = new ObjectOutputStream(out);
            if(value instanceof SerializableCode) {
                WRITE_REPLACE.set(value, false);
                substream.writeObject(value);
                WRITE_REPLACE.set(value, true);
            } else {
                substream.writeObject(value);
            }
            substream.flush();
        }

        private void readObject(java.io.ObjectInputStream in) throws IOException, ClassNotFoundException {
            synchronized (LOADED) {
                this.bytecode = (byte[]) in.readObject();
                IClass<?> type = LOADED.computeIfAbsent(bytecode, IClassLoader.getSystem()::defineIClass);
                this.value = new CodeInputStream(in, type).readObject();
            }
        }

    }

    private static class CodeInputStream extends ObjectInputStream {

        private final IClass<?> clazz;

        public CodeInputStream(InputStream in, IClass<?> clazz) throws IOException {
            super(in);
            this.clazz = clazz;
        }

        protected Class<?> resolveClass(ObjectStreamClass classDesc) throws ClassNotFoundException, IOException {
            return classDesc.getName().equals(clazz.name()) ? clazz.raw() : super.resolveClass(classDesc);
        }
    }

}
