package main

import (
	"context"
	"fmt"

	"github.com/jefflinse/rhizome"
)

func main() {
	g := rhizome.New[int]()

	g.AddNode("double", func(_ context.Context, n int) (int, error) {
		return n * 2, nil
	})
	g.AddNode("add-ten", func(_ context.Context, n int) (int, error) {
		return n + 10, nil
	})

	g.AddEdge(rhizome.Start, "double")
	g.AddEdge("double", "add-ten")
	g.AddEdge("add-ten", rhizome.End)

	compiled, _ := g.Compile()
	result, _ := compiled.Run(context.Background(), 5)

	fmt.Println(result) // 20
}

type state struct {
	Value  int
	Status string
}

g := rhizome.New[*state]()

g.AddNode("classify", func(_ context.Context, s *state) (*state, error) {
	if s.Value >= 100 {
		s.Status = "high"
	} else {
		s.Status = "low"
	}
	return s, nil
})
g.AddNode("handle-high", handleHigh)
g.AddNode("handle-low", handleLow)

g.AddEdge(rhizome.Start, "classify")
g.AddConditionalEdge("classify", func(_ context.Context, s *state) (string, error) {
	if s.Status == "high" {
		return "handle-high", nil
	}
	return "handle-low", nil
}, "handle-high", "handle-low")
g.AddEdge("handle-high", rhizome.End)
g.AddEdge("handle-low", rhizome.End)

compiled, _ := g.Compile()
result, _ := compiled.Run(ctx, &state{Value: 150})

g := rhizome.New[int]()

g.AddNode("increment", func(_ context.Context, n int) (int, error) {
	return n + 1, nil
})

g.AddEdge(rhizome.Start, "increment")
g.AddConditionalEdge("increment", func(_ context.Context, n int) (string, error) {
	if n >= 5 {
		return rhizome.End, nil
	}
	return "increment", nil
}, "increment", rhizome.End)

compiled, _ := g.Compile(rhizome.WithMaxNodeExecs(20)) // compile-time default
result, _ := compiled.Run(context.Background(), 0)     // result: 5

// Per-Run override:
result, _ = compiled.Run(context.Background(), 0, rhizome.WithRunMaxNodeExecs[int](50))

logger := func(ctx context.Context, node string, state int, next rhizome.NodeFunc[int]) (int, error) {
	fmt.Printf("entering %s\n", node)
	result, err := next(ctx, state)
	fmt.Printf("leaving %s\n", node)
	return result, err
}

result, _ := compiled.Run(ctx, 0, rhizome.WithMiddleware(logger))

result, err := compiled.Run(ctx, initial, rhizome.WithMiddleware(
    rhizome.Recover[*State](),               // trap panics, return ErrNodePanic
    rhizome.Retry[*State](                   // retry transient failures
        rhizome.WithMaxAttempts(3),
    ),
    rhizome.Timeout[*State](30*time.Second), // per-attempt deadline
))

type MyState struct {
	Step int
	Logs []string
}

func (s *MyState) MarshalBinary() ([]byte, error)    { return json.Marshal(s) }
func (s *MyState) UnmarshalBinary(data []byte) error { return json.Unmarshal(data, s) }

store := &rhizome.MemoryStore{} // or any CheckpointStore implementation

g := rhizome.New[*MyState]()
// ... AddNode/AddEdge ...

compiled, err := g.Compile(rhizome.WithCheckpointing(store))
if err != nil {
	// Returns ErrCheckpointRequiresSnapshotter if *MyState does not satisfy Snapshotter.
	panic(err)
}

// Thread IDs correlate runs with checkpoints; required when checkpointing is enabled.
final, err := compiled.Run(ctx, &MyState{}, rhizome.WithThreadID[*MyState]("conversation-123"))

// If the run was interrupted, resume from the last checkpoint in a fresh state instance:
resumed, err := compiled.Resume(ctx, "conversation-123", &MyState{})

g.AddNode("confirm", func(ctx context.Context, s *State) (*State, error) {
    resp, err := rhizome.Interrupt(ctx, rhizome.InterruptRequest{
        Kind:    "approve",
        Payload: s.Proposal,
    })
    if err != nil {
        return s, err
    }
    s.Approved = resp.Value.(bool)
    return s, nil
})

handler := func(ctx context.Context, req rhizome.InterruptRequest) (rhizome.InterruptResponse, error) {
    // Blocking is expected here — the graph waits.
    // Any blocking call should select on ctx.Done() to honor cancellation.
    approved := promptUser(req.Payload)
    return rhizome.InterruptResponse{Value: approved}, nil
}

final, err := compiled.Run(ctx, &State{},
    rhizome.WithInterruptHandler[*State](handler))