package org.jruby;

import org.jruby.anno.JRubyMethod;
import org.jruby.internal.runtime.methods.AttrReaderMethod;
import org.jruby.ir.operands.UndefinedValue;
import org.jruby.ir.runtime.IRRuntimeHelpers;
import org.jruby.runtime.Block;
import org.jruby.runtime.ClassIndex;
import org.jruby.runtime.ObjectAllocator;
import org.jruby.runtime.ThreadContext;
import org.jruby.runtime.Visibility;
import org.jruby.runtime.builtin.IRubyObject;
import org.jruby.runtime.ivars.VariableAccessor;
import org.jruby.runtime.ivars.VariableTableManager;
import org.jruby.runtime.marshal.MarshalDumper;
import org.jruby.runtime.marshal.MarshalLoader;
import org.jruby.util.io.RubyInputStream;
import org.jruby.util.io.RubyOutputStream;

import java.util.LinkedHashSet;
import java.util.Map;
import java.util.stream.Collectors;

import static org.jruby.RubyBasicObject.rbInspect;
import static org.jruby.RubyHash.newSmallHash;
import static org.jruby.api.Convert.asFixnum;
import static org.jruby.api.Convert.toLong;
import static org.jruby.api.Convert.toSymbol;
import static org.jruby.api.Create.allocArray;
import static org.jruby.api.Create.newArray;
import static org.jruby.api.Create.newEmptyArray;
import static org.jruby.api.Create.newEmptyHash;
import static org.jruby.api.Create.newSmallHash;
import static org.jruby.api.Create.newString;
import static org.jruby.api.Define.defineClass;
import static org.jruby.api.Error.argumentError;
import static org.jruby.api.Error.keywordError;
import static org.jruby.api.Error.typeError;
import static org.jruby.ir.runtime.IRRuntimeHelpers.setCallInfo;
import static org.jruby.runtime.Arity.checkArgumentCount;
import static org.jruby.runtime.Helpers.invokedynamic;
import static org.jruby.runtime.ThreadContext.CALL_KEYWORD;
import static org.jruby.runtime.ThreadContext.hasNonemptyKeywords;
import static org.jruby.runtime.invokedynamic.MethodNames.HASH;
import static org.jruby.util.RubyStringBuilder.str;

public class RubyData {

    private static final String MEMBERS_KEY = "__members__";
    private static final String ACCESSORS_KEY = "__accessors__";

    public static RubyClass createDataClass(ThreadContext context, RubyClass Object) {
        RubyClass Data = defineClass(context, "Data", Object, ObjectAllocator.NOT_ALLOCATABLE_ALLOCATOR)
                .classIndex(ClassIndex.DATA)
                .defineMethods(context, RubyData.class);

        Data.getSingletonClass().undefMethods(context, "new");

        return Data;
    }

    @JRubyMethod(meta = true, rest = true)
    public static RubyClass define(ThreadContext context, IRubyObject self, IRubyObject[] args, Block block) {
        LinkedHashSet<RubySymbol> keySet = new LinkedHashSet<>();
        for (int i = 0 ; i < args.length ; i++) {
            RubySymbol mem = toSymbol(context, args[i]);
            if (mem.validAttrsetName()) {
                throw argumentError(context, "invalid data member: " + mem);
            }
            if (keySet.contains(mem)) {
                throw argumentError(context, "duplicate member: " + mem);
            }
            keySet.add(mem);
        }

        RubyClass dataClass = newDataStruct(context, (RubyClass) self, keySet);

        if (block.isGiven()) {
            dataClass.module_eval(context, block);
        }

        return dataClass;
    }

