/*
 * Copyright 2011 SpringSource
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.codehaus.groovy.grails.cli.support;

import groovy.lang.GroovySystem;
import groovy.lang.MetaClass;
import java.lang.ref.WeakReference;
import java.util.ArrayList;
import java.util.List;
import groovy.lang.MetaClassRegistry;
import groovy.lang.MetaClassRegistryChangeEvent;
import groovy.lang.MetaClassRegistryChangeEventListener;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import org.codehaus.groovy.runtime.metaclass.MetaClassRegistryImpl;

/**
 * Allows clean-up of changes made to the MetaClassRegistry.
 *
 * @author Graeme Rocher
 * @since 2.0
 */
@SuppressWarnings("rawtypes")
public class MetaClassRegistryCleaner implements MetaClassRegistryChangeEventListener {

    private Map<Class, Object> alteredClasses = new ConcurrentHashMap<Class, Object> ();
    private Map<IdentityWeakReference, Object> alteredInstances = new ConcurrentHashMap<IdentityWeakReference, Object>();
    private static final Object NO_CUSTOM_METACLASS = new Object();

    public void updateConstantMetaClass(MetaClassRegistryChangeEvent cmcu) {
        MetaClass oldMetaClass = cmcu.getOldMetaClass();
        Class classToUpdate = cmcu.getClassToUpdate();
        Object instanceToUpdate = cmcu.getInstance();
        if (instanceToUpdate == null) {
            updateMetaClassOfClass(oldMetaClass, classToUpdate);
        } else {
            updateMetaClassOfInstance(oldMetaClass, instanceToUpdate);
        }
    }

    private void updateMetaClassOfInstance(MetaClass oldMetaClass, Object instanceToUpdate) {
        IdentityWeakReference key = new IdentityWeakReference(instanceToUpdate);
        if (oldMetaClass != null) {
            Object current = alteredInstances.get(key);
            if (current == null || current == NO_CUSTOM_METACLASS) {
                alteredInstances.put(key, oldMetaClass);
            }
        } else {
            alteredInstances.put(key, NO_CUSTOM_METACLASS);
        }
    }

    private void updateMetaClassOfClass(MetaClass oldMetaClass, Class classToUpdate) {
        if (oldMetaClass != null) {
            Object current = alteredClasses.get(classToUpdate);
            if(current == null || current == NO_CUSTOM_METACLASS) {
                alteredClasses.put(classToUpdate, oldMetaClass);
            }
        }
        else {
            alteredClasses.put(classToUpdate, NO_CUSTOM_METACLASS);
        }
    }

    public void clean() {
        MetaClassRegistryImpl registry = (MetaClassRegistryImpl) GroovySystem.getMetaClassRegistry();
        cleanMetaClassOfClass(registry);
        cleanMetaClassOfInstance(registry);
    }

    private void cleanMetaClassOfInstance(MetaClassRegistryImpl registry) {
        List<IdentityWeakReference> keys = new ArrayList<IdentityWeakReference>(alteredInstances.keySet());
        for (IdentityWeakReference key : keys) {
            Object instance = key.get();
            if (instance != null) {
                Object alteredMetaClass = alteredInstances.get(key);
                if (alteredMetaClass == NO_CUSTOM_METACLASS) {
                    alteredMetaClass = null;
                }
                registry.setMetaClass(instance, (MetaClass) alteredMetaClass);
            }
        }
        alteredInstances.clear();
    }

    private void cleanMetaClassOfClass(MetaClassRegistryImpl registry) {
        Set<Class> classes = new HashSet<Class>(alteredClasses.keySet());
        for (Class aClass : classes) {
            Object alteredMetaClass = alteredClasses.get(aClass);
            if(alteredMetaClass == NO_CUSTOM_METACLASS) {
                registry.removeMetaClass(aClass);
            }
            else {
                registry.setMetaClass(aClass, (MetaClass) alteredMetaClass);
            }
        }
        alteredClasses.clear();
    }

    private static final class IdentityWeakReference extends WeakReference<Object> {

        private int hash;

        public IdentityWeakReference(Object referent) {
            super(referent);
            hash = System.identityHashCode(referent);
        }

        public int hashCode() {
            return hash;
        }

        public boolean equals(Object obj) {
            if (this == obj) {
                return true;
            }
            IdentityWeakReference ref = (IdentityWeakReference) obj;
            return this.get() == ref.get();
        }

    }

}

@@ -1628,16 +1628,53 @@
 bject();
+%0A    private static boolean cleaning;
 %0A%0A    pu
@@ -1740,24 +1740,52 @@
 ent cmcu) %7B%0A
+        if(!cleaning) %7B%0A    
         Meta
@@ -1837,16 +1837,20 @@
         
+    
 Class cl
@@ -1888,32 +1888,36 @@
 date();%0A        
+ 
+   
 Object instanceT
@@ -1946,32 +1946,36 @@
 ance();%0A        
+    
 if (instanceToUp
@@ -1994,32 +1994,36 @@
 ) %7B%0A            
+  
+  
 updateMetaClassO
@@ -2059,32 +2059,36 @@
 pdate);%0A        
+    
 %7D else %7B%0A       
@@ -2084,32 +2084,36 @@
 e %7B%0A            
+  
+  
 updateMetaClassO
@@ -2147,32 +2147,46 @@
 tanceToUpdate);%0A
+            %7D%0A
         %7D%0A    %7D%0A
@@ -3153,24 +3153,71 @@
 d clean() %7B%0A
+        try %7B%0A            cleaning = true;%0A    
         Meta
@@ -3306,32 +3306,36 @@
 stry();%0A        
+    
 cleanMetaClassOf
@@ -3351,32 +3351,36 @@
 istry);%0A        
+  
+  
 cleanMetaClassOf
@@ -3395,24 +3395,84 @@
 (registry);%0A
+        %7D finally %7B%0A            cleaning = false;%0A        %7D%0A
     %7D%0A%0A    p

fix for GRAILS-8690 "Unit test side-effects in interactive console"