package loop

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"sync"

	"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/tool"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// ToolExecutionMode tokens accepted by Options.ToolExecution. Empty selects
// toolExecParallel.
const (
	toolExecParallel   = "parallel"
	toolExecSequential = "sequential"
)

// criticalToolLoopTerminalReason is surfaced as Result.Text and on every blocked
// call when the critical-tool-loop safety net stops the run.
const criticalToolLoopTerminalReason = "critical tool-loop detected: an identical tool call was repeated across rounds; the run is stopping"

// ToolCallVeto is returned by a CallAdmit hook to block a single call.
type ToolCallVeto struct {
	Block  bool
	Reason string
}

// ToolOutcome is the normalized result of executing one tool call. It is what an
// after-call / after-outcome hook observes and may rewrite before the loop
// renders it back to the model.
type ToolOutcome struct {
	// Result is the value the tool returned (nil on failure or skip).
	Result any
	// Err is non-nil when the call failed to execute; it is nil on success.
	Err error
	// ErrKind classifies a failure as tool.KindArgumentValidation or
	// tool.KindExecution (the task's errorKind status classification). Empty on
	// success.
	ErrKind string
	// Skipped is true when the call never executed (duplicate / blocked / vetoed).
	Skipped bool
	// Reason is a human-readable explanation for a skipped call.
	Reason string
	// Terminate, when set by a hook, signals that the whole round should stop
	// once every call in the batch is also marked terminate.
	Terminate bool
}

// ToolLoopHooks groups the optional per-batch / per-call tool-loop hooks. Every
// field is nil-tolerant. These are the extension points the loop keeps open for
// "敏感操作代审" (design/11 batch admission) and the future async-tool path, so
// the audit layer and the loop are decoupled.
type ToolLoopHooks struct {
	// BatchAdmit is called once per batch before any call executes. Returning
	// (false, reason) blocks the whole batch: every call becomes a skipped result
	// carrying reason. nil = always admit.
	BatchAdmit func(ctx context.Context, calls []agent.ToolCall) (bool, string)
	// CallAdmit is called per call before it executes. A non-nil veto with
	// Block=true blocks that single call. nil = always admit.
	CallAdmit func(ctx context.Context, call agent.ToolCall) *ToolCallVeto
	// AfterCall rewrites a call's outcome after the tool executes. A nil return
	// keeps the original. It runs before AfterOutcome.
	AfterCall func(ctx context.Context, call agent.ToolCall, outcome ToolOutcome) *ToolOutcome
	// AfterOutcome observes/rewrites the finalized outcome (after AfterCall and
	// after any skip/block decision) immediately before the result is rendered
	// back to the model. nil = keep the outcome.
	AfterOutcome func(ctx context.Context, call agent.ToolCall, outcome ToolOutcome) *ToolOutcome
}

// toolRunResult is the outcome of running one tool batch.
type toolRunResult struct {
	messages  []model.Message
	executed  int
	terminate bool
}

// toolLoopState carries the per-run state of the critical-tool-loop safety net.
// It is scoped to a single Run call, so concurrent runs never share it.
type toolLoopState struct {
	criticalSeen bool            // a critical loop was already reported this run
	prevFP       map[string]bool // fingerprints seen in the immediately previous batch
}

// toolCallAction is the execution decision for one call after classification.
type toolCallAction struct {
	run    bool   // execute the call
	reason string // when run==false, the human reason fed back to the model
}

