main.go

package main

import (
	"encoding/json"
	"errors"
	ast "fizz/proto"
	"flag"
	"fmt"
	"github.com/fizzbee-io/fizzbee/lib"
	"github.com/fizzbee-io/fizzbee/modelchecker"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime/pprof"
	"slices"
	"strings"
	"sync/atomic"
	"syscall"
	"time"
)

var isPlayground bool
var simulation bool
var internalProfile bool
var saveStates bool
var copyAst bool
var outputDir string
var seed int64
var maxRuns int
var simFirstTraces int
var explorationStrategy string
var traceFile string
var preinitHookFile string
var trace string
var preinitHook string
var traceExtend int

var isTest bool
var experimentalProcessedQueue bool
var experimentalNoGraph bool
var experimentalNoStateReturns bool
var noSymmetryReduction bool

func main() {
	args := parseFlags()
	// Check if the correct number of arguments is provided
	if len(args) != 1 {
		fmt.Println("Usage:", os.Args[0], "<json_file>")
		os.Exit(1)
	}

	// --experimental_no_graph implies --experimental_processed_queue (the
	// drop-graph optimization only makes sense in the queue-of-processed
	// model). Auto-enable rather than erroring out.
	if experimentalNoGraph && !experimentalProcessedQueue {
		fmt.Println("Note: --experimental_no_graph auto-enables --experimental_processed_queue.")
		experimentalProcessedQueue = true
	}

	// Get the input JSON file name from command line argument
	jsonFilename := args[0]
	dirPath := filepath.Dir(jsonFilename)

	f := loadInputJSON(jsonFilename)

	// --experimental_no_graph cannot support liveness checks (no graph to
	// run SCC/cycle detection on). Refuse upfront if the spec declares any
	// always-eventually / eventually-always / eventually invariants.
	if experimentalNoGraph {
		for _, inv := range f.Invariants {
			for _, op := range inv.TemporalOperators {
				if op == "eventually" || op == "always-eventually" || op == "eventually-always" {
					fmt.Printf("--experimental_no_graph refused: spec declares liveness invariant '%s' (operator '%s'). The no-graph mode keeps no state graph, so cycle-based liveness checks cannot run.\n", inv.Name, op)
					os.Exit(1)
				}
			}
		}
	}

	sourceFileName := filepath.Join(dirPath, f.SourceInfo.GetFileName())
	//fmt.Println("dirPath:", dirPath)
	// Calculate the relative path
	var stateConfig *ast.StateSpaceOptions
	if traceFile != "" || trace != "" || traceExtend > 0 {
		stateConfig = getStateConfigForTraceChecking(stateConfig)
	} else {
		stateConfig = loadStateOptions(dirPath, f.GetFrontMatter())
		fmt.Printf("StateSpaceOptions: %+v\n", stateConfig)
		applyDefaultStateOptions(stateConfig)
	}

	// no_graph mode cannot drive crash simulation (the crash variants
	// rely on the in-memory graph: Attach/visited/Inbound). Warn loudly
	// when crash_on_yield is on; the per-yield-point crashProcess /
	// crashThread calls are no-oped below in the processor.
	if experimentalNoGraph && stateConfig.GetOptions().GetCrashOnYield() {
		fmt.Println("WARNING: --experimental_no_graph disables crash-on-yield simulation. Crash variants will NOT be explored. Set crash_on_yield: false to silence this warning, or omit --experimental_no_graph to keep crash simulation.")
	}

	outDir, err := createOutputDir(dirPath, isTest)
	if err != nil {
		return
	}
	err = dumpStateSpaceOptions(stateConfig, outDir)
	if err != nil {
		fmt.Printf("Error writing state space options: %v\n", err)
		return
	}

	if f.Composition != nil {
		runCompositionalModelChecking(f, dirPath, outDir)
	} else if len(f.Refinements) > 0 {
		if len(f.Refinements) > 1 {
			fmt.Println("Multiple refinements found. Only single refinement is supported currently. Contact us if you need support for multiple refinements.")
			return
		}
		runRefinementModelChecking(f, dirPath, outDir, jsonFilename)
	} else {
		// For single spec model checking, copy AST to the output directory
		if copyAst {
			err = copyAstToOutputDir(jsonFilename, outDir)
			if err != nil {
				fmt.Printf("Error copying AST file: %v\n", err)
				return
			}
		}
		modelCheckSingleSpec(f, stateConfig, dirPath, outDir, sourceFileName, jsonFilename, nil)
	}
}

func dumpStateSpaceOptions(stateConfig *ast.StateSpaceOptions, outDir string) error {
	stateConfigJson, err := protojson.MarshalOptions{
		Multiline: true,
	}.Marshal(stateConfig)
	if err != nil {
		return err
	}
	err = os.WriteFile(filepath.Join(outDir, "state_config.json"), stateConfigJson, 0644)
	return err
}

func copyAstToOutputDir(jsonFilename string, outDir string) error {
	// Read the AST JSON file
	astContent, err := os.ReadFile(jsonFilename)
	if err != nil {
		return fmt.Errorf("failed to read AST file: %w", err)
	}

	// Create the destination path with fixed name spec_ast.json
	destPath := filepath.Join(outDir, "spec_ast.json")

	// Write the AST to the output directory
	err = os.WriteFile(destPath, astContent, 0644)
	if err != nil {
		return fmt.Errorf("failed to write AST file: %w", err)
	}

	return nil
}

func runCompositionalModelChecking(f *ast.File, dirPath string, outDir string) {
	if simulation {
		fmt.Println("Simulation mode not supported for composition")
		return
	}
	roots := make(map[string]*modelchecker.Node)
	rootsList := make([]*modelchecker.Node, len(f.Composition.GetSpecs()))
	// create a map to store the interface preserving states
	// key is the hash, value is the list of (list of Nodes)
	joinHashes := make(modelchecker.JoinHashes)
	for i, spec := range f.Composition.GetSpecs() {
		fmt.Printf("\nComposed Spec: %s\n", spec.Name)
		fileRef, fnRef, err := ParseFunctionRef(spec.Expr.GetPyExpr())
		if err != nil {
			fmt.Println("Error parsing function reference:", err)
			return
		} else if fileRef != spec.Name {
			fmt.Println("File reference does not match spec name:", fileRef, "!=", spec.Name)
			return
		}
		// Get the new json file name as dirPath + "/" + name + ".json"
		composedJsonFileName := filepath.Join(dirPath, spec.Name+".json")
		composedFile := loadInputJSON(composedJsonFileName)
		composedSourceFileName := filepath.Join(dirPath, composedFile.SourceInfo.GetFileName())
		composedStateConfig := loadStateOptions(dirPath, composedFile.GetFrontMatter())
		applyDefaultStateOptions(composedStateConfig)
		composedOutDir := filepath.Join(outDir, spec.Name)
		if err := os.MkdirAll(composedOutDir, 0755); err != nil {
			fmt.Println("Error creating directory:", err)
			return
		}
		if copyAst {
			err = copyAstToOutputDir(composedJsonFileName, composedOutDir)
			if err != nil {
				fmt.Printf("Error copying AST file: %v\n", err)
				return
			}
		}
		err = dumpStateSpaceOptions(composedStateConfig, composedOutDir)
		if err != nil {
			fmt.Printf("Error writing state space options: %v\n", err)
			return
		}
		fmt.Println("Model checking composed spec:", composedSourceFileName)
		root := modelCheckSingleSpec(composedFile, composedStateConfig, dirPath, composedOutDir, composedSourceFileName, composedJsonFileName, nil)
		if root == nil {
			fmt.Println("Error in model checking composed spec:", spec.Name, "Aborting")
			return
		}
		roots[spec.Name] = root
		rootsList[i] = root
		populateJoinHashes(joinHashes, i, root, fileRef, fnRef)
	}

	err := ComposeTransitions(f, joinHashes, rootsList, outDir)
	if err != nil {
		fmt.Println("Error composing transitions:", err)
		return
	}
	return
}

