package com.yammer.metrics.spring.config;

import org.springframework.aop.framework.ProxyConfig;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.parsing.BeanComponentDefinition;
import org.springframework.beans.factory.parsing.CompositeComponentDefinition;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.xml.BeanDefinitionParser;
import org.springframework.beans.factory.xml.ParserContext;
import org.springframework.util.StringUtils;
import org.w3c.dom.Element;

import com.yammer.metrics.core.HealthCheckRegistry;
import com.yammer.metrics.core.MetricsRegistry;
import com.yammer.metrics.spring.ExceptionMeteredAnnotationBeanPostProcessor;
import com.yammer.metrics.spring.GaugeAnnotationBeanPostProcessor;
import com.yammer.metrics.spring.HealthCheckBeanPostProcessor;
import com.yammer.metrics.spring.MeteredAnnotationBeanPostProcessor;
import com.yammer.metrics.spring.TimedAnnotationBeanPostProcessor;

import static org.springframework.beans.factory.config.BeanDefinition.*;

public class AnnotationDrivenBeanDefinitionParser implements BeanDefinitionParser {

    public BeanDefinition parse(Element element, ParserContext parserContext) {
        final Object source = parserContext.extractSource(element);

        final CompositeComponentDefinition compDefinition = new CompositeComponentDefinition(element.getTagName(),
                                                                                             source);
        parserContext.pushContainingComponent(compDefinition);

        String metricsBeanName = element.getAttribute("metrics-registry");
        if (!StringUtils.hasText(metricsBeanName)) {
            metricsBeanName = registerComponent(parserContext,
                                                build(MetricsRegistry.class,
                                                      source,
                                                      ROLE_APPLICATION));
        }

        String healthCheckBeanName = element.getAttribute("health-check-registry");
        if (!StringUtils.hasText(healthCheckBeanName)) {
            healthCheckBeanName = registerComponent(parserContext,
                                                    build(HealthCheckRegistry.class,
                                                          source,
                                                          ROLE_APPLICATION));
        }

        String scope = element.getAttribute("scope");
        if (!StringUtils.hasText(scope)) {
            scope = null;
        }

        ProxyConfig proxyConfig = new ProxyConfig();

        if (StringUtils.hasText(element.getAttribute("expose-proxy"))) {
            proxyConfig.setExposeProxy(Boolean.valueOf(element.getAttribute("expose-proxy")));
        }

        if (StringUtils.hasText(element.getAttribute("proxy-target-class"))) {
            proxyConfig.setProxyTargetClass(Boolean.valueOf(element.getAttribute("proxy-target-class")));
        }

        registerComponent(parserContext,
                          build(ExceptionMeteredAnnotationBeanPostProcessor.class,
                                source,
                                ROLE_INFRASTRUCTURE
                          )
                          .addConstructorArgReference(metricsBeanName)
                          .addConstructorArgValue(proxyConfig)
                          .addConstructorArgValue(scope));

        registerComponent(parserContext,
                          build(MeteredAnnotationBeanPostProcessor.class,
                                source,
                                ROLE_INFRASTRUCTURE
                          )
                          .addConstructorArgReference(metricsBeanName)
                          .addConstructorArgValue(proxyConfig)
                          .addConstructorArgValue(scope));

        registerComponent(parserContext,
                          build(TimedAnnotationBeanPostProcessor.class,
                                source,
                                ROLE_INFRASTRUCTURE
                          )
                          .addConstructorArgReference(metricsBeanName)
                          .addConstructorArgValue(proxyConfig)
                          .addConstructorArgValue(scope));

        registerComponent(parserContext,
                          build(GaugeAnnotationBeanPostProcessor.class,
                                source,
                                ROLE_INFRASTRUCTURE
                          )
                          .addConstructorArgReference(metricsBeanName)
                          .addConstructorArgValue(scope));

        registerComponent(parserContext,
                          build(HealthCheckBeanPostProcessor.class,
                                source,
                                ROLE_INFRASTRUCTURE
                          )
                          .addConstructorArgReference(healthCheckBeanName));

        parserContext.popAndRegisterContainingComponent();

        return null;
    }

    private BeanDefinitionBuilder build(Class<?> klazz, Object source, int role) {
        final BeanDefinitionBuilder beanDefBuilder = BeanDefinitionBuilder.rootBeanDefinition(klazz);
        beanDefBuilder.setRole(role);
        beanDefBuilder.getRawBeanDefinition().setSource(source);
        return beanDefBuilder;
    }

    private String registerComponent(ParserContext parserContext, BeanDefinitionBuilder beanDefBuilder) {
    	final BeanDefinition beanDef = beanDefBuilder.getBeanDefinition();
        final String beanName = parserContext.getReaderContext().registerWithGeneratedName(beanDef);
        parserContext.registerComponent(new BeanComponentDefinition(beanDef, beanName));
        return beanName;
    }

}

@@ -580,16 +580,91 @@
 ement;%0A%0A
+import com.yammer.metrics.HealthChecks;%0Aimport com.yammer.metrics.Metrics;%0A
 import c
@@ -1941,24 +1941,16 @@
 (Metrics
-Registry
 .class,%0A
@@ -2074,32 +2074,117 @@
 OLE_APPLICATION)
+%0A                                                .setFactoryMethod(%22defaultRegistry%22)
 );%0A        %7D%0A%0A  
@@ -2458,24 +2458,17 @@
 lthCheck
-Registry
+s
 .class,%0A
@@ -2608,16 +2608,105 @@
 ICATION)
+%0A                                                    .setFactoryMethod(%22defaultRegistry%22)
 );%0A     

Use default registry if not explicitly set in metrics:annotation-driven.