package com.example.statemachinedemo;

import java.util.*;

public class ExportUtils {

    public static String generatePlantUML(List<Transition> transitions,
                                                                Set<String> startStates,
                                                                Set<String> endStates,
                                                                boolean useLambdaGuards) {
    StringBuilder sb = new StringBuilder();
    sb.append("@startuml\n");
    sb.append("skinparam state {\n");
    sb.append("  BackgroundColor<<choice>> LightYellow\n");
    sb.append("}\n\n");

    // 1. Print start states
    for (String start : startStates) {
        sb.append("[*] --> ").append(simplify(start)).append("\n");
    }
    sb.append("\n");

    // 2. Find all states that have choice transitions originating from them
    Set<String> statesWithChoice = new HashSet<>();
    for (Transition t : transitions) {
        if ("withChoice".equals(t.type)) {
            for (String source : t.sourceStates) {
                statesWithChoice.add(simplify(source));
            }
        }
    }

    // 3. Declare choice pseudostates & connect from original state
    for (String state : statesWithChoice) {
        String choiceState = state + "_choice";
        sb.append("state ").append(choiceState).append(" <<choice>>\n");
        sb.append(state).append(" --> ").append(choiceState).append("\n");
    }
    sb.append("\n");

    // Palette of colors for choice states (NO leading '#')
    List<String> choiceColors = List.of("FF6347", "4682B4", "32CD32", "FFD700", "6A5ACD", "FF69B4");

    // Assign a color per withChoice source state cycling through palette
    Map<String, String> choiceStateColorMap = new HashMap<>();
    int colorIndex = 0;
    for (String state : statesWithChoice) {
        choiceStateColorMap.put(state, choiceColors.get(colorIndex % choiceColors.size()));
        colorIndex++;
    }

    Set<String> transitionLines = new HashSet<>();

    // Helper to get color for non-choice transitions (NO leading '#')
        java.util.function.Function<String, String> getTransitionColor = (type) -> switch (type) {
            case "withExternal" -> "1E90FF"; // DodgerBlue
            case "withInternal" -> "32CD32"; // LimeGreen
            case "withLocal" -> "FFA500";    // Orange
            case "withJunction" -> "FF69B4"; // HotPink
            case "withJoin" -> "8A2BE2";     // BlueViolet
            case "withFork" -> "20B2AA";     // LightSeaGreen
            default -> "000000";             // Black fallback
        };
    // 4. Output transitions
    for (Transition t : transitions) {
        if (t.sourceStates == null || t.sourceStates.isEmpty()) continue;

        boolean allowNoTarget = "withJoin".equals(t.type) || "withFork".equals(t.type);

        List<String> targets;
        if (t.targetStates == null || t.targetStates.isEmpty()) {
            if (!allowNoTarget) {
                // Use sourceStates as targets (self-loop)
                targets = t.sourceStates;
            } else {
                // withJoin or withFork with no targets => skip emitting transitions
                continue;
            }
        } else {
            targets = t.targetStates;
        }

        for (String rawSource : t.sourceStates) {
            String source = simplify(rawSource);

            for (String rawTarget : targets) {
                String target = simplify(rawTarget);

                String transitionSource = source;

                // Reroute choice transitions through pseudostate
                if ("withChoice".equals(t.type) && statesWithChoice.contains(source)) {
                    transitionSource = source + "_choice";
                }

                // Prepare guard text
                String guardText = "";
                if (t.guard != null && !t.guard.isBlank()) {
                    if (useLambdaGuards && t.isLambdaGuard) {
                        guardText = "lambda";
                    } else {
                        guardText = t.guard.replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim();
                    }
                }

                // Prepare actions text (comma separated)
                String actionsText = "";
                if (t.actions != null && !t.actions.isEmpty()) {
                    StringBuilder actionBuilder = new StringBuilder();
                    for (int i = 0; i < t.actions.size(); i++) {
                        if (i > 0) actionBuilder.append(", ");
                        String action = t.actions.get(i);
                        boolean isLambda = t.isLambdaActions.get(i);
                        if (useLambdaGuards && isLambda) {
                            actionBuilder.append("lambda");
                        } else {
                            actionBuilder.append(action);
                        }
                    }
                    actionsText = actionBuilder.toString();
                }

                // Build label with event, guard, and actions
                StringBuilder label = new StringBuilder();
                if (t.event != null && !t.event.isBlank()) {
                    label.append(simplify(t.event));
                }
                if (!guardText.isBlank()) {
                    if (label.length() > 0) label.append(" ");
                    label.append("[").append(guardText).append("]");
                }
                if (!actionsText.isBlank()) {
                    if (label.length() > 0) label.append(" / ");
                    label.append(actionsText);
                }

                // Add order info for choice/junction transitions
                if (("withChoice".equals(t.type) || "withJunction".equals(t.type)) && t.order != null) {
                    if (label.length() > 0) label.append(" ");
                    label.append("(order=").append(t.order).append(")");
                }

                // Determine color for this transition
                String color;
                if ("withChoice".equals(t.type)) {
                    color = choiceStateColorMap.getOrDefault(source, "000000");
                } else {
                    color = getTransitionColor.apply(t.type);
                }

                // Add color with single '#' prefix
                String line = transitionSource + " -[#"+color+",bold]-> " + target + (label.isEmpty() ? "" : (" : " + label));

                // Skip duplicates
                if (transitionLines.add(line)) {
                    sb.append(line).append("\n");
                }
            }
        }
    }

    sb.append("\n");

    // 5. Mark end states with transition to [*]
    for (String end : endStates) {
        sb.append(simplify(end)).append(" --> [*]\n");
    }

    sb.append("@enduml\n");
    return sb.toString();
}


