bridge22

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;

    // ===== TUNING =====
    private static final int MAX_ROWS = 5;
    private static final int MAX_STRING = 40;
    private static final int MAX_EDGES = 1200;

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

    @Override
    public void start() {

        StringBuilder dot = new StringBuilder("digraph G {\n");

        dot.append("rankdir=TB;\n");
        dot.append("nodesep=0.4;\n");
        dot.append("ranksep=1.2;\n");
        dot.append("node [shape=box, fontname=\"Arial\", fontsize=10];\n");
        dot.append("edge [dir=forward];\n\n");

        Map<String, Set<String>> graph = new HashMap<>();
        Map<String, BeanDefinition> defs = new HashMap<>();

        // ===== BUILD GRAPH =====
        ApplicationContext current = startContext;

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

            for (String name : factory.getBeanDefinitionNames()) {

                defs.put(name, factory.getBeanDefinition(name));
                graph.putIfAbsent(name, new HashSet<>());

                for (String dep : factory.getDependenciesForBean(name)) {
                    graph.putIfAbsent(dep, new HashSet<>());

                    // UNDIRECTED LINK
                    graph.get(name).add(dep);
                    graph.get(dep).add(name);
                }
            }
            current = current.getParent();
        }

        // ===== CONNECTED COMPONENTS =====
        List<List<String>> components = findComponents(graph);

        int edgeCount = 0;

        // ===== RENDER =====
        for (List<String> comp : components) {

            // vertical stack inside component
            dot.append("{\n");

            for (String bean : comp) {
                renderNode(dot, bean, defs.get(bean));
            }

            // force vertical chain
            for (int i = 0; i < comp.size() - 1; i++) {
                dot.append("\"").append(comp.get(i)).append("\" -> \"")
                   .append(comp.get(i + 1))
                   .append("\" [style=invis];\n");
            }

            dot.append("}\n");

            // real edges
            for (String from : comp) {
                for (String to : graph.getOrDefault(from, Collections.emptySet())) {
                    if (edgeCount++ > MAX_EDGES) break;
                    dot.append("\"").append(from).append("\" -> \"")
                       .append(to).append("\";\n");
                }
            }
        }

        dot.append("}\n");

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

    // ===== CONNECTED COMPONENTS (DFS) =====
    private List<List<String>> findComponents(Map<String, Set<String>> graph) {
        List<List<String>> result = new ArrayList<>();
        Set<String> visited = new HashSet<>();

        for (String node : graph.keySet()) {
            if (visited.contains(node)) continue;

            List<String> comp = new ArrayList<>();
            Deque<String> stack = new ArrayDeque<>();
            stack.push(node);

            while (!stack.isEmpty()) {
                String cur = stack.pop();
                if (!visited.add(cur)) continue;

                comp.add(cur);

                for (String next : graph.getOrDefault(cur, Collections.emptySet())) {
                    if (!visited.contains(next)) {
                        stack.push(next);
                    }
                }
            }

            // sort for stable vertical layout
            comp.sort(String::compareTo);
            result.add(comp);
        }

        // sort components: large ones first
        result.sort((a, b) -> Integer.compare(b.size(), a.size()));

        return result;
    }

    // ===== NODE =====
    private void renderNode(StringBuilder dot, String name, BeanDefinition def) {

        StringBuilder label = new StringBuilder();
        label.append(shorten(name));

        if (def != null) {
            int rowCount = 0;

            for (PropertyValue pv : def.getPropertyValues().getPropertyValues()) {
                if (rowCount++ >= MAX_ROWS) {
                    label.append("\\n...");
                    break;
                }

                label.append("\\n")
                     .append(shorten(pv.getName()))
                     .append("=")
                     .append(shorten(extractValue(pv.getValue())));
            }
        }

        dot.append("\"").append(name).append("\" ")
           .append("[label=\"").append(escape(label.toString())).append("\", ")
           .append("style=filled, fillcolor=\"").append(getColor(name)).append("\"];\n");
    }

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

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

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

        return value.toString();
    }

    private String shorten(String s) {
        if (s == null) return "";
        s = trim(s);
        if (s.length() > MAX_STRING) {
            return s.substring(0, MAX_STRING - 3) + "...";
        }
        return s;
    }

    private String escape(String s) {
        return s.replace("\"", "'").replace("\n", "\\n");
    }

    private String getColor(String name) {
        String lower = name.toLowerCase();
        if (lower.contains("comet")) return "#FFF904";
        if (lower.contains("citi")) return "#CBE6C7";
        return "#E1F5FE";
    }

    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"));
    }

    @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(); }
}