    @JRubyMethod(keywords = true, rest = true)
    public static void initialize(ThreadContext context, IRubyObject self, IRubyObject[] args) {
        ThreadContext.resetCallInfo(context); // we don't directly use callInfo here

        RubyBasicObject selfObj = (RubyBasicObject) self;

        selfObj.checkFrozen();
        RubyArray<RubySymbol> members = getStructMembers(self);
        int numMembers = members.size();

        if (args.length == 0) {
            if (numMembers > 0) {
                throw keywordError(context, "missing", members);
            }

            selfObj.setFrozen(true);
            return;
        }

        if (args.length > 1 || !(args[0] instanceof RubyHash)) {
            throw argumentError(context, args.length, 0, 0);
        }

        RubyHash hash = (RubyHash) args[0];

        if (hash.size() < numMembers) {
            RubyArray missing = (RubyArray) members.op_diff(context, hash.keys(context));
            throw keywordError(context, "missing", missing);
        }

        RubyArray[] unknownKeywordsPtr = {null};
        RubyClass metaClass = selfObj.getMetaClass();
        VariableTableManager vtm = getVariableManagerFromClass(metaClass);
        Map<String, VariableAccessor> variableAccessors = vtm.getVariableAccessorsForRead();

        hash.visitAll(context, (c, h, k, v, i) -> {
            String keyString = toSymbol(context, k).idString();
            VariableAccessor variableAccessor = variableAccessors.get(keyString);
            if (variableAccessor != null) {
                variableAccessor.set(self, v);
            } else {
                RubyArray unknownKeywords = unknownKeywordsPtr[0];
                if (unknownKeywords == null) {
                    unknownKeywordsPtr[0] = unknownKeywords = newEmptyArray(context);
                }
                unknownKeywords.append(context, k);
            }
        });
        selfObj.setFrozen(true);
        if (unknownKeywordsPtr[0] != null) {
            throw keywordError(context, "unknown", unknownKeywordsPtr[0]);
        }
    }

    @JRubyMethod(name = "initialize_copy")
    public static IRubyObject initialize_copy(ThreadContext context, IRubyObject copy, IRubyObject original) {
        if (original != copy) {
            original.getMetaClass().getVariableTableManager().syncVariables((RubyBasicObject) copy, original);
        }

        copy.setFrozen(true);

        return copy;
    }

    @JRubyMethod(name = "==")
    public static IRubyObject op_equal(ThreadContext context, IRubyObject self, IRubyObject other) {
        return checkDataEquality(context, (RubyBasicObject) self, (RubyBasicObject) other, "==", RubyData::equalData);
    }

    @JRubyMethod(name = "eql?")
    public static IRubyObject eql(ThreadContext context, IRubyObject self, IRubyObject other) {
        return checkDataEquality(context, (RubyBasicObject) self, (RubyBasicObject) other, "eql?", RubyData::eqlData);
    }

    @JRubyMethod
    public static IRubyObject hash(ThreadContext context, IRubyObject self) {
        RubyBasicObject selfObj = (RubyBasicObject) self;

        int h = selfObj.getType().hashCode();
        VariableAccessor[] accessors = getStructAccessors(selfObj);
        for (int i = 0; i < accessors.length; i++) {
            h = (h << 1) | (h < 0 ? 1 : 0);
            IRubyObject hash = context.safeRecurse(RubyData::hashData, selfObj, (IRubyObject) accessors[i].get(selfObj), "hash", true);
            h ^= toLong(context, hash);
        }

        return asFixnum(context, h);
    }

    @JRubyMethod(alias = "to_s")
    public static IRubyObject inspect(ThreadContext context, IRubyObject self) {
        return context.safeRecurse(
                RubyData::inspectData,
                newString(context, "#<data "), self, "hash", true);
    }

    @JRubyMethod
    public static RubyHash to_h(ThreadContext context, IRubyObject self, Block block) {
        RubyArray<RubySymbol> members = getStructMembers(self);
        VariableAccessor[] accessors = getStructAccessors(self);
        RubyHash h = newSmallHash(context);

        for (int i = 0 ; i < accessors.length; i++) {
            RubySymbol k = members.eltOk(i);
            IRubyObject v = (IRubyObject) accessors[i].get(self);
            if (block.isGiven()) {
                h.fastASetSmallPair(context, block.yieldSpecific(context, k, v));
            } else {
                h.fastASetSmall(k, v);
            }
        }
        return h;
    }

    @JRubyMethod
    public static RubyArray members(ThreadContext context, IRubyObject self) {
        return getStructMembers(self).aryDup();
    }

    @JRubyMethod
    public static RubyArray deconstruct(ThreadContext context, IRubyObject self) {
        VariableAccessor[] accessors = getStructAccessors(self);
        RubyArray ary = allocArray(context, accessors.length);
        for (int i = 0; i < accessors.length; i++) {
            ary.append(context, (IRubyObject) accessors[i].get(self));
        }
        return ary;
    }