// runToolCalls executes a batch of model-requested tool calls. It classifies
// duplicate invocations (the critical-tool-loop safety net), applies the optional
// batch/per-call hooks, runs the surviving calls in parallel (bounded) or
// serially, and feeds every outcome back as a RoleTool message in the original
// call order. It returns the number of calls actually invoked (executed), the
// updated message slice, and whether the round should terminate.
//
// Invocation is delegated to the dispatched tool, which applies the per-tool
// deadline; a missing dispatcher is fed back as an error so the loop never
// panics or crashes on an unhandled call.
func (l *Loop) runToolCalls(ctx context.Context, messages []model.Message, calls []agent.ToolCall, st *toolLoopState) (toolRunResult, error) {
	d := l.dispatcher()
	res := toolRunResult{messages: messages}
	if len(calls) == 0 {
		return res, nil
	}

	actions, terminate := classifyToolBatch(st, calls)
	if terminate {
		res.terminate = true
	}

	// Security gate (design/11 §3): a group/stranger session may invoke only
	// tools whose editable marker allows that environment (default closed). A
	// disallowed call never executes; it is fed back as a not_allowed soft result
	// and recorded to the audit sink. It is the final authority over the batch —
	// a later audit hook may tighten but never re-admit a denied tool.
	if d != nil {
		l.applyToolSecurityGate(ctx, d, calls, actions)
	}

	// BatchAdmit hook (audit / political review): block the whole batch.
	if l.hooks.BatchAdmit != nil {
		if ok, reason := l.hooks.BatchAdmit(ctx, calls); !ok {
			if reason == "" {
				reason = "tool batch blocked by policy"
			}
			for i := range calls {
				actions[i] = toolCallAction{run: false, reason: reason}
			}
		}
	}

	// CallAdmit hook: per-call veto over the calls we intend to run.
	if l.hooks.CallAdmit != nil {
		for i := range calls {
			if !actions[i].run {
				continue
			}
			if veto := l.hooks.CallAdmit(ctx, calls[i]); veto != nil && veto.Block {
				reason := veto.Reason
				if reason == "" {
					reason = "tool call blocked by policy"
				}
				actions[i] = toolCallAction{run: false, reason: reason}
			}
		}
	}

	// Decide serial vs parallel: explicit sequential mode, or any call whose tool
	// is declared sequential (a single sequential marker makes the whole group
	// serial, matching OpenClaw's launchParallelToolCalls).
	sequential := l.toolExec == toolExecSequential
	if !sequential && d != nil {
		for i := range calls {
			if actions[i].run && d.IsSequential(calls[i].Name) {
				sequential = true
				break
			}
		}
	}

	outcomes := make([]ToolOutcome, len(calls))
	if sequential {
		for i := range calls {
			if actions[i].run {
				outcomes[i] = l.invokeOne(ctx, d, calls[i])
				res.executed++
			} else {
				outcomes[i] = skippedOutcome(actions[i].reason)
			}
		}
	} else {
		idxs := make([]int, 0, len(calls))
		for i := range calls {
			if actions[i].run {
				idxs = append(idxs, i)
			} else {
				outcomes[i] = skippedOutcome(actions[i].reason)
			}
		}
		l.invokeParallel(ctx, d, calls, idxs, outcomes)
		res.executed += len(idxs)
	}

	// AfterCall hook (only for calls that actually executed).
	if l.hooks.AfterCall != nil {
		for i := range calls {
			if outcomes[i].Skipped {
				continue
			}
			if after := l.hooks.AfterCall(ctx, calls[i], outcomes[i]); after != nil {
				outcomes[i] = *after
			}
		}
	}

	// AfterOutcome hook (runs on every finalized outcome, skipped or not).
	if l.hooks.AfterOutcome != nil {
		for i := range calls {
			if after := l.hooks.AfterOutcome(ctx, calls[i], outcomes[i]); after != nil {
				outcomes[i] = *after
			}
		}
	}

	// A hook may ask to terminate the round by marking every call terminate
	// (OpenClaw's shouldTerminateToolBatch).
	if !res.terminate && shouldTerminateBatch(outcomes) {
		res.terminate = true
	}

	// Render the assistant tool-call turn (native providers require it so the
	// following tool_result can be correlated) plus every result, in call order so
	// the model sees a stable feed.
	res.messages = appendToolCallHistory(res.messages, calls, outcomes)
	return res, nil
}