func runRefinementModelChecking(f *ast.File, dirPath string, outDir string, jsonFilename string) {
	if simulation {
		fmt.Println("Simulation mode not supported for refinement")
		return
	}

	if len(f.Refinements) == 0 {
		fmt.Println("No refinement specification found")
		return
	}

	// For now, assume only one refinement block
	refinement := f.Refinements[0]

	// Create map of roots for abstract specs
	roots := make(map[string]*modelchecker.Node)
	rootsList := make([]*modelchecker.Node, len(refinement.GetSpecs()))
	joinHashes := make(modelchecker.JoinHashes)

	var implSpecIndex int = -1

	// First, run model checking for each abstract spec (not "_")
	for i, spec := range refinement.GetSpecs() {
		if spec.Name == "_" {
			implSpecIndex = i
			continue
		}

		fmt.Printf("\nRefined Spec: %s\n", spec.Name)
		fileRef, fnRef, err := ParseFunctionRef(spec.Expr.GetPyExpr())
		if err != nil {
			fmt.Println("Error parsing function reference:", err)
			return
		} else if fileRef != spec.Name {
			fmt.Println("File reference does not match spec name:", fileRef, "!=", spec.Name)
			return
		}

		abstractJsonFile := filepath.Join(dirPath, spec.Name+".json")
		abstractFile := loadInputJSON(abstractJsonFile)
		abstractSourceFile := filepath.Join(dirPath, abstractFile.SourceInfo.GetFileName())

		stateConfig := loadStateOptions(dirPath, abstractFile.GetFrontMatter())
		applyDefaultStateOptions(stateConfig)

		abstractOutDir := filepath.Join(outDir, spec.Name)
		if err := os.MkdirAll(abstractOutDir, 0755); err != nil {
			fmt.Println("Error creating directory:", err)
			return
		}
		if copyAst {
			err = copyAstToOutputDir(abstractJsonFile, abstractOutDir)
			if err != nil {
				fmt.Printf("Error copying AST file: %v\n", err)
				return
			}
		}
		err = dumpStateSpaceOptions(stateConfig, abstractOutDir)
		if err != nil {
			fmt.Printf("Error writing state space options: %v\n", err)
			return
		}

		fmt.Println("Model checking abstract spec:", abstractSourceFile)
		root := modelCheckSingleSpec(abstractFile, stateConfig, dirPath, abstractOutDir, abstractSourceFile, abstractJsonFile, nil)
		if root == nil {
			fmt.Println("Error in model checking abstract spec:", spec.Name)
			return
		}

		roots[spec.Name] = root
		rootsList[i] = root
		populateJoinHashes(joinHashes, i, root, fileRef, fnRef)
	}

	if implSpecIndex == -1 {
		fmt.Println("Error: no implementation (_) spec found in refinement block")
		return
	}

	// Now run model checking on the implementation (this file), using joinHashes
	implOutDir := filepath.Join(outDir, "_")
	if err := os.MkdirAll(implOutDir, 0755); err != nil {
		fmt.Println("Error creating directory:", err)
		return
	}

	if copyAst {
		err := copyAstToOutputDir(jsonFilename, implOutDir)
		if err != nil {
			fmt.Printf("Error copying AST file: %v\n", err)
			return
		}
	}

	fmt.Println("Model checking implementation spec (refined):", f.SourceInfo.GetFileName())
	stateConfig := loadStateOptions(dirPath, f.GetFrontMatter())
	applyDefaultStateOptions(stateConfig)
	err := dumpStateSpaceOptions(stateConfig, implOutDir)
	if err != nil {
		fmt.Printf("Error writing state space options: %v\n", err)
		return
	}

	root := modelCheckSingleSpec(f, stateConfig, dirPath, implOutDir, f.SourceInfo.GetFileName(), jsonFilename, joinHashes)
	if root == nil {
		fmt.Println("Error in model checking implementation spec")
		return
	}

	return
}

type ComposedNode = struct {
	*modelchecker.Node
	Nodes []*modelchecker.Node
}

type ComposedTransition = struct {
	From *ComposedNode
	To   *ComposedNode
}

func cartesianProduct(sets []modelchecker.TransitionSet, handler func(ts []lib.Pair[*modelchecker.Node, *modelchecker.Node]) (
	*ComposedNode, *ComposedTransition)) (*ComposedNode, *ComposedTransition) {

	var helper func(int, []lib.Pair[*modelchecker.Node, *modelchecker.Node])
	var failedNode *ComposedNode
	var failedTransition *ComposedTransition
	helper = func(depth int, current []lib.Pair[*modelchecker.Node, *modelchecker.Node]) {
		if depth == len(sets) {
			failedNode, failedTransition = handler(current)
			return
		}
		for t := range sets[depth] {
			helper(depth+1, append(current, t))
			if failedNode != nil || failedTransition != nil {
				return // Stop further processing if a failure is found
			}
		}
	}
	helper(0, []lib.Pair[*modelchecker.Node, *modelchecker.Node]{})
	if failedTransition != nil {
		fmt.Println("Failed Transition:", failedTransition.From.Nodes, "->", failedTransition.To.Nodes)
	} else if failedNode != nil {
		fmt.Println("Failed Node:", failedNode.Nodes)
	}
	return failedNode, failedTransition
}

func ComposeTransitions(f *ast.File, joinHashes modelchecker.JoinHashes, roots []*modelchecker.Node, outDir string) error {
	hasTransitionInvariant := false
	for _, invariant := range f.Invariants {
		if invariant.Block != nil && slices.Contains(invariant.TemporalOperators, "transition") {
			hasTransitionInvariant = true
			break
		}
	}
	for key, sets := range joinHashes {
		if len(sets) < 2 {
			return fmt.Errorf("need at least 2 transition sets for key %s", key)
		}

		failedNode, failedTransition := cartesianProduct(sets, func(ts []lib.Pair[*modelchecker.Node, *modelchecker.Node]) (*ComposedNode, *ComposedTransition) {
			// Check if all transitions are stuttering
			allStuttering := true
			for _, t := range ts {
				if t.First.Heap.HashCode() != t.Second.Heap.HashCode() {
					allStuttering = false
					break
				}
			}
			if allStuttering {
				return nil, nil // Skip all-stuttering transitions
			}
			composedFrom := make(map[string]*modelchecker.Heap)
			composedTo := make(map[string]*modelchecker.Heap)
			from := &ComposedNode{Nodes: make([]*modelchecker.Node, len(ts))}
			to := &ComposedNode{Nodes: make([]*modelchecker.Node, len(ts))}

			for i, t := range ts {
				from.Nodes[i] = t.First
				to.Nodes[i] = t.Second

				composedFrom[f.Composition.GetSpecs()[i].GetName()] = t.First.Heap
				composedTo[f.Composition.GetSpecs()[i].GetName()] = t.Second.Heap
			}

			composed := &ComposedTransition{From: from, To: to}
			// Handle composed transition
			processTo := modelchecker.NewProcess("yield", []*ast.File{f}, nil)
			processTo.Enable()
			processTo.Heap = modelchecker.NewComposedHeap(composedTo)
			failedInvariants := modelchecker.CheckInvariants(processTo)
			to.Node = modelchecker.NewNode(processTo)
			if len(failedInvariants) > 0 && len(failedInvariants[0]) > 0 {
				processTo.FailedInvariants = failedInvariants
				fmt.Println("Composed state failed invariants:", failedInvariants)
				processTo.FailedInvariants = failedInvariants
				return to, nil
			}
			if hasTransitionInvariant {
				processFrom := modelchecker.NewProcess("yield", []*ast.File{f}, nil)
				processFrom.Enable()
				processFrom.Heap = modelchecker.NewComposedHeap(composedFrom)
				processTo.Parent = processFrom
				failedTransitionInvariants := modelchecker.CheckTransitionInvariants(processTo)
				if len(failedTransitionInvariants) > 0 && len(failedTransitionInvariants[0]) > 0 {
					//processTo.FailedInvariants = failedInvariants
					fmt.Println("Composed transition failed invariants:", failedTransitionInvariants)
					link := &modelchecker.Link{
						Node:             to.Node,
						Name:             "composed-transition",
						FailedInvariants: failedTransitionInvariants,
					}
					from.Node = modelchecker.NewNode(processFrom)
					from.Outbound = append(from.Outbound, link)
					to.Inbound = append(to.Inbound, link)
					return to, composed
				}
			}

			return nil, nil
		})
		if failedNode != nil || failedTransition != nil {
			if failedTransition != nil {
				fmt.Println("Failed Composed Node:", failedTransition)
				printComposedTransition(failedTransition)
				dumpComposedFailedTransition(failedTransition, f.Composition, roots, outDir)
			} else /*if failedNode != nil*/ {
				fmt.Println("Failed Composed Node:", failedNode.Nodes)
				dumpComposedFailedNode(failedNode, f.Composition, roots, outDir)
			}
			return fmt.Errorf("failed to compose transitions for key %s", key)
		}
	}
	return nil
}