    @JRubyMethod
    public static RubyHash deconstruct_keys(ThreadContext context, IRubyObject self, IRubyObject keys) {
        if (keys.isNil()) {
            return to_h(context, self, Block.NULL_BLOCK);
        }
        if (!(keys instanceof RubyArray)) {
            throw typeError(context, "wrong argument type " + keys.getType() + " (expected Array or nil)");
        }
        RubyArray keysAry = (RubyArray) keys;
        RubyArray<RubySymbol> members = getStructMembers(self);
        VariableAccessor[] accessors = getStructAccessors(self);
        if (accessors.length < keysAry.size()) {
            return newEmptyHash(context);
        }
        RubyHash h = newSmallHash(context);
        for (int i = 0; i < keysAry.size(); i++) {
            IRubyObject key = keysAry.eltOk(i);
            int memberIndex = members.indexOf(key);
            if (memberIndex == -1) {
                return h;
            }
            h.fastASetSmall(key, (IRubyObject) accessors[memberIndex].get(self));
        }
        return h;
    }

    @JRubyMethod(keywords = true, optional = 1, checkArity = false)
    public static IRubyObject with(ThreadContext context, IRubyObject self, IRubyObject[] args) {
        IRubyObject kwargs = IRRuntimeHelpers.receiveKeywords(context, args, false, true, false);
        if (!(kwargs instanceof RubyHash)) {
            checkArgumentCount(context, args.length, 0, 0);
            return self;
        }

        checkArgumentCount(context, args.length - 1, 0, 0);

        RubyHash kwargsHash = (RubyHash) kwargs;
        RubyHash h = to_h(context, self, Block.NULL_BLOCK);
        h.addAll(context, kwargsHash);
        context.callInfo = CALL_KEYWORD;
        return DataMethods.rbNew(context, self.getMetaClass(), h);
    }

    public static class DataMethods {
        @JRubyMethod(name = {"new", "[]"}, keywords = true, rest = true)
        public static IRubyObject rbNew(ThreadContext context, IRubyObject self, IRubyObject[] values) {
            RubyClass klass = (RubyClass) self;

            RubyHash init;
            IRubyObject maybeKwargs = IRRuntimeHelpers.receiveKeywords(context, values, true, true, false);
            if (maybeKwargs instanceof RubyHash) {
                init = (RubyHash) maybeKwargs;
            } else {
                RubyArray<RubySymbol> members = getMembersFromClass(klass);

                checkArgumentCount(context, values.length, 0, members.size());

                init = newSmallHash(context.runtime);

                for (int i = 0; i < values.length; i++) {
                    RubySymbol sym = members.eltOk(i);
                    init.fastASetSmall(sym, values[i]);
                }
            }

            IRubyObject dataObject = klass.getAllocator().allocate(context.runtime, klass);

            context.callInfo = ThreadContext.CALL_KEYWORD;

            // TODO: avoid initialize and hash overhead for known types
            dataObject.getMetaClass().getBaseCallSite(RubyClass.CS_IDX_INITIALIZE)
                    .call(context, self, dataObject, init);


            return dataObject;
        }

        @JRubyMethod(name = {"new", "[]"}, keywords = true)
        public static IRubyObject rbNew(ThreadContext context, IRubyObject self) {
            ThreadContext.resetCallInfo(context);

            RubyClass klass = (RubyClass) self;
            IRubyObject dataObject = klass.getAllocator().allocate(context.runtime, klass);

            dataObject.getMetaClass().getBaseCallSite(RubyClass.CS_IDX_INITIALIZE)
                    .call(context, self, dataObject);

            return dataObject;
        }

        @JRubyMethod(name = {"new", "[]"}, keywords = true)
        public static IRubyObject rbNew(ThreadContext context, IRubyObject self, IRubyObject hashOrElt) {
            int callInfo = ThreadContext.resetCallInfo(context);

            RubyClass klass = (RubyClass) self;
            RubyHash init;
            if (hasNonemptyKeywords(callInfo)) {
                if (!(hashOrElt instanceof RubyHash)) {
                    throw argumentError(context, 1, 0, 0);
                }

                init = (RubyHash) hashOrElt;
            } else {
                RubyArray<RubySymbol> members = getMembersFromClass(klass);

                checkArgumentCount(context, 1, 0, members.size());

                init = newSmallHash(context.runtime);

                RubySymbol sym = members.eltOk(0);
                init.fastASetSmall(sym, hashOrElt);
                callInfo = CALL_KEYWORD;
            }

            IRubyObject dataObject = klass.getAllocator().allocate(context.runtime, klass);

            context.callInfo = callInfo;

            // TODO: avoid initialize and hash overhead for known types
            dataObject.getMetaClass().getBaseCallSite(RubyClass.CS_IDX_INITIALIZE)
                    .call(context, self, dataObject, init);

            return dataObject;
        }

