bridge119

package com.citi.get.cet.comet.tools;

import org.springframework.beans.PropertyValue;
import org.springframework.beans.factory.config.*;
import org.springframework.context.*;
import org.springframework.context.SmartLifecycle;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.*;

import static org.apache.commons.lang.StringUtils.trim;

public class SpringVisualizer implements ApplicationContextAware, SmartLifecycle {

    private ConfigurableApplicationContext startContext;
    private boolean isRunning = false;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) {
        this.startContext = (ConfigurableApplicationContext) applicationContext;
    }

    @Override
    public void start() {
        StringBuilder dot = new StringBuilder("digraph G {\n");

        // ===== LAYOUT (VERTICAL + COMPACT) =====
        dot.append("    layout=dot;\n");
        dot.append("    rankdir=TB;\n");           // vertical flow
        dot.append("    splines=ortho;\n");        // cleaner edges
        dot.append("    nodesep=0.3;\n");          // tighter horizontally
        dot.append("    ranksep=1.2;\n");          // more vertical spacing
        dot.append("    ordering=out;\n");

        dot.append("    graph [fontname=\"Arial\", fontsize=10];\n");
        dot.append("    node [shape=plaintext, fontname=\"Arial\", fontsize=10];\n");
        dot.append("    edge [dir=forward];\n\n");

        Set<String> processed = new HashSet<>();
        Map<String, List<String>> clusters = new HashMap<>();
        List<String> isolatedBeans = new ArrayList<>();

        ApplicationContext current = this.startContext;

        // ===== FIRST PASS: COLLECT =====
        while (current != null) {
            ConfigurableListableBeanFactory factory =
                    ((ConfigurableApplicationContext) current).getBeanFactory();

            for (String name : factory.getBeanDefinitionNames()) {
                if (!processed.add(name)) continue;

                // cluster by package
                String pkg = name.contains(".")
                        ? name.substring(0, name.lastIndexOf('.'))
                        : "default";

                clusters.computeIfAbsent(pkg, k -> new ArrayList<>()).add(name);

                // detect isolated
                if (factory.getDependenciesForBean(name).length == 0 &&
                    factory.getDependentBeans(name).length == 0) {
                    isolatedBeans.add(name);
                }
            }
            current = current.getParent();
        }

        // ===== SECOND PASS: RENDER NODES =====
        current = this.startContext;
        processed.clear();

        while (current != null) {
            ConfigurableListableBeanFactory factory =
                    ((ConfigurableApplicationContext) current).getBeanFactory();

            for (String name : factory.getBeanDefinitionNames()) {
                if (!processed.add(name)) continue;

                BeanDefinition def = factory.getBeanDefinition(name);

                String color = name.toLowerCase().contains("comet") ? "#FFF904" :
                               (name.toLowerCase().contains("citi") ? "#CBE6C7" : "#E1F5FE");

                dot.append(String.format("    \"%s\" [label=<", name));
                dot.append("<TABLE BORDER=\"0\" CELLBORDER=\"1\" CELLSPACING=\"0\" CELLPADDING=\"4\" BGCOLOR=\"")
                        .append(color).append("\">");

                // header
                dot.append("<TR><TD COLSPAN=\"2\" BGCOLOR=\"#DDDDDD\" CELLPADDING=\"6\"><B>")
                        .append(xmlEscape(name))
                        .append("</B></TD></TR>");

                // properties
                for (PropertyValue pv : def.getPropertyValues().getPropertyValues()) {
                    dot.append("<TR>");
                    dot.append("<TD WIDTH=\"120\" ALIGN=\"LEFT\"><FONT FACE=\"Courier\">")
                            .append(xmlEscape(pv.getName()))
                            .append("</FONT></TD>");
                    dot.append("<TD WIDTH=\"180\" ALIGN=\"LEFT\">")
                            .append(extractValue(pv.getValue()))
                            .append("</TD>");
                    dot.append("</TR>");
                }

                // constructor args
                ConstructorArgumentValues cav = def.getConstructorArgumentValues();

                for (Map.Entry<Integer, ConstructorArgumentValues.ValueHolder> e :
                        cav.getIndexedArgumentValues().entrySet()) {

                    dot.append("<TR>");
                    dot.append("<TD><I>arg[").append(e.getKey()).append("]</I></TD>");
                    dot.append("<TD>").append(extractValue(e.getValue().getValue())).append("</TD>");
                    dot.append("</TR>");
                }

                for (ConstructorArgumentValues.ValueHolder vh : cav.getGenericArgumentValues()) {
                    String type = vh.getType() != null
                            ? vh.getType().substring(vh.getType().lastIndexOf('.') + 1)
                            : "gen";

                    dot.append("<TR>");
                    dot.append("<TD><I>arg:").append(type).append("</I></TD>");
                    dot.append("<TD>").append(extractValue(vh.getValue())).append("</TD>");
                    dot.append("</TR>");
                }

                dot.append("</TABLE>>];\n");

                // edges
                for (String dep : factory.getDependenciesForBean(name)) {
                    dot.append(String.format("    \"%s\" -> \"%s\";\n", name, dep));
                }
            }
            current = current.getParent();
        }

        // ===== CLUSTERS =====
        int clusterId = 0;
        for (Map.Entry<String, List<String>> entry : clusters.entrySet()) {
            dot.append("subgraph cluster_").append(clusterId++).append(" {\n");
            dot.append("label=\"").append(xmlEscape(entry.getKey())).append("\";\n");
            dot.append("style=filled;\ncolor=\"#F5F5F5\";\n");

            for (String bean : entry.getValue()) {
                dot.append("\"").append(bean).append("\";\n");
            }
            dot.append("}\n");
        }

        // ===== STACK ISOLATED =====
        dot.append("    { rank=same;\n");
        for (String bean : isolatedBeans) {
            dot.append("        \"").append(bean).append("\";\n");
        }
        dot.append("    }\n");

        dot.append("}\n");

        writeToFile(dot.toString());
        this.isRunning = true;
    }

    private String extractValue(Object value) {
        if (value == null) return "null";

        if (value instanceof TypedStringValue) {
            return xmlEscape(trim(((TypedStringValue) value).getValue()));
        }

        if (value instanceof BeanReference) {
            return "@" + ((BeanReference) value).getBeanName();
        }

        if (value instanceof BeanDefinition) {
            String cls = ((BeanDefinition) value).getBeanClassName();
            return cls != null
                    ? "<i>" + cls.substring(cls.lastIndexOf('.') + 1) + "</i>"
                    : "<i>InnerBean</i>";
        }

        if (value instanceof Iterable) {
            List<String> out = new ArrayList<>();
            int i = 0;
            for (Object o : (Iterable<?>) value) {
                if (i++ > 10) { out.add("..."); break; }
                out.add(extractValue(o));
            }
            return String.join("<BR/>", out);
        }

        String s = value.toString();
        if (s.length() > 30) {
            s = s.substring(0, 27) + "...";
        }
        return xmlEscape(s);
    }

    private void writeToFile(String content) {
        try {
            File root = findProjectRoot();
            File buildDir = new File(root, "build");
            buildDir.mkdirs();

            File file = new File(buildDir, "spring-beans.dot");
            try (FileWriter writer = new FileWriter(file)) {
                writer.write(content);
            }

            System.out.println("\n==== GRAPH GENERATED ====");
            System.out.println(file.toURI());
            System.out.println("========================\n");

        } catch (IOException e) {
            System.err.println("Failed to write DOT file: " + e.getMessage());
        }
    }

    private File findProjectRoot() {
        File current = new File(System.getProperty("user.dir"));
        while (current != null) {
            if (new File(current, "build.gradle").exists() ||
                new File(current, "settings.gradle").exists()) {
                return current;
            }
            current = current.getParentFile();
        }
        return new File(System.getProperty("user.dir"));
    }

    private String xmlEscape(String input) {
        if (input == null) return "";
        return input.replace("&", "&amp;")
                .replace("<", "&lt;")
                .replace(">", "&gt;");
    }

    @Override public int getPhase() { return Integer.MAX_VALUE; }
    @Override public boolean isAutoStartup() { return true; }
    @Override public void stop() { isRunning = false; }
    @Override public boolean isRunning() { return isRunning; }
    @Override public void stop(Runnable callback) { stop(); callback.run(); }
}