func printComposedTransition(composed *ComposedTransition) {
	fmt.Printf("Composed transition: ")
	for i := range composed.From.Nodes {
		fmt.Printf("(%s -> %s) ", composed.From.Nodes[i].Heap.String(), composed.To.Nodes[i].Heap.String())
	}
	fmt.Println()
}

func ParseFunctionRef(input string) (string, string, error) {
	input = strings.TrimSpace(input)

	if input == "" {
		return "", "", errors.New("input is empty")
	}

	parts := strings.Split(input, ".")
	switch len(parts) {
	//case 1:
	//	funcName := strings.TrimSpace(parts[0])
	//	if funcName == "" {
	//		return "", "", errors.New("function name is required")
	//	}
	//	return "", funcName, nil
	case 2:
		fileRef := strings.TrimSpace(parts[0])
		funcName := strings.TrimSpace(parts[1])
		if fileRef == "" || funcName == "" {
			return "", "", errors.New("file reference or function name is empty")
		}
		return fileRef, funcName, nil
	default:
		return "", "", errors.New("invalid function reference format. Expected format: 'fileName.funcName'")
	}
}

func populateJoinHashes(joinHashes modelchecker.JoinHashes, i int, root *modelchecker.Node, fileRef string, fnRef string) {
	visited := make(map[*modelchecker.Node]bool)

	var dfs func(from *modelchecker.Node, current *modelchecker.Node)
	dfs = func(from *modelchecker.Node, current *modelchecker.Node) {

		if from != nil && current.IsYieldNode() {
			addToJoinHashes(joinHashes, i, from, current, fnRef)
		}
		if visited[current] {
			return
		}
		visited[current] = true

		if current.IsYieldNode() {
			// Add stuttering (to, to) link for new Nodes discovered
			addToJoinHashes(joinHashes, i, current, current, fnRef)
			from = current // update 'from' for child traversal
		}

		for _, child := range current.Outbound {
			if child.Node != nil {
				dfs(from, child.Node)
			}
		}
	}

	dfs(nil, root)
}

func addToJoinHashes(joinHashes modelchecker.JoinHashes, i int, from, to *modelchecker.Node, fnName string) {
	//fmt.Println("Adding to join hashes:", from.Heap.String(), "->", to.Heap.String(), "for function", fnName)
	fromState := modelchecker.ExecFunction(from.Process.CloneForAssert(nil, 0), fnName).String()
	toState := modelchecker.ExecFunction(to.Process.CloneForAssert(nil, 0), fnName).String()

	key := modelchecker.HashKey(fromState + "," + toState)

	if _, exists := joinHashes[key]; !exists {
		joinHashes[key] = make([]modelchecker.TransitionSet, i+1)
	}
	for len(joinHashes[key]) <= i {
		joinHashes[key] = append(joinHashes[key], make(modelchecker.TransitionSet))
	}
	if joinHashes[key][i] == nil {
		joinHashes[key][i] = make(modelchecker.TransitionSet)
	}
	pair := lib.NewPair(from, to)
	joinHashes[key][i][pair] = true
}

