package loop

import (
	"context"
	"fmt"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// Run drives the agent loop: inject relevant memory into the first-turn context,
// then run up to MaxRounds steps. Each step runs the model (streaming through
// OnStream when configured) and any requested tool calls are executed under a
// per-invoke timeout and fed back as RoleTool messages so the model can continue.
// A checkpoint is persisted at every step, every step is traced as a
// state_transition span (with an llm_call child and tool_call spans for each
// invocation), and after each step the memory Extractor runs with errors logged
// rather than fatal.
//
// It returns the final answer, or a partial Result with StopReasonMaxRounds when
// the budget is exhausted without a final answer. Step errors are wrapped with %w
// so callers can classify them with errors.Is.
func (l *Loop) Run(ctx context.Context, sessionID string, messages []model.Message) (*Result, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if l.agent == nil {
		return nil, fmt.Errorf("loop: run: %w", ErrNilAgent)
	}

	msgs := l.injectMemory(ctx, messages)
	res := &Result{}
	for step := 1; step <= l.maxRounds; step++ {
		if err := ctx.Err(); err != nil {
			return nil, fmt.Errorf("loop: run: step %d: %w", step, err)
		}

		rc, roundSpan := l.span(ctx, "loop.step", trace.SpanKindStateTransition)
		outcome, err := l.step(rc, sessionID, msgs)
		if err != nil {
			l.spanProp(roundSpan, "error", err.Error())
			l.finishSpan(roundSpan, trace.SpanStatusError)
			return nil, fmt.Errorf("loop: run: step %d: %w", step, err)
		}
		l.spanProp(roundSpan, "step", step)
		l.spanProp(roundSpan, "stop_reason", outcome.stopReason)
		l.finishSpan(roundSpan, trace.SpanStatusOK)

		res.Steps = step
		if outcome.final {
			res.Text = outcome.text
			res.StopReason = outcome.stopReason
			l.saveCheckpoint(ctx, sessionID, step, msgs, outcome.stopReason)
			l.extractMemory(ctx, msgs, step)
			return res, nil
		}

		newMsgs, executed, err := l.runToolCalls(rc, msgs, outcome.toolCalls)
		if err != nil {
			return nil, fmt.Errorf("loop: run: step %d: tools: %w", step, err)
		}
		msgs = newMsgs
		res.ToolCallsExecuted += executed
		l.saveCheckpoint(ctx, sessionID, step, msgs, agent.StopReasonToolCalls)
		l.extractMemory(ctx, msgs, step)
	}

	// Budget exhausted without a final answer.
	res.Steps = l.maxRounds
	res.StopReason = StopReasonMaxRounds
	l.saveCheckpoint(ctx, sessionID, l.maxRounds, msgs, StopReasonMaxRounds)
	l.extractMemory(ctx, msgs, l.maxRounds)
	return res, nil
}