        @JRubyMethod(name = "members")
        public static IRubyObject members(ThreadContext context, IRubyObject self) {
            return getMembersFromClass((RubyClass) self).aryDup();
        }

        @JRubyMethod(name = "inspect")
        public static IRubyObject inspect(ThreadContext context, IRubyObject klass) {
            IRubyObject inspect = ((RubyClass) klass).rubyName(context);
            // TODO: keyword init like struct
//            if (RTEST(rb_struct_s_keyword_init(klass))) {
//                rb_str_cat_cstr(inspect, "(keyword_init: true)");
//            }
            return inspect;
        }
    }

    // TODO: Mostly copied from RubyStruct; unify.
    public static void marshalTo(ThreadContext context, RubyOutputStream out, IRubyObject data, MarshalDumper output) {
        output.registerObject(data);
        output.dumpDefaultObjectHeader(context, out, 'S', data.getMetaClass());

        RubyArray<RubySymbol> members = getStructMembers(data);
        VariableAccessor[] accessors = getStructAccessors(data);
        int size = members.size();
        output.writeInt(out, size);

        for (int i = 0; i < size; i++) {
            RubySymbol name = members.eltInternal(i);
            output.dumpObject(context, out, name);
            output.dumpObject(context, out, (IRubyObject) accessors[i].get(data));
        }
    }

    // TODO: Mostly copied from RubyStruct; unify.
    public static IRubyObject unmarshalFrom(ThreadContext context, RubyInputStream in, MarshalLoader input, RubyClass rbClass) {
        final RubyArray<RubySymbol> members = getMembersFromClass(rbClass);
        final VariableAccessor[] accessors = getAccessorsFromClass(rbClass);

        final int len = input.unmarshalInt(context, in);

        final IRubyObject result = input.entry(rbClass.allocate(context));

        for (int i = 0; i < len; i++) {
            RubySymbol slot = input.symbol(context, in);
            RubySymbol elem = members.eltInternal(i);
            if (!elem.equals(slot)) {
                throw typeError(context, str(context.runtime, "struct ", rbClass,
                        " not compatible (:", slot, " for :", elem, ")").toString());
            }
            accessors[i].set(result, input.unmarshalObject(context, in));
        }

        result.setFrozen(true);

        return result;
    }

    private static RubyClass newDataStruct(ThreadContext context, RubyClass superClass, LinkedHashSet<RubySymbol> keySet) {
        Ruby runtime = context.runtime;

        RubyClass subclass = RubyClass.newClass(runtime, superClass);

        VariableTableManager vtm = subclass.getVariableTableManager();
        VariableAccessor[] accessors = new VariableAccessor[keySet.size()];
        int i = 0;
        for (RubySymbol sym : keySet) {
            accessors[i++] = vtm.getVariableAccessorForWrite(sym.idString());
        }

        ObjectAllocator allocator =
                runtime.getObjectSpecializer()
                        .specializeForVariables(
                                subclass,
                                keySet.stream().map(RubySymbol::idString).collect(Collectors.toSet()));

        subclass.allocator(allocator);

        RubyArray members = newArray(context, keySet);
        members.freeze(context);
        subclass.setInternalVariable(MEMBERS_KEY, members);
        subclass.setInternalVariable(ACCESSORS_KEY, accessors);

        RubyClass dataSClass = subclass.getSingletonClass();

        dataSClass.undefMethods(context, "define")
                .defineMethods(context, DataMethods.class);

        for (RubySymbol sym : keySet) {
            VariableAccessor accessor = vtm.getVariableAccessorForWrite(sym.idString());
            // TODO: AttrReader expects to potentially see many variable tables; this could be simplified
            subclass.addMethod(context, sym.idString(), new AttrReaderMethod(subclass, Visibility.PUBLIC, accessor));
        }

        return subclass;
    }

