package org.bdware.analysis.taint;

import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.MethodInsnNode;

import org.objectweb.asm.tree.InvokeDynamicInsnNode;

public class HeapObject extends TaintValue {
	private Map<String, TaintValue> fields = new HashMap<String, TaintValue>();
	static int allocatID = 0;
	int id;

	public HeapObject() {
		super(1);
		id = allocatID++;
	}

	public HeapObject(TaintValue t) {
		this();
		if (t != null) {
			merge(t);
			size = t.size;
		} else {
			isTainted = 0;
			size = 1;
		}
	}

	public static HeapObject getRootObject() {
		HeapObject ret = new HeapObject();
		ret.setProp("Global", new HeapObject());
		ret.setProp("getScope", new HeapObject());
		return ret;
	}

	@Override
	public String toString() {
		return countSize() + "|" + super.toString();
	}

	private String countSize() {
		objSet = new HashSet<>();
		return countSizeInternal();
	}

	private String countSizeInternal() {
		if (objSet.contains(this))
			return "{H" + id + "}";
		objSet.add(this);
		String ret = "{H" + id + ",";
		for (String field : fields.keySet()) {
			TaintValue tv = fields.get(field);
			if (tv instanceof HeapObject)
				ret += field + ":" + ((HeapObject) tv).countSizeInternal() + ",";
			else
				ret += field + ":" + tv.isTainted;
		}
		ret += isTainted + "}";
		return ret;
	}

	private static HashSet<HeapObject> objSet;

	public static TaintValue operate(AbstractInsnNode insn, List<? extends TaintValue> values,
			TaintValue calculatedVal) {
		TaintValue retVal = null;
		String desc;
		switch (insn.getOpcode()) {
		case Opcodes.INVOKEVIRTUAL:
			desc = ((MethodInsnNode) insn).owner;
			if (isGetScopeObject((MethodInsnNode) insn) && values.size() != 0) {
				TaintValue tv = values.get(0);
				if (tv instanceof HeapObject) {
					HeapObject heapObject = (HeapObject) tv;
					retVal = heapObject.getProp("getScope");
					if (retVal == null) {
						return calculatedVal;
					} else
						return retVal;
				}
			}
			break;
		case Opcodes.INVOKESPECIAL:
		case Opcodes.INVOKESTATIC:
		case Opcodes.INVOKEINTERFACE:
			break;
		case Opcodes.INVOKEDYNAMIC:
			desc = ((InvokeDynamicInsnNode) insn).name;
			desc = desc.replaceAll(".*:", "");
			if (isDynGet((InvokeDynamicInsnNode) insn)) {
				TaintValue tval = values.get(0);
				if (tval instanceof HeapObject) {
					HeapObject heapObject = (HeapObject) tval;
					retVal = heapObject.getProp(desc);
					if (retVal != null) {
						if (TaintConfig.isDebug)
							System.out.println("[HeapObject] get:" + desc + " taintVal:" + retVal);
						return retVal;
					} else {
						if (TaintConfig.isDebug)
							System.out.println("[HeapObject] get:" + desc + " taintVal: empty");
						heapObject.setProp(desc, new HeapObject());
						return heapObject.getProp(desc);
					}
				} else {
					// TODO
				}
			} else if (isDynSet((InvokeDynamicInsnNode) insn)) {
				// TaintValue tv = values.get(0);
				if (TaintConfig.isDebug)
					System.out.println("[HeapObject] set:" + desc + " taintVal:" + values.get(1));
				TaintValue tval = values.get(0);
				if (tval instanceof HeapObject) {
					HeapObject heapObject = (HeapObject) tval;
					// heapObject.merge(values.get(1));
					heapObject.setProp(desc,
							values.get(1) instanceof HeapObject ? values.get(1) : new HeapObject(values.get(1)));
				} else {
					HeapObject heapObject = new HeapObject(tval);
					heapObject.merge(tval);
					// heapObject.merge(values.get(1));
					heapObject.setProp(desc, values.get(1));
				}
			} else {
				for (TaintValue tv : values) {
					if (tv instanceof HeapObject) {
						calculatedVal.isTainted |= ((HeapObject) tv).wholeTaint();
					}
				}
			}
			break;
		default:
		}
		return calculatedVal;

	}

	private static boolean isDynGet(InvokeDynamicInsnNode insn) {
		String desc = insn.name;
		if (desc.contains("getProp") && desc.contains("getMethod") && desc.contains("getElem")) {
			return true;
		}
		return false;
	}

	private static boolean isDynSet(InvokeDynamicInsnNode insn) {
		String desc = insn.name;
		if (desc.contains("setProp") && desc.contains("setElem")) {
			// System.out.println(desc);
			return true;
		}
		return false;
	}

	private static boolean isGetScopeObject(MethodInsnNode insn) {
		if (insn.owner.contains("wrp/jdk/nashorn/internal/runtime/ScriptFunction") && insn.name.equals("getScope"))
			return true;
		return false;
	}

	public TaintValue getProp(String str) {
		return fields.get(str);
	}

	public TaintValue setProp(String str, TaintValue val) {
		return fields.put(str, val);
	}

	@Override
	public HeapObject clone() {
		return this;
	}

	public HeapObject actualClone() {
		ccObj = new HashMap<>();
		return cloneInternal();
	}

	static Map<HeapObject, HeapObject> ccObj;

	private HeapObject cloneInternal() {
		if (ccObj.containsKey(this))
			return ccObj.get(this);
		HeapObject ho = new HeapObject();
		ccObj.put(this, ho);
		ho.isTainted = isTainted;
		for (String field : fields.keySet()) {
			TaintValue target = fields.get(field);
			if (target instanceof HeapObject)
				ho.fields.put(field, ((HeapObject) target).cloneInternal());
			else
				ho.fields.put(field, target.clone());
		}
		return ho;
	}

	// TODO consider the "Cite circle"
	public void heapMerge(TaintValue target) {
		merge(target);
		if (target instanceof HeapObject) {
			HeapObject targetObject = (HeapObject) target;
			for (String field : targetObject.fields.keySet()) {
				if (fields.containsKey(field)) {
					TaintValue t2 = fields.get(field);
					TaintValue target2 = targetObject.fields.get(field);
					// System.out.println("[HeapObject heapMerge]:" + field);
					if (t2 instanceof HeapObject) {
						((HeapObject) t2).heapMerge(target2);
					} else if (target2 instanceof HeapObject) {
						HeapObject next = ((HeapObject) target2).clone();
						next.merge(t2);
						fields.put(field, next);
					} else
						t2.merge(target2);

				} else
					fields.put(field, targetObject.fields.get(field).clone());

			}
		}
	}

	public synchronized long wholeTaint() {
		objSet = new HashSet<>();
		return wholeTaintInternal();
	}

	private long wholeTaintInternal() {
		if (objSet.contains(this))
			return 0L;
		objSet.add(this);
		long ret = isTainted;
		for (TaintValue tv : fields.values()) {
			if (tv instanceof HeapObject)
				ret |= ((HeapObject) tv).wholeTaintInternal();
			else
				ret |= tv.isTainted;
		}
		return ret;
	}

}