// replayTraceWithFullState writes the action-name path to a trace file
// (so the user has it on disk too) and re-execs the binary with
// --trace-file pointing at it. The replay runs in default mode (with
// graph) so the standard error-report machinery prints the per-step
// state — converting the lightweight no_graph trace into the full
// rich output that the user would have seen without --experimental_no_graph.
// Best-effort: failures to write/exec just print a fallback command.
func replayTraceWithFullState(path []string, outDir, jsonFilename string) {
	if len(path) == 0 || jsonFilename == "" {
		return
	}
	traceFile := filepath.Join(outDir, "trace.txt")
	body := strings.Join(path, "\n") + "\n"
	if err := os.WriteFile(traceFile, []byte(body), 0644); err != nil {
		fmt.Printf("(replay skipped: %v)\n", err)
		return
	}
	self, err := os.Executable()
	if err != nil {
		fmt.Printf("(replay skipped: %v)\nReplay manually with:\n  fizz --trace-file %s <spec.fizz>\n", err, traceFile)
		return
	}
	// Preserve flags that change which spec state space the replay sees.
	// preinit-hook flags override spec constants; if we drop them the
	// replay reads a different spec (different NUM_WRITERS, etc.) and the
	// trace won't match anything. Everything else is safe to leave at
	// default for the replay — guided trace fixes the path either way.
	args := []string{"--trace-file", traceFile}
	if preinitHookFile != "" {
		args = append(args, "--preinit-hook-file", preinitHookFile)
	}
	if preinitHook != "" {
		args = append(args, "--preinit-hook", preinitHook)
	}
	args = append(args, jsonFilename)
	fmt.Println()
	fmt.Println("=== Replaying trace with full per-step state ===")
	cmd := exec.Command(self, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	_ = cmd.Run()
}

func modelCheckSingleSpec(f *ast.File, stateConfig *ast.StateSpaceOptions, dirPath string, outDir string, sourceFileName string, jsonFilename string, hashes modelchecker.JoinHashes) *modelchecker.Node {
	// Parse trace if provided (either from file or string)
	var guidedTrace *modelchecker.GuidedTrace
	if traceFile != "" {
		var err error
		guidedTrace, err = modelchecker.ParseTraceFile(traceFile)
		if err != nil {
			fmt.Printf("Error parsing trace file: %v\n", err)
			return nil
		}
		fmt.Printf("Loaded trace with %d links\n", len(guidedTrace.LinkNames))

		// Trace mode is incompatible with simulation
		if simulation {
			fmt.Println("Error: --trace-file and --simulation cannot be used together")
			return nil
		}
	} else if trace != "" {
		var err error
		guidedTrace, err = modelchecker.ParseTraceString(trace)
		if err != nil {
			fmt.Printf("Error parsing trace string: %v\n", err)
			return nil
		}
		fmt.Printf("Loaded trace with %d links\n", len(guidedTrace.LinkNames))

		// Trace mode is incompatible with simulation
		if simulation {
			fmt.Println("Error: --trace and --simulation cannot be used together")
			return nil
		}
	} else if traceExtend > 0 {
		// No trace specified but --trace-extend given: explore traceExtend steps
		// from the initial state (empty trace). guidedTrace with no LinkNames but
		// ExtendDepth > 0 signals "explore N steps from Init" to ShouldScheduleNode.
		if simulation {
			fmt.Println("Error: --trace-extend and --simulation cannot be used together")
			return nil
		}
		guidedTrace = &modelchecker.GuidedTrace{}
	}

	// Set trace extension depth if configured
	if guidedTrace != nil && traceExtend > 0 {
		guidedTrace.ExtendDepth = traceExtend
		fmt.Printf("Trace will extend %d action(s) beyond the guided path\n", traceExtend)
	}

	// Resolve preinit hook content (either from file or string)
	var preinitHookContentResolved string
	if preinitHookFile != "" {
		content, err := os.ReadFile(preinitHookFile)
		if err != nil {
			fmt.Printf("Error reading preinit hook file: %v\n", err)
			return nil
		}
		preinitHookContentResolved = string(content)
	} else if preinitHook != "" {
		preinitHookContentResolved = preinitHook
	}

	//maxRuns := 10000
	if !simulation || seed != 0 {
		maxRuns = 1
	}
	if simulation && seed != 0 {
		fmt.Println("Seed:", seed)
	}
	if simulation && maxRuns == 0 {
		fmt.Println("MaxRuns: unlimited")
	}
	stopped := false
	runs := 0
	// Aggregates exists-assertion witnesses across simulation runs (each
	// run uses a fresh Processor). Simulation is not exhaustive, so a
	// never-witnessed exists at the end is a coverage WARNING, not a
	// failure.
	var simUnsatisfiedExists []*modelchecker.InvariantPosition
	simExistsTracked := false
	var p1 *modelchecker.Processor
	var holder atomic.Pointer[modelchecker.Processor]
	var lastRootNode *modelchecker.Node

	setupSignalHandler(&holder, &stopped)

	// Periodic simulation progress: print "\rSimulated N runs" to stderr
	// every 100 runs when stderr is a TTY. Silent otherwise so log files,
	// pipes, and parallel-worker captures stay clean. simProgressEmitted
	// tracks whether we ever wrote a \r line; if so, callers below must
	// emit \n to stderr before any non-progress output so the next message
	// doesn't append to the progress line.
	simProgressOnTty := false
	simProgressEmitted := false
	if simulation {
		if fi, err := os.Stderr.Stat(); err == nil && (fi.Mode()&os.ModeCharDevice) != 0 {
			simProgressOnTty = true
		}
	}
	clearProgressLine := func() {
		if simProgressEmitted {
			fmt.Fprintln(os.Stderr)
			simProgressEmitted = false
		}
	}

	i := 0
	for !stopped && (maxRuns <= 0 || i < maxRuns) {
		i++

		p1 = modelchecker.NewProcessor([]*ast.File{f}, stateConfig, simulation, seed, dirPath, explorationStrategy, isTest, hashes, guidedTrace, preinitHookContentResolved)
		p1.SetExperimentalProcessedQueue(experimentalProcessedQueue)
		p1.SetExperimentalNoGraph(experimentalNoGraph)
		p1.SetExperimentalNoStateReturns(experimentalNoStateReturns)
		p1.SetDisableSymmetryReduction(noSymmetryReduction)
		holder.Store(p1)

		rootNode, failedNode, endTime, err := startModelChecker(p1)
		runs++
		lastRootNode = rootNode

		if simulation && guidedTrace == nil {
			// Keep only the exists positions that NO run so far has
			// witnessed: intersect with this run's unsatisfied set.
			if !simExistsTracked {
				simUnsatisfiedExists = p1.UnsatisfiedExists()
				simExistsTracked = true
			} else if len(simUnsatisfiedExists) > 0 {
				cur := make(map[modelchecker.InvariantPosition]bool)
				for _, pos := range p1.UnsatisfiedExists() {
					cur[*pos] = true
				}
				remaining := simUnsatisfiedExists[:0]
				for _, pos := range simUnsatisfiedExists {
					if cur[*pos] {
						remaining = append(remaining, pos)
					}
				}
				simUnsatisfiedExists = remaining
			}
		}

		if simProgressOnTty && runs%100 == 0 {
			// \033[K clears any leftover characters from a prior longer
			// line (e.g. if maxRuns digit count shrinks across runs).
			if maxRuns > 0 {
				fmt.Fprintf(os.Stderr, "\r\033[KSimulated %d / %d runs", runs, maxRuns)
			} else {
				fmt.Fprintf(os.Stderr, "\r\033[KSimulated %d runs", runs)
			}
			simProgressEmitted = true
		}

		// Early deadlock detection (NEW path only): startProcessedQueue
		// aborts the run on the first detected deadlock and exposes it
		// here. Report and exit before any post-run graph traversal — no
		// point doing the rest of the work, and dumpFailedNode reads
		// node.Inbound which is still populated for this node. In
		// no-graph mode the Inbound back-pointer was cleared, so fall
		// back to the lightweight pathTail trace.
		if earlyDead := p1.GetEarlyDeadlock(); earlyDead != nil && !simulation {
			fmt.Println("DEADLOCK detected (early)")
			fmt.Println("FAILED: Model checker failed")
			if experimentalNoGraph {
				path := earlyDead.PathNames()
				fmt.Printf("Trace (%d steps):\n", len(path))
				for _, step := range path {
					fmt.Printf("  %s\n", step)
				}
				replayTraceWithFullState(path, outDir, jsonFilename)
			} else {
				dumpFailedNode(sourceFileName, earlyDead, rootNode, outDir)
			}
			return nil
		}

		if !simulation && !experimentalNoGraph {
			if writeDotFileIfNeeded(p1, rootNode, outDir, "graph.dot") {
				return nil
			}
		} else if i <= simFirstTraces {
			// One file per captured run so the first N traces are
			// preserved instead of overwriting each other on a shared
			// graph.dot path. The "last run" still lands at graph.dot
			// below (line 903) so existing tooling reading that canonical
			// name keeps working.
			writeDotFileIfNeeded(p1, rootNode, outDir, fmt.Sprintf("graph_run_%d.dot", i))
		}

		if err != nil {
			printTraceAndExit(err)
		}

		// In no-graph mode there's no in-memory state graph to walk for
		// dot generation, exists-witness checking, liveness, or deadlock
		// detection via traverseBFS. Print a minimal summary and return.
		// (Liveness was already refused at startup; deadlock detection in
		// no-graph mode happens at expansion time inside startProcessedQueue.)
		if experimentalNoGraph {
			if failedNode != nil {
				if failedNode.FailedInvariants != nil && len(failedNode.FailedInvariants) > 0 && len(failedNode.FailedInvariants[0]) > 0 {
					fmt.Println("FAILED: Model checker failed. Invariant:", f.Invariants[failedNode.FailedInvariants[0][0]].Name)
				} else {
					fmt.Println("FAILED: Model checker failed")
				}
				path := failedNode.PathNames()
				fmt.Printf("Trace (%d steps):\n", len(path))
				for _, step := range path {
					fmt.Printf("  %s\n", step)
				}
				replayTraceWithFullState(path, outDir, jsonFilename)
				return nil
			}
			fmt.Printf("Visited entries: %d  Unique states: %d\n", p1.GetVisitedNodesCount(), p1.GetUniqueYieldCount())
			if p1.Stopped() {
				fmt.Println("Model checker stopped")
				return nil
			}
			// exists assertions: no retained graph to walk, but the
			// processor accumulated witness flags across every explored
			// state. Exploration here is exhaustive, so a missing witness
			// is proof of unreachability — same verdict and format as the
			// full-graph CheckSimpleExistsWitness path. Skipped in trace
			// mode for the same reason as the full-graph path: a guided
			// trace only explores a slice of the state space.
			if guidedTrace == nil {
				if unsatisfied := p1.UnsatisfiedExists(); len(unsatisfied) > 0 {
					fmt.Println("\nFAILED: Expected states never reached")
					for i2, invariant := range unsatisfied {
						fmt.Printf("Invariant %d: %s\n", i2, f.Invariants[invariant.InvariantIndex].Name)
					}
					return nil
				}
			}
			fmt.Println("PASSED: Model checker completed successfully")
			return rootNode
		}

		// Check if trace was fully executed
		if guidedTrace != nil && !guidedTrace.IsExhausted() {
			fmt.Printf("WARNING: Trace execution incomplete. Expected %d links, executed %d links.\n",
				len(guidedTrace.LinkNames), guidedTrace.GetCurrentIndex())
			fmt.Println("The trace may contain links that don't match the model or are unreachable.")
			return nil
		}

		//fmt.Println("root", root)
		if failedNode == nil {
			var failurePath []*modelchecker.Link
			var failedInvariant *modelchecker.InvariantPosition
			nodes, messages, deadlock, yieldsCount := modelchecker.GetAllNodes(rootNode, stateConfig.GetOptions().GetMaxActions())

			if writeCommunicationFileIfNeeded(messages, outDir) {
				return nil
			}

			if deadlock != nil && stateConfig.GetDeadlockDetection() && !p1.Stopped() && !simulation {
				fmt.Println("DEADLOCK detected")
				fmt.Println("FAILED: Model checker failed")
				if simulation {
					fmt.Println("seed:", p1.Seed)
				}
				dumpFailedNode(sourceFileName, deadlock, rootNode, outDir)
				return nil
			}
			if !simulation {
				fmt.Println("Valid Nodes:", len(nodes), "Unique states:", yieldsCount)
				// Skip the exists-witness check in trace mode: a guided trace
				// only explores a slice of the state space, so `exists`
				// invariants legitimately can't be satisfied yet. Letting the
				// failure short-circuit here would also skip writing the
				// nodes/links pb files, which the trace consumer needs.
				if guidedTrace == nil {
					invariants := modelchecker.CheckSimpleExistsWitness(nodes)
					if len(invariants) > 0 {
						fmt.Println("\nFAILED: Expected states never reached")
						for i2, invariant := range invariants {
							fmt.Printf("Invariant %d: %s\n", i2, f.Invariants[invariant.InvariantIndex].Name)
						}
						if !isTest {
							fmt.Println("Time taken to check invariant: ", time.Now().Sub(endTime))
						}
						return nil
					}
				}
			}
			if !simulation && !p1.Stopped() && guidedTrace == nil {
				if stateConfig.GetLiveness() == "" || stateConfig.GetLiveness() == "enabled" || stateConfig.GetLiveness() == "true" || stateConfig.GetLiveness() == "strict" || stateConfig.GetLiveness() == "strict/bfs" {
					failurePath, failedInvariant = modelchecker.CheckStrictLiveness(rootNode, nodes)
				} else if stateConfig.GetLiveness() == "eventual" || stateConfig.GetLiveness() == "nondeterministic" {
					failurePath, failedInvariant = modelchecker.CheckFastLiveness(nodes)
				}
				fmt.Printf("IsLive: %t\n", failedInvariant == nil)
				if !isTest {
					fmt.Printf("Time taken to check liveness: %v\n", time.Now().Sub(endTime))
				}
			}

			if failedInvariant == nil && !simulation {
				if p1.Stopped() {
					fmt.Println("Model checker stopped")
					return nil
				}
				fmt.Println("PASSED: Model checker completed successfully")
				//Nodes, _, _ := modelchecker.GetAllNodes(rootNode)
				if saveStates || !isPlayground {
					nodeFiles, linkFileNames, err := modelchecker.GenerateProtoOfJson(nodes, outDir+"/")
					if err != nil {
						fmt.Println("Error generating proto files:", err)
						return rootNode
					}
					fmt.Printf("Writen %d node files and %d link files to dir %s\n", len(nodeFiles), len(linkFileNames), outDir)
				}
				return rootNode
			} else if failedInvariant != nil {
				fmt.Println("FAILED: Liveness check failed")
				if failedInvariant.FileIndex > 0 {
					fmt.Printf("Only one file expected. Got %d\n", failedInvariant.FileIndex)
				} else {
					fmt.Printf("Invariant: %s\n", f.Invariants[failedInvariant.InvariantIndex].Name)
				}
				GenerateFailurePath(sourceFileName, failurePath, failedInvariant, outDir)
				_, _, err = modelchecker.GenerateErrorPathProtoOfJson(failurePath, outDir+"/")
				if err != nil {
					fmt.Println("Error writing files", err)
				}
				return nil
			}

		} else if failedNode != nil {
			clearProgressLine()
			// Always-assertion failures live on the node's Process.FailedInvariants.
			// Transition-assertion failures live on the *inbound link* — see
			// processor.go where CheckTransitionInvariants writes to
			// node.Inbound[0].FailedInvariants. Without checking the link, a
			// transition-assertion failure in simulation falls through to the
			// misleading "Deadlock/stuttering" message, and in model-checking
			// mode nothing about the failure is printed at all.
			if failedNode.FailedInvariants != nil && len(failedNode.FailedInvariants) > 0 && len(failedNode.FailedInvariants[0]) > 0 {
				fmt.Println("FAILED: Model checker failed. Invariant: ", f.Invariants[failedNode.FailedInvariants[0][0]].Name)
			} else if len(failedNode.Inbound) > 0 && failedNode.Inbound[0].FailedInvariants != nil && len(failedNode.Inbound[0].FailedInvariants[0]) > 0 {
				fmt.Println("FAILED: Model checker failed. Transition Invariant:", f.Invariants[failedNode.Inbound[0].FailedInvariants[0][0]].Name)
			} else if simulation {
				fmt.Println("FAILED: Model checker failed. Deadlock/stuttering detected")
			} else {
				fmt.Println("FAILED: Model checker failed (no failed-invariant metadata on the failing node; see trace below).")
			}
			if simulation {
				fmt.Println("seed:", p1.Seed)
			}
			dumpFailedNode(sourceFileName, failedNode, rootNode, outDir)
			return nil
		}
	}
	clearProgressLine()
	fmt.Println("Stopped after", runs, "runs at ", time.Now())
	if simulation && len(simUnsatisfiedExists) > 0 {
		fmt.Printf("WARNING: %d exists assertion(s) never satisfied in %d simulation run(s):\n", len(simUnsatisfiedExists), runs)
		for _, pos := range simUnsatisfiedExists {
			fmt.Printf("  exists assertion %s\n", f.Invariants[pos.InvariantIndex].Name)
		}
		fmt.Println("Simulation is not exhaustive: this may be a coverage gap rather than unreachability. Model check to verify.")
	}
	if simulation && p1 != nil {
		if runs-simFirstTraces > 1 {
			fmt.Println("Not printing intermediate traces (only the last trace is shown)")
		}
		// Always write the canonical graph.dot so downstream tooling
		// (skills/scripts/docs that grep `out/*/graph.dot`) finds it
		// even when --simulation_first_traces >= runs. When the last
		// run was also captured as graph_run_<runs>.dot, this is a
		// duplicate write — same content under the canonical name.
		writeDotFileIfNeeded(p1, lastRootNode, outDir, "graph.dot")
	}
	return nil
}

func writeCommunicationFileIfNeeded(messages []string, outDir string) bool {
	if len(messages) > 0 && !simulation {
		graphDot := modelchecker.GenerateCommunicationGraph(messages)
		dotFileName := filepath.Join(outDir, "communication.dot")
		// Write the content to the file
		err := os.WriteFile(dotFileName, []byte(graphDot), 0644)
		if err != nil {
			fmt.Println("Error writing to file:", err)
			return true
		}
		if !isPlayground {
			fmt.Printf("Writen communication diagram dotfile: %s\nTo generate svg, run: \n"+
				"dot -Tsvg %s -o communication.svg && open communication.svg\n", dotFileName, dotFileName)
		}
	}
	return false
}

func printTraceAndExit(err error) {
	var modelErr *modelchecker.ModelError
	if errors.As(err, &modelErr) {
		fmt.Println("Stack Trace:")
		fmt.Println(modelErr.SprintStackTrace())
	} else {
		fmt.Println("Error:", err)
	}
	os.Exit(1)
}

// writeDotFileIfNeeded writes the graph .dot file for the given root node
// to outDir/fileName when the visited-node count is small enough to be
// useful. fileName lets callers separate per-run outputs in simulation
// mode (graph_run_1.dot, graph_run_2.dot, …) from the canonical
// "last result" graph.dot used by the model-check path.
func writeDotFileIfNeeded(p1 *modelchecker.Processor, rootNode *modelchecker.Node, outDir string, fileName string) bool {
	if p1.GetVisitedNodesCount() < 250 || (simulation && p1.GetVisitedNodesCount() < 1000) {
		dotString := modelchecker.GenerateDotFile(rootNode, make(map[*modelchecker.Node]bool))
		dotFileName := filepath.Join(outDir, fileName)
		// Write the content to the file
		err := os.WriteFile(dotFileName, []byte(dotString), 0644)
		if err != nil {
			fmt.Println("Error writing to file:", err)
			return true
		}
		if !isPlayground {
			fmt.Printf("Writen graph dotfile: %s\nTo generate svg, run: \n"+
				"dot -Tsvg %s -o graph.svg && open graph.svg\n", dotFileName, dotFileName)
		}
	} else if !simulation {
		fmt.Printf("Skipping dotfile generation. Too many Nodes: %d\n", p1.GetVisitedNodesCount())
	}
	return false
}

func setupSignalHandler(holder *atomic.Pointer[modelchecker.Processor], stopped *bool) {
	c := make(chan os.Signal)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		fmt.Println("\nInterrupted. Stopping state exploration")
		*stopped = true
		p1 := holder.Load()
		if p1 != nil {
			p1.Stop()
		}
	}()
}