    public static String generateDot(List<Transition> transitions,
                                     Set<String> startStates,
                                     Set<String> endStates,
                                     boolean useLambdaGuards) {
        StringBuilder sb = new StringBuilder();
        sb.append("digraph stateMachine {\n");
        sb.append("  rankdir=LR;\n");
        sb.append("  node [shape=circle, style=filled, fillcolor=lightgray];\n\n");

        // 1. Fake start node
        sb.append("  start [shape=point, label=\"\"];\n");

        // 2. Find states with choice transitions
        Set<String> statesWithChoice = new HashSet<>();
        for (Transition t : transitions) {
            if ("withChoice".equals(t.type)) {
                for (String source : t.sourceStates) {
                    statesWithChoice.add(simplify(source));
                }
            }
        }

        // 3. Declare normal states and choice pseudostates
        Set<String> allStates = new HashSet<>();
        for (Transition t : transitions) {
            if (t.sourceStates != null) {
                for (String s : t.sourceStates) allStates.add(simplify(s));
            }
            if (t.targetStates != null) {
                for (String tgt : t.targetStates) allStates.add(simplify(tgt));
            }
        }

        // Print all normal states (except choice pseudostates for now)
        for (String state : allStates) {
            if (!statesWithChoice.contains(state + "_choice")) {
                String shape = endStates.contains(state) ? "doublecircle" : "circle";
                sb.append("  ").append(state)
                        .append(" [shape=").append(shape).append(", fillcolor=lightgray];\n");
            }
        }
        sb.append("\n");

        // Declare choice pseudostates as diamonds
        for (String state : statesWithChoice) {
            String choiceNode = state + "_choice";
            sb.append("  ").append(choiceNode)
                    .append(" [shape=diamond, label=\"\", fillcolor=lightyellow];\n");
        }
        sb.append("\n");

        // 4. Connect start node(s) to start states
        for (String start : startStates) {
            sb.append("  start -> ").append(simplify(start)).append(";\n");
        }
        sb.append("\n");

        // 5. Connect normal states to their choice pseudostates
        for (String state : statesWithChoice) {
            String choiceNode = state + "_choice";
            sb.append("  ").append(state).append(" -> ").append(choiceNode).append(";\n");
        }
        sb.append("\n");

        // 6. Transitions: reroute withChoice transitions via choice pseudostate
        for (Transition t : transitions) {
            if (t.sourceStates == null || t.sourceStates.isEmpty()) continue;

            boolean allowNoTarget = "withJoin".equals(t.type) || "withFork".equals(t.type);

            List<String> targets;
            if (t.targetStates == null || t.targetStates.isEmpty()) {
                if (!allowNoTarget) {
                    // self-loop: targets = sources
                    targets = t.sourceStates;
                } else {
                    // withJoin/withFork with no targets: skip
                    continue;
                }
            } else {
                targets = t.targetStates;
            }

            for (String rawSource : t.sourceStates) {
                String source = simplify(rawSource);

                for (String rawTarget : targets) {
                    String target = simplify(rawTarget);
                    String edgeSource = source;

                    if ("withChoice".equals(t.type) && statesWithChoice.contains(source)) {
                        edgeSource = source + "_choice";
                    }

                    // Prepare guard text
                    String guardText = "";
                    if (t.guard != null && !t.guard.isBlank()) {
                        if (useLambdaGuards && t.isLambdaGuard) {
                            guardText = "lambda";
                        } else {
                            guardText = t.guard.replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim();
                        }
                    }

                    // Prepare actions text
                    String actionsText = "";
                    if (t.actions != null && !t.actions.isEmpty()) {
                        StringBuilder actionBuilder = new StringBuilder();
                        for (int i = 0; i < t.actions.size(); i++) {
                            if (i > 0) actionBuilder.append(", ");
                            String action = t.actions.get(i);
                            boolean isLambda = t.isLambdaActions.get(i);
                            if (useLambdaGuards && isLambda) {
                                actionBuilder.append("lambda");
                            } else {
                                actionBuilder.append(action);
                            }
                        }
                        actionsText = actionBuilder.toString();
                    }

                    // Build label with event, guard, and actions
                    StringBuilder label = new StringBuilder();
                    if (t.event != null && !t.event.isBlank()) {
                        label.append(simplify(t.event));
                    }
                    if (!guardText.isBlank()) {
                        if (label.length() > 0) label.append(" ");
                        label.append("[").append(guardText).append("]");
                    }
                    if (!actionsText.isBlank()) {
                        if (label.length() > 0) label.append(" / ");
                        label.append(actionsText);
                    }

                    // Add order info for choice/junction transitions
                    if (("withChoice".equals(t.type) || "withJunction".equals(t.type)) && t.order != null) {
                        if (label.length() > 0) label.append(" ");
                        label.append("(order=").append(t.order).append(")");
                    }

                    String edgeLabel = label.length() > 0 ? " [label=\"" + label.toString() + "\"]" : "";

                    sb.append("  ").append(edgeSource).append(" -> ").append(target).append(edgeLabel).append(";\n");
                }
            }
        }

        sb.append("}\n");
        return sb.toString();
    }