// classifyToolBatch partitions a batch into run/skip decisions, enforcing the
// critical-tool-loop safety net:
//
//   - A call identical (same name + arguments) to one earlier in the same batch
//     is skipped with a duplicate hint (never executes).
//   - A call identical to one from a previous round is a critical-loop repeat:
//     the first occurrence blocks the repeated call and sets criticalSeen, so the
//     model gets a turn to reassess; a second consecutive occurrence terminates
//     the whole run.
//
// It returns (actions, terminate). When terminate is true every call is blocked
// with criticalToolLoopTerminalReason and the round must end.
func classifyToolBatch(st *toolLoopState, calls []agent.ToolCall) ([]toolCallAction, bool) {
	seen := make(map[string]bool, len(calls))
	actions := make([]toolCallAction, len(calls))
	repeat := false
	for i, call := range calls {
		fp := toolFingerprint(call)
		if seen[fp] {
			// Duplicate of a call already in this batch.
			if st.prevFP[fp] {
				actions[i] = toolCallAction{run: false, reason: repeatedCallReason()}
				repeat = true
			} else {
				actions[i] = toolCallAction{run: false, reason: duplicateCallReason()}
			}
			continue
		}
		seen[fp] = true
		if st.prevFP[fp] {
			// Repeated from the previous round — a critical loop.
			actions[i] = toolCallAction{run: false, reason: repeatedCallReason()}
			repeat = true
		} else {
			actions[i] = toolCallAction{run: true}
		}
	}

	// Remember which fingerprints were live this round for the next round's
	// cross-round repeat detection.
	st.prevFP = seen

	terminate := false
	if repeat {
		if st.criticalSeen {
			terminate = true
			for i := range actions {
				actions[i] = toolCallAction{run: false, reason: criticalToolLoopTerminalReason}
			}
		} else {
			st.criticalSeen = true
		}
	}
	return actions, terminate
}

// invokeOne executes a single tool call under its per-tool deadline, emitting a
// tool span, streaming tool_execution_update progress and normalizing the
// outcome.
func (l *Loop) invokeOne(ctx context.Context, d *tool.Dispatcher, tc agent.ToolCall) ToolOutcome {
	if d == nil {
		return ToolOutcome{Err: fmt.Errorf("no tool dispatcher configured for %q", tc.Name), ErrKind: tool.KindExecution}
	}
	l.emitToolExecution(ctx, tc.Name, "start")
	_, toolSpan := l.span(ctx, "tool."+tc.Name, trace.SpanKindTool)
	result, err := d.Invoke(ctx, tc.Name, tc.Args)
	phase := "end"
	if err != nil {
		phase = "end:error"
	}
	l.emitToolExecution(ctx, tc.Name, phase)
	out := ToolOutcome{Result: result, Err: err, ErrKind: errorKind(err)}
	if err != nil {
		l.spanProp(toolSpan, "error", err.Error())
		l.spanProp(toolSpan, "error_kind", out.ErrKind)
		l.finishSpan(toolSpan, trace.SpanStatusError)
	} else {
		l.spanProp(toolSpan, "result", result)
		l.finishSpan(toolSpan, trace.SpanStatusOK)
	}
	return out
}

// emitToolExecution forwards a tool_execution_update progress event to the
// frontend when the loop is on the streaming path (OnStream configured); it is a
// no-op otherwise (muted internal passes skip it too). phase is a short label
// like "start" or "end:error".
func (l *Loop) emitToolExecution(ctx context.Context, name, phase string) {
	if l.onStream == nil || isMutedBroadcast(ctx) {
		return
	}
	ev := model.StreamEvent{Kind: model.StreamKindToolExecution, Delta: phase + ":" + name}
	ev.SessionID = loopSession(ctx)
	ev.TurnID = loopTurn(ctx)
	l.onStream(ev)
}

// invokeParallel runs calls at idxs concurrently, capped at the loop's
// parallelism, writing each outcome back into out at its call index so callers
// can render results in the original order. The semaphore bounds concurrency and
// in-flight goroutines release their slot on completion, so a canceled context
// (which makes every Invoke return promptly) cannot leak goroutines.
func (l *Loop) invokeParallel(ctx context.Context, d *tool.Dispatcher, calls []agent.ToolCall, idxs []int, out []ToolOutcome) {
	if len(idxs) == 0 {
		return
	}
	sem := make(chan struct{}, l.toolParallelism)
	var wg sync.WaitGroup
	for _, i := range idxs {
		wg.Add(1)
		sem <- struct{}{}
		go func(i int) {
			defer wg.Done()
			defer func() { <-sem }()
			out[i] = l.invokeOne(ctx, d, calls[i])
		}(i)
	}
	wg.Wait()
}