func applyDefaultStateOptions(stateConfig *ast.StateSpaceOptions) {
	if stateConfig.Options.MaxActions == 0 {
		stateConfig.Options.MaxActions = 100
	}
	if stateConfig.Options.MaxConcurrentActions == 0 {
		stateConfig.Options.MaxConcurrentActions = min(2, stateConfig.Options.MaxActions)
	}
	if stateConfig.DeadlockDetection == nil {
		deadlockDetection := true
		stateConfig.DeadlockDetection = &deadlockDetection
	}
	if stateConfig.Options.CrashOnYield == nil {
		crashOnYield := true
		stateConfig.Options.CrashOnYield = &crashOnYield
	}
}

func loadStateOptions(dirPath string, f *ast.FrontMatter) *ast.StateSpaceOptions {
	configFileName := filepath.Join(dirPath, "fizz.yaml")
	fmt.Println("configFileName:", configFileName)
	stateConfig, err := modelchecker.ReadOptionsFromYaml(configFileName)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			if isPlayground {
				deadlockDetection := true
				crashOnYield := true
				stateConfig = &ast.StateSpaceOptions{
					Options:           &ast.Options{MaxActions: 100, MaxConcurrentActions: 2, CrashOnYield: &crashOnYield},
					Liveness:          "strict",
					DeadlockDetection: &deadlockDetection,
				}
			} else {
				fmt.Println("fizz.yaml not found. Using default options")
				stateConfig = &ast.StateSpaceOptions{Options: &ast.Options{MaxActions: 100, MaxConcurrentActions: 2}}
			}
		} else {
			fmt.Println("Error reading fizz.yaml:", err)
			os.Exit(1)
		}

	}
	if f.GetYaml() != "" {
		fmStateConfig, err := modelchecker.ReadOptionsFromYamlString(f.GetYaml())
		if err != nil {
			fmt.Println("Error parsing YAML frontmatter:", err)
			os.Exit(1)
		}
		proto.Merge(stateConfig, fmStateConfig)
	}
	return stateConfig
}