    private static IRubyObject checkDataEquality(ThreadContext context, RubyBasicObject self, RubyBasicObject other, String name, ThreadContext.RecursiveFunctionEx<RubyBasicObject> func) {
        RubyBasicObject selfObj = self;
        RubyBasicObject otherObj = other;

        if (selfObj == otherObj) return context.tru;
        RubyClass metaClass = otherObj.getMetaClass();
        if (!metaClass.isKindOfModule(context.runtime.getData())) return context.fals;
        if (metaClass.getRealClass() != selfObj.getMetaClass().getRealClass()) return context.fals;
//        if (RSTRUCT_LEN(s) != RSTRUCT_LEN(s2)) {
//            rb_bug("inconsistent struct"); /* should never happen */
//        }

        return context.safeRecurse(func, selfObj, otherObj, name, true);
    }

    private static IRubyObject inspectData(ThreadContext context, RubyString prefix, IRubyObject s, boolean recur) {
        RubyString cname = s.getMetaClass().rubyName(context);
        RubyString str = prefix;
        char first = cname.charAt(0);

        if (recur || first != '#') {
            str.append(cname);
        }
        if (recur) {
            return str.catString(":...>");
        }

        RubyArray<RubySymbol> members = getStructMembers(s);
        VariableAccessor[] accessors = getStructAccessors(s);
        int len = members.size();

        for (int i = 0; i < len; i++) {
            if (i > 0) {
                str.catString(", ");
            } else if (first != '#') {
                str.catString(" ");
            }

            RubySymbol id = members.eltOk(i);
            if (id.validLocalVariableName() || id.validConstantName()) {
                str.append(id.fstring());
            } else {
                str.append(id.inspect(context));
            }
            str.catString("=");
            str.append(rbInspect(context, (IRubyObject) accessors[i].get(s)));
        }
        str.catString(">");

        return str;
    }

    private static RubyArray<RubySymbol> getStructMembers(IRubyObject s) {
        RubyClass metaClass = s.getMetaClass();
        return getMembersFromClass(metaClass);
    }

    private static RubyArray<RubySymbol> getMembersFromClass(RubyClass metaClass) {
        while (metaClass != null) {
            RubyArray<RubySymbol> members = (RubyArray<RubySymbol>) metaClass.getInternalVariable(MEMBERS_KEY);

            if (members != null) return members;

            metaClass = metaClass.getSuperClass();
        }

        throw new RuntimeException("non-Data attempted to access Data members");
    }

    private static VariableAccessor[] getStructAccessors(IRubyObject s) {
        RubyClass metaClass = s.getMetaClass();
        return getAccessorsFromClass(metaClass);
    }

    private static VariableAccessor[] getAccessorsFromClass(RubyClass metaClass) {
        while (metaClass != null) {
            VariableAccessor[] accessors = (VariableAccessor[]) metaClass.getInternalVariable(ACCESSORS_KEY);

            if (accessors != null) return accessors;

            metaClass = metaClass.getSuperClass();
        }

        throw new RuntimeException("non-Data attempted to access Data accessors");
    }

    private static VariableTableManager getVariableManagerFromClass(RubyClass metaClass) {
        while (metaClass != null) {
            VariableAccessor[] accessors = (VariableAccessor[]) metaClass.getInternalVariable(ACCESSORS_KEY);

            if (accessors != null) return metaClass.getVariableTableManager();

            metaClass = metaClass.getSuperClass();
        }

        throw new RuntimeException("non-Data attempted to access Data accessors");
    }

    private static IRubyObject hashData(ThreadContext ctx, RubyBasicObject state, IRubyObject obj, boolean recur) {
        return recur ? asFixnum(ctx, 0) : invokedynamic(ctx, obj, HASH);
    }

    private static IRubyObject eqlData(ThreadContext c, RubyBasicObject s, IRubyObject o, boolean recur) {
        if (recur) return c.tru;
        VariableAccessor[] accessors = getStructAccessors(s);
        for (int i = 0; i < accessors.length; i++) {
            VariableAccessor accessor = accessors[i];
            if (!RubyBasicObject.eqlInternal(c, ((RubyBasicObject) accessor.get(s)), (IRubyObject) accessor.get(o))) {
                return c.fals;
            }
        }

        return c.tru;
    }

    private static IRubyObject equalData(ThreadContext c, RubyBasicObject s, IRubyObject o, boolean recur) {
        if (recur) return c.tru;
        VariableAccessor[] accessors = getStructAccessors(s);
        for (int i = 0; i < accessors.length; i++) {
            VariableAccessor accessor = accessors[i];
            if (!RubyBasicObject.equalInternal(c, ((RubyBasicObject) accessor.get(s)), (IRubyObject) accessor.get(o))) {
                return c.fals;
            }
        }

        return c.tru;
    }
}