package main

import (
	"fmt"
	"log"

	pf "github.com/The-Pocket/PocketFlow-Go" // Adjust import path
)

// Define node logic using PocketFlow's functional style

// myStartNode creates a node that starts the workflow.
func myStartNode() pf.BaseNode {
	return pf.NewNode().
		SetExec(func(prepResult any, params pf.SharedContext) (any, error) {
			log.Println("Starting workflow...")
			// Exec result can be used by Post to determine action
			return "started_data", nil
		}).
		SetPost(func(ctx pf.SharedContext, prepResult any, execResult any, params pf.SharedContext) (string, error) {
			// Use execResult to decide the next step
			log.Printf("Start node finished with data: %v\n", execResult)
			ctx["start_result"] = execResult // Optional: Update shared context
			return "started", nil            // Action name to trigger the next node
		})
}

// myEndNode creates a node that ends the workflow.
func myEndNode() pf.BaseNode {
	return pf.NewNode().
		SetPrep(func(ctx pf.SharedContext, params pf.SharedContext) (any, error) {
			// Prep can access the shared context
			startData := ctx["start_result"]
			prepMsg := fmt.Sprintf("Preparing to end workflow, received: %v", startData)
			log.Println(prepMsg)
			return prepMsg, nil // Prep result passed to Exec
		}).
		SetExec(func(prepResult any, params pf.SharedContext) (any, error) {
			prepMsg := prepResult.(string) // Assume prep result is string
			log.Printf("Ending workflow with: %s\n", prepMsg)
			// End nodes often don't need to return data
			return nil, nil
		})
	// Default Post (returns DefaultAction) is fine here
}

func main() {
	// Create instances of your nodes
	startNode := myStartNode()
	endNode := myEndNode()

	// Connect the nodes: start -> end (when action is "started")
	startNode.Next("started", endNode)

	// Create a flow with the start node
	flow := pf.NewFlow(startNode)

	// Create a context and run the flow
	context := make(pf.SharedContext)
	log.Println("Executing workflow...")
	finalAction, err := flow.Run(context)
	if err != nil {
		log.Fatalf("Workflow failed: %v\n", err)
	}

	log.Printf("Workflow completed successfully. Final action: %s\n", finalAction)
	log.Printf("Final Context: %v\n", context)
}