func getStateConfigForTraceChecking(stateConfig *ast.StateSpaceOptions) *ast.StateSpaceOptions {
	// Trace-specific config: ignore frontmatter and yaml, use defaults for trace mode
	deadlockDetection := false
	crashOnYield := false
	stateConfig = &ast.StateSpaceOptions{
		Options:           &ast.Options{MaxActions: 100, MaxConcurrentActions: 1, CrashOnYield: &crashOnYield},
		Liveness:          "false",
		DeadlockDetection: &deadlockDetection,
	}
	return stateConfig
}

func loadInputJSON(jsonFilename string) *ast.File {
	// Read the content of the JSON file
	jsonContent, err := os.ReadFile(jsonFilename)
	if err != nil {
		fmt.Println("Error reading JSON file:", err)
		os.Exit(1)
	}
	f := &ast.File{}
	err = protojson.Unmarshal(jsonContent, f)
	if err != nil {
		fmt.Println("Error unmarshalling JSON:", err)
		os.Exit(1)
	}
	return f
}

func parseFlags() []string {
	flag.BoolVar(&isPlayground, "playground", false, "is for playground")
	flag.BoolVar(&simulation, "simulation", false, "Runs in simulation mode (DFS). Default=false for no simulation (BFS)")
	flag.BoolVar(&internalProfile, "internal_profile", false, "Enables CPU and memory profiling of the model checker")
	flag.BoolVar(&saveStates, "save_states", false, "Save states to disk")
	flag.BoolVar(&copyAst, "copy-ast", false, "Copy the AST JSON file to the output directory")
	flag.StringVar(&outputDir, "output-dir", "", "Custom output directory path (if not specified, uses default SPEC_DIR/out/run_timestamp)")
	flag.Int64Var(&seed, "seed", 0, "Seed for random number generator used in simulation mode")
	flag.IntVar(&maxRuns, "max_runs", 0, "Maximum number of simulation runs/paths to explore. Default=0 for unlimited")
	flag.IntVar(&simFirstTraces, "simulation_first_traces", 0, "In simulation mode, write dot files for the first N runs (default 0). The last run is always written.")
	flag.StringVar(&explorationStrategy, "exploration_strategy", "bfs", "Exploration strategy for exhaustive model checking. Options: bfs (default), dfs, random.")
	flag.StringVar(&traceFile, "trace-file", "", "Path to trace file for guided execution")
	flag.StringVar(&preinitHookFile, "preinit-hook-file", "", "Path to Starlark file to execute after preinit but before freezing globals")
	flag.StringVar(&trace, "trace", "", "Trace content as a string for guided execution (multiline supported)")
	flag.IntVar(&traceExtend, "trace-extend", 0, "After trace ends, explore this many additional action depths (0 = stop at trace end)")
	flag.StringVar(&preinitHook, "preinit-hook", "", "Starlark code as a string to execute after preinit but before freezing globals (multiline supported)")
	flag.BoolVar(&isTest, "test", false, "Testing mode (prevents printing timestamps and other non-deterministic behavior. Default=false")
	flag.BoolVar(&experimentalProcessedQueue, "experimental_processed_queue", false, "EXPERIMENTAL: queue holds processed (yield-point) nodes instead of unprocessed action-starts. Dedupes successors before they enter the queue, reducing peak queue memory (~10x on BFS, ~3x on DFS). State space and assertion outcomes are identical under BFS. Under DFS/Random combined with max_actions, the changed exploration order may prune a different subset of states than the default path within the bound — raise max_actions above the spec's natural diameter to avoid the divergence. Default=false.")
	flag.BoolVar(&experimentalNoGraph, "experimental_no_graph", false, "EXPERIMENTAL: drops the in-memory state graph after processing each yield-point. Keeps symmetry reduction, dedup, unique-state count, and safety/transition assertions. Does NOT support liveness assertions — refuses to run if any are present. Auto-enables --experimental_processed_queue. Massive RSS reduction at the cost of trace replayability (only a lightweight action-name chain is kept for failure reporting). Default=false.")
	flag.BoolVar(&experimentalNoStateReturns, "experimental_no_state_returns", false, "EXPERIMENTAL: omit Process.Returns from HashCode, JSON state, and dot-file state labels. Returns remain available on Link.Returns (the per-transition field). Use this to verify downstream consumers (graph, MBT, explorer) have migrated to reading return values from links. Planned to become the default. Default=false.")
	flag.BoolVar(&noSymmetryReduction, "no-symmetry-reduction", false, "Disable symmetry reduction: dedup uses only the plain state hash, so every persisted state keeps its concrete symmetric values/role identities (no canonical renaming). Larger state space. Required when generating state graphs for MBT replay from specs that use symmetric roles or symmetry values. Default=false.")
	flag.Parse()

	// Validate that both file and string versions are not provided
	if traceFile != "" && trace != "" {
		fmt.Println("Error: cannot specify both --trace-file and --trace")
		os.Exit(1)
	}
	if preinitHookFile != "" && preinitHook != "" {
		fmt.Println("Error: cannot specify both --preinit-hook-file and --preinit-hook")
		os.Exit(1)
	}

	args := flag.Args()
	return args
}

func startModelChecker(p1 *modelchecker.Processor) (*modelchecker.Node, *modelchecker.Node, time.Time, error) {
	if simulation {
		rootNode, failedNode, _ := p1.Start()
		return rootNode, failedNode, time.Now(), nil
	}
	if internalProfile {
		startCpuProfile()
		defer pprof.StopCPUProfile()
	}
	startTime := time.Now()
	rootNode, failedNode, err := p1.Start()
	endTime := time.Now()
	if !isTest {
		fmt.Printf("Time taken for model checking: %v\n", endTime.Sub(startTime))
	}
	if internalProfile {
		startHeapProfile()
	}
	return rootNode, failedNode, endTime, err
}

func startCpuProfile() {
	// Start CPU profiling
	f, err := os.Create("cpu.pprof")
	if err != nil {
		panic(err)
	}
	err = pprof.StartCPUProfile(f)
	if err != nil {
		panic(err)
	}
}

func startHeapProfile() {
	f, err := os.Create("mem.pprof")
	if err != nil {
		panic(err)
	}
	err = pprof.WriteHeapProfile(f)
	if err != nil {
		panic(err)
	}
}

func dumpFailedNode(srcFileName string, failedNode *modelchecker.Node, rootNode *modelchecker.Node, outDir string) {

	failurePath := modelchecker.ExtractFailurePath(failedNode, rootNode)
	GenerateFailurePath(srcFileName, failurePath, nil, outDir)
	_, _, err := modelchecker.GenerateErrorPathProtoOfJson(failurePath, outDir+"/")
	if err != nil {
		fmt.Println("Error writing files", err)
	}
}

func GenerateComposedFailurePath(failed *ComposedNode, composition *ast.Composition, roots []*modelchecker.Node, outDir string) {
	// Collect failure paths from each composing spec
	componentPaths := make([][]*modelchecker.Link, len(failed.Nodes))
	subgraphs := make([]string, len(failed.Nodes))
	lastNodeIDs := make([]string, len(failed.Nodes))

	builder := strings.Builder{}
	builder.WriteString("digraph G {\n")

	for i, node := range failed.Nodes {
		specName := composition.GetSpecs()[i].GetName()
		componentPaths[i] = modelchecker.ExtractFailurePath(node, roots[i])
		subgraph := modelchecker.GenerateFailurePath(componentPaths[i], specName, nil)
		subgraphs[i] = subgraph
		builder.WriteString(subgraph)

		// ID of the last node in this path
		lastNodeIDs[i] = fmt.Sprintf("\"%s_%d\"", specName, len(componentPaths[i])-1)
	}

	// Add the merged virtual node
	mergedID := "\"merged\""
	builder.WriteString(fmt.Sprintf("  %s [label=\"merged\", shape=doubleoctagon, style=filled, fillcolor=lightgray];\n", mergedID))
	fmt.Println(failed)
	// Add edges from each last node to merged node
	for _, id := range lastNodeIDs {
		builder.WriteString(fmt.Sprintf("  %s -> %s [arrowhead=none, style=dashed, color=gray];\n", id, mergedID))
	}

	builder.WriteString("}\n")
	dotStr := builder.String()
	err := writeErrorDotFile(outDir, dotStr)
	if err != nil {
		return
	}
}

func dumpComposedFailedNode(failed *ComposedNode, composition *ast.Composition, roots []*modelchecker.Node, outDir string) {
	GenerateComposedFailurePath(failed, composition, roots, outDir)
}

func dumpComposedFailedTransition(failed *ComposedTransition, composition *ast.Composition, roots []*modelchecker.Node, outDir string) {
	GenerateTransitionFailurePath(failed.From, failed.To, composition, roots, outDir)
}

