engine.go

package flowengine

import (
	"context"
	"errors"
	"fmt"
	"strings"
)

// ===========================================================================
// Engine - FlowEngine Execution Engine
// ===========================================================================

// Engine is the flowengine execution engine.
type Engine struct {
	adapter        Adapter
	sessions       *SessionManager
	defaultRetries int
}

// EngineOption configures Engine.
type EngineOption func(*Engine)

// WithDefaultRetries sets the default retry count.
func WithDefaultRetries(n int) EngineOption {
	return func(e *Engine) {
		e.defaultRetries = n
	}
}

// New creates a new Engine with the given adapter.
func New(adapter Adapter, opts ...EngineOption) *Engine {
	e := &Engine{
		adapter:        adapter,
		sessions:       NewSessionManager(),
		defaultRetries: 2,
	}
	for _, opt := range opts {
		opt(e)
	}
	return e
}

// Run executes a Runnable.
func (e *Engine) Run(ctx context.Context, r Runnable, metadata Metadata) (*FlowContext, *Response, error) {
	fc := NewFlowContext(ctx, e, metadata)
	resp, err := r.Run(fc)
	return fc, resp, err
}

// RunWithSession executes a Runnable within a session context.
func (e *Engine) RunWithSession(ctx context.Context, sessionID string, r Runnable, metadata Metadata) (*FlowContext, *Response, error) {
	fc := NewFlowContext(ctx, e, metadata)
	session := e.sessions.Get(sessionID)
	fc.WithSession(session)
	resp, err := r.Run(fc)
	return fc, resp, err
}

// DeleteSession deletes a session.
func (e *Engine) DeleteSession(ctx context.Context, sessionID string) error {
	session := e.sessions.GetIfExists(sessionID)
	if session == nil {
		return nil
	}

	if sa, ok := e.adapter.(SessionAdapter); ok {
		if rootID := session.RootResponseID(); rootID != "" {
			_ = sa.DeleteSession(ctx, rootID)
		}
	}

	e.sessions.Delete(sessionID)
	return nil
}

// Stream executes a step and returns a streaming response channel.
func (e *Engine) Stream(ctx context.Context, step *Step, metadata Metadata) (<-chan string, error) {
	if step == nil {
		return nil, errors.New("flowengine: step is nil")
	}
	if step.Actions != nil && !step.Actions.IsEmpty() {
		return nil, errors.New("flowengine: stream only supports steps without actions")
	}

	fc := NewFlowContext(ctx, e, metadata)
	system, err := step.BuildSystemPrompt(fc)
	if err != nil {
		return nil, err
	}
	user, err := step.BuildUserPrompt(fc, "")
	if err != nil {
		return nil, err
	}

	return e.adapter.Stream(ctx, system, user)
}

// ===========================================================================
// Step Execution
// ===========================================================================

func (e *Engine) executeStep(fc *FlowContext, step *Step) (*Response, error) {
	maxRetries := step.MaxRetries
	if maxRetries == 0 {
		maxRetries = e.defaultRetries
	}

	tools := step.ToolDefs()

	if len(tools) > 0 {
		return e.executeWithToolsRetry(fc, step, maxRetries)
	}

	return e.executeTextOnly(fc, step)
}

func (e *Engine) executeWithToolsRetry(fc *FlowContext, step *Step, maxRetries int) (*Response, error) {
	toolAdapter, ok := e.adapter.(ToolAdapter)
	if !ok {
		return nil, errors.New("flowengine: adapter does not support tools")
	}

	maxAttempts := maxRetries + 1
	var lastErr error
	var lastResp *Response
	var hint string

	for attempt := 1; attempt <= maxAttempts; attempt++ {
		select {
		case <-fc.Done():
			return nil, fc.Err()
		default:
		}

		system, err := step.BuildSystemPrompt(fc)
		if err != nil {
			return nil, err
		}
		user, err := step.BuildUserPrompt(fc, hint)
		if err != nil {
			return nil, err
		}

		tools := step.ToolDefs()
		session := fc.Session()

		var call *ToolCall
		var responseID string

		if session != nil {
			previousID := session.LastResponseID()
			call, responseID, err = toolAdapter.CallWithToolsAndSession(fc, system, user, previousID, tools)
			if err != nil {
				lastErr = err
				hint = e.buildErrorHint(err, step)
				continue
			}
		} else {
			call, err = toolAdapter.CallWithTools(fc, system, user, tools)
			if err != nil {
				lastErr = err
				hint = e.buildErrorHint(err, step)
				continue
			}
		}

		if call == nil {
			lastErr = errors.New("flowengine: tool call is nil")
			hint = e.buildErrorHint(lastErr, step)
			continue
		}

		resp := &Response{
			Type:   ResponseMethod,
			Action: call.Name,
			Params: call.Arguments,
		}
		lastResp = resp

		if !step.ValidateAction(call.Name) {
			lastErr = fmt.Errorf("unknown action: %s", call.Name)
			hint = fmt.Sprintf("Action '%s' is invalid. Available actions: %s", call.Name, strings.Join(step.ActionNames(), ", "))
			continue
		}

		if step.Validator != nil {
			if validErr := step.Validator(resp); validErr != nil {
				hint = fmt.Sprintf("Validation failed: %v\nPlease fix and retry.", validErr)
				lastErr = &ValidationError{
					Err:      validErr,
					Response: resp,
					Attempt:  attempt,
				}
				continue
			}
		}

		if session != nil {
			session.RecordNode(step.Name, responseID)
		}

		return resp, nil
	}

	if lastResp != nil {
		return lastResp, &ValidationError{
			Err:      lastErr,
			Response: lastResp,
			Attempt:  maxAttempts,
		}
	}

	return nil, fmt.Errorf("flowengine: all %d attempts failed: %w", maxAttempts, lastErr)
}

func (e *Engine) executeTextOnly(fc *FlowContext, step *Step) (*Response, error) {
	system, err := step.BuildSystemPrompt(fc)
	if err != nil {
		return nil, err
	}
	user, err := step.BuildUserPrompt(fc, "")
	if err != nil {
		return nil, err
	}

	session := fc.Session()
	sessionAdapter, isSessionAdapter := e.adapter.(SessionAdapter)

	var raw string
	var responseID string

	if session != nil && isSessionAdapter {
		previousID := session.LastResponseID()
		raw, responseID, err = sessionAdapter.CallWithSession(fc, system, user, previousID)
		if err != nil {
			return nil, err
		}
		session.RecordNode(step.Name, responseID)
		session.AppendMessage("user", user)
		session.AppendMessage("assistant", raw)
	} else {
		raw, err = e.adapter.Call(fc, system, user)
		if err != nil {
			return nil, err
		}
	}

	return &Response{
		Type:    ResponseMessage,
		Message: raw,
		Raw:     raw,
	}, nil
}

func (e *Engine) buildErrorHint(err error, step *Step) string {
	var sb strings.Builder
	sb.WriteString("Previous call failed: ")
	sb.WriteString(err.Error())
	sb.WriteString("\n\n")
	sb.WriteString("Available actions: ")
	sb.WriteString(strings.Join(step.ActionNames(), ", "))
	sb.WriteString("\nPlease select a valid action and retry.")
	return sb.String()
}

// Adapter returns the underlying adapter.
func (e *Engine) Adapter() Adapter {
	return e.adapter
}

// Sessions returns the session manager.
func (e *Engine) Sessions() *SessionManager {
	return e.sessions
}

// Close stops background session cleanup.
func (e *Engine) Close() {
	if e.sessions != nil {
		e.sessions.Stop()
	}
}