// returned by classifyToolBatch for the two skip reasons.
func duplicateCallReason() string {
	return "duplicate call with identical arguments already issued in this batch"
}

func repeatedCallReason() string {
	return "already used with identical arguments in a previous round; do not repeat it"
}

// skippedOutcome builds a Skipped outcome carrying the human reason.
func skippedOutcome(reason string) ToolOutcome {
	if reason == "" {
		reason = "tool call skipped"
	}
	return ToolOutcome{Skipped: true, Reason: reason}
}

// shouldTerminateBatch reports whether every call in the batch asked to terminate
// (OpenClaw's shouldTerminateToolBatch), so a hook can stop the round.
func shouldTerminateBatch(outcomes []ToolOutcome) bool {
	if len(outcomes) == 0 {
		return false
	}
	for _, o := range outcomes {
		if !o.Terminate {
			return false
		}
	}
	return true
}

// toolFingerprint canonicalizes a tool call (name + args) so identical calls can
// be detected. Args are JSON-marshaled, so equivalent objects with different key
// orders collapse to the same fingerprint; a marshal failure falls back to %v.
func toolFingerprint(call agent.ToolCall) string {
	b, err := json.Marshal(call.Args)
	if err != nil {
		b = []byte(fmt.Sprintf("%v", call.Args))
	}
	return call.Name + "\x00" + string(b)
}

// errorKind classifies an invocation error into a tool.Kind* label so the model
// can distinguish an argument-validation failure (re-issue a corrected call) from
// an execution failure. A nil error yields "" (success).
func errorKind(err error) string {
	if err == nil {
		return ""
	}
	var te *tool.ToolError
	if errors.As(err, &te) {
		return te.Kind()
	}
	if errors.Is(err, tool.ErrInvalidArgs) {
		return tool.KindArgumentValidation
	}
	return tool.KindExecution
}

// dispatcher returns the loop's tool dispatcher, falling back to the agent's own
// when the loop was not given one.
func (l *Loop) dispatcher() *tool.Dispatcher {
	if l.tools != nil {
		return l.tools
	}
	if l.agent != nil {
		return l.agent.Tools()
	}
	return nil
}

// toolResultMessage renders a tool call's outcome as a RoleTool message ready to
// be fed back to the model. Skipped calls carry their reason; failures carry a
// structured {kind,message} so the LLM can react appropriately; success carries
// the value.
func toolResultMessage(tc agent.ToolCall, out ToolOutcome) model.Message {
	return model.Message{Role: model.RoleTool, ToolCallName: tc.Name,
		Content: []model.Content{{Type: model.ContentTypeText, Text: toolResultText(tc, out)}}}
}

// toolResultText renders the text payload of a tool result message.
func toolResultText(tc agent.ToolCall, out ToolOutcome) string {
	switch {
	case out.Skipped:
		if out.Reason == "" {
			return fmt.Sprintf("tool %s skipped", tc.Name)
		}
		return fmt.Sprintf("tool %s skipped: %s", tc.Name, out.Reason)
	case out.Err != nil:
		if kind := out.ErrKind; kind != "" {
			if payload, err := json.Marshal(map[string]string{"kind": kind, "message": out.Err.Error()}); err == nil {
				return fmt.Sprintf("tool %s failed: %s", tc.Name, payload)
			}
		}
		return fmt.Sprintf("tool %s failed: %v", tc.Name, out.Err)
	default:
		return fmt.Sprintf("tool %s result: %s", tc.Name, formatToolResult(out.Result))
	}
}

// formatToolResult renders a tool return value as text for the model. Strings are
// passed through (an empty string becomes "empty"); everything else is JSON
// encoded, falling back to a plain %v render on an encoding error.
func formatToolResult(result any) string {
	switch v := result.(type) {
	case nil:
		return "null"
	case string:
		if v == "" {
			return "empty"
		}
		return v
	default:
		b, err := json.Marshal(v)
		if err != nil {
			return fmt.Sprintf("%v", v)
		}
		return string(b)
	}
}