func GenerateTransitionFailurePath(from, to *ComposedNode, composition *ast.Composition, roots []*modelchecker.Node, outDir string) {
	lastFromNodeIDs := make([]string, len(from.Nodes))
	lastToNodeIDs := make([]string, len(from.Nodes))

	builder := strings.Builder{}
	builder.WriteString("digraph G {\n")

	for i := range from.Nodes {
		specName := composition.GetSpecs()[i].GetName()

		// Path from root to `from`
		path := modelchecker.ExtractFailurePath(from.Nodes[i], roots[i])
		lastFromIndex := len(path) - 1

		// Handle transition to `to`
		if from.Nodes[i] == to.Nodes[i] {
			// Add stutter node manually
			lastFromNodeIDs[i] = fmt.Sprintf("\"%s_%d\"", specName, lastFromIndex)
			lastToNodeIDs[i] = lastFromNodeIDs[i]
			builder.WriteString(modelchecker.GenerateFailurePath(path, specName, nil))
			builder.WriteString(fmt.Sprintf("  %s -> %s [label=\"stutter\", color=\"red\", style=dashed];\n", lastFromNodeIDs[i], lastToNodeIDs[i]))
			continue
		}

		// Extract transition path and append
		transition := extractTransitionPath(from.Nodes[i], to.Nodes[i])
		for _, link := range transition {
			link.FailedInvariants = map[int][]int{0: {0}} // Clear failed invariants for transition links
		}
		path = append(path, transition...)

		// Generate full path subgraph
		builder.WriteString(modelchecker.GenerateFailurePath(path, specName, nil))

		// Set FROM node id (before transition path)
		lastFromNodeIDs[i] = fmt.Sprintf("\"%s_%d\"", specName, lastFromIndex)

		// Set TO node id (end of entire path)
		lastToIndex := len(path) - 1
		lastToNodeIDs[i] = fmt.Sprintf("\"%s_%d\"", specName, lastToIndex)
	}

	// Virtual merge node for FROM
	mergedFromID := "\"composite_before\""
	addCompositeNode(builder, mergedFromID, lastFromNodeIDs)

	// Virtual merge node for TO
	mergedToID := "\"composite_after\""
	addCompositeNode(builder, mergedToID, lastToNodeIDs)
	builder.WriteString(fmt.Sprintf("  %s -> %s [style=dashed, color=red];\n", mergedFromID, mergedToID))
	builder.WriteString("}\n")
	dotStr := builder.String()
	err := writeErrorDotFile(outDir, dotStr)
	if err != nil {
		return
	}
}

func addCompositeNode(builder strings.Builder, mergedFromID string, lastFromNodeIDs []string) {
	builder.WriteString(fmt.Sprintf("  %s [label=\"Composite before\", shape=doubleoctagon, style=filled, fillcolor=orange];\n", mergedFromID))
	for _, id := range lastFromNodeIDs {
		builder.WriteString(fmt.Sprintf("  %s -> %s [arrowhead=none, style=dashed, color=gray];\n", id, mergedFromID))
	}
}

func extractTransitionPath(from *modelchecker.Node, to *modelchecker.Node) []*modelchecker.Link {
	type pathEntry struct {
		node *modelchecker.Node
		path []*modelchecker.Link
	}

	visited := make(map[*modelchecker.Node]bool)
	queue := []pathEntry{{node: from, path: []*modelchecker.Link{}}}

	for len(queue) > 0 {
		current := queue[0]
		queue = queue[1:]

		// Prevent revisiting
		if visited[current.node] {
			continue
		}
		visited[current.node] = true

		for _, link := range current.node.Outbound {
			nextNode := link.Node

			// Skip self-loops
			if nextNode == current.node {
				continue
			}

			newPath := append(append([]*modelchecker.Link{}, current.path...), link)

			// If this is the destination node
			if nextNode == to {
				return newPath
			}

			// If it's not a yield node, we can go deeper
			if !nextNode.IsYieldNode() {
				queue = append(queue, pathEntry{node: nextNode, path: newPath})
			} else {
				// If it's a yield node but not the target, stop — we don't want to overshoot
				continue
			}
		}
	}

	// No path found
	return []*modelchecker.Link{}
}

func GenerateFailurePath(srcFileName string, failurePath []*modelchecker.Link, invariant *modelchecker.InvariantPosition, outDir string) {
	for _, link := range failurePath {
		node := link.Node
		stepName := link.Name

		fmt.Printf("------\n%s\n", stepName)

		nodeStr := node.Heap.ToJson()
		nodeStr = strings.ReplaceAll(nodeStr, lib.SymmetryPrefix, "")
		fmt.Printf("--\nstate: %s\n", nodeStr)
		if len(node.Returns) > 0 {
			fmt.Printf("returns: %s\n", strings.ReplaceAll(node.Returns.String(), lib.SymmetryPrefix, ""))
		}
		if len(node.Roles) > 0 {
			for _, role := range node.Roles {
				roleStr := strings.ReplaceAll(role.Fields.String(), lib.SymmetryPrefix, "")
				fmt.Printf("%s: %s\n", role.RefStringShort(), roleStr)
			}
		}
	}
	fmt.Println("------")
	if !isPlayground {
		errJsonFileName := filepath.Join(outDir, "error-graph.json")
		bytes, err := json.MarshalIndent(failurePath, "", "  ")
		if err != nil {
			fmt.Println("Error creating json:", err)
		}
		err = os.WriteFile(errJsonFileName, bytes, 0644)
		if err != nil {
			fmt.Println("Error writing to file:", err)
			return
		}
		fmt.Printf("Writen graph json: %s\n", errJsonFileName)
	}

	dotStr := modelchecker.GenerateFailurePath(failurePath, "", invariant)
	err := writeErrorDotFile(outDir, dotStr)
	if err != nil {
		return
	}
	err = modelchecker.GenerateFailurePathHtml(srcFileName, failurePath, invariant, outDir)
	if err != nil {
		return
	}
	if !isPlayground {
		fmt.Printf("Writen error states as html: %s/error-states.html\nTo open: \n"+
			"open %s/error-states.html\n", outDir, outDir)
	}
}

func writeErrorDotFile(outDir string, dotStr string) error {
	dotFileName := filepath.Join(outDir, "error-graph.dot")
	// Write the content to the file
	err := os.WriteFile(dotFileName, []byte(dotStr), 0644)
	if err != nil {
		fmt.Println("Error writing to file:", err)
		return err
	}
	if !isPlayground {
		fmt.Printf("Writen graph dotfile: %s\nTo generate an image file, run: \n"+
			"dot -Tsvg %s -o error-graph.svg && open error-graph.svg\n", dotFileName, dotFileName)
	}
	return nil
}

func createOutputDir(dirPath string, testing bool) (string, error) {
	// If custom output directory is specified, use it directly
	if outputDir != "" {
		// Convert to absolute path if it's relative
		absOutputDir, err := filepath.Abs(outputDir)
		if err != nil {
			fmt.Println("Error resolving output directory path:", err)
			return "", err
		}

		// Create the directory
		if err := os.MkdirAll(absOutputDir, 0755); err != nil {
			fmt.Println("Error creating output directory:", err)
			return absOutputDir, err
		}

		return absOutputDir, nil
	}

	// Default behavior: create timestamped directory
	var newDirName string
	if testing {
		newDirName = "run_test"
	} else {
		// Create the directory name with current date and time
		dateTimeStr := time.Now().Format("2006-01-02_15-04-05") // Format: YYYY-MM-DD_HH-MM-SS
		newDirName = fmt.Sprintf("run_%s", dateTimeStr)
	}

	// Create the full path for the new directory
	newDirPath := filepath.Join(dirPath, "out", newDirName)

	// Create the directory
	if err := os.MkdirAll(newDirPath, 0755); err != nil {
		fmt.Println("Error creating directory:", err)
		return newDirPath, err
	}

	// Define the symlink path
	latestSymlinkPath := filepath.Join(dirPath, "out", "latest")

	// Remove the existing symlink if it exists
	if _, err := os.Lstat(latestSymlinkPath); err == nil {
		if err := os.Remove(latestSymlinkPath); err != nil {
			fmt.Println("Error removing existing symlink:", err)
			return newDirPath, err
		}
	}
	// Convert to absolute path
	absNewDirPath, err := filepath.Abs(newDirPath)
	if err != nil {
		fmt.Println("Error resolving absolute path:", err)
		return "", err
	}
	// Create the new symlink
	if err := os.Symlink(absNewDirPath, latestSymlinkPath); err != nil {
		fmt.Println("Error creating symlink:", err)
		return newDirPath, err
	}
	// Still returning the newDirPath instead of the symlink path
	// So, all the output logs will still point to the newDirPath.
	// This reduces issues when multiple executions are run in parallel.
	return newDirPath, nil
}