    public static String exportToSCXML(List<Transition> transitions,
                                       Set<String> startStates,
                                       Set<String> endStates,
                                       boolean useLambdaGuards) {
        StringBuilder sb = new StringBuilder();
        sb.append("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
        sb.append("<scxml xmlns=\"http://www.w3.org/2005/07/scxml\" version=\"1.0\">\n\n");

        // 1. Collect all states
        Set<String> allStates = new HashSet<>();
        for (Transition t : transitions) {
            if (t.sourceStates != null) {
                for (String s : t.sourceStates) allStates.add(s);
            }
            if (t.targetStates != null) {
                for (String tState : t.targetStates) allStates.add(tState);
            }
        }

        // 2. Output <state> elements with transitions
        for (String state : allStates) {
            String simpleState = simplify(state);
            sb.append("  <state id=\"").append(simpleState).append("\">\n");

            // If this state is a start state, add <initial> and transition to it
            if (startStates.contains(state)) {
                sb.append("    <initial>\n");
                sb.append("      <transition target=\"").append(simpleState).append("\"/>\n");
                sb.append("    </initial>\n");
            }

            // Add transitions from this state
            for (Transition t : transitions) {
                if (t.sourceStates == null || !t.sourceStates.contains(state)) continue;

                boolean allowNoTarget = "withJoin".equals(t.type) || "withFork".equals(t.type);

                List<String> targets;
                if (t.targetStates == null || t.targetStates.isEmpty()) {
                    if (!allowNoTarget) {
                        // self-loop target = source
                        targets = List.of(state);
                    } else {
                        // withJoin/withFork with no target: skip
                        continue;
                    }
                } else {
                    targets = t.targetStates;
                }

                for (String rawTarget : targets) {
                    String target = simplify(rawTarget);
                    if (target == null || target.isEmpty()) continue;

                    String event = t.event != null ? simplify(t.event) : null;
                    String guard = null;
                    if (t.guard != null && !t.guard.isBlank()) {
                        if (useLambdaGuards && t.isLambdaGuard) {
                            guard = "lambda";
                        } else {
                            guard = t.guard.replaceAll("[\\n\\r]+", " ").trim();
                        }
                    }

                    sb.append("    <transition");
                    if (event != null && !event.isBlank()) sb.append(" event=\"").append(event).append("\"");
                    sb.append(" target=\"").append(target).append("\"");
                    if (guard != null && !guard.isBlank()) sb.append(" cond=\"").append(escapeXml(guard)).append("\"");

                    // Add order info for choice/junction transitions as XML comment inside transition
                    if (("withChoice".equals(t.type) || "withJunction".equals(t.type)) && t.order != null) {
                        sb.append(">\n");
                        sb.append("      <!-- order=").append(t.order).append(" -->\n");
                        sb.append("    </transition>\n");
                    } else {
                        sb.append("/>\n");
                    }
                }
            }

            sb.append("  </state>\n\n");
        }

        sb.append("</scxml>\n");
        return sb.toString();
    }

    // Helper: extract last enum part
    private static String simplify(String full) {
        if (full == null) return "";
        int dot = full.lastIndexOf('.');
        return dot >= 0 ? full.substring(dot + 1) : full;
    }

    private static String escapeXml(String s) {
        if (s == null) return "";
        return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
                .replace("\"", "&quot;").replace("'", "&apos;");
    }

}