open-code-review/internal/agent/agent.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"os/exec"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/open-code-review/open-code-review/internal/config/allowlist"
	"github.com/open-code-review/open-code-review/internal/config/rules"
	"github.com/open-code-review/open-code-review/internal/config/template"
	"github.com/open-code-review/open-code-review/internal/config/toolsconfig"
	"github.com/open-code-review/open-code-review/internal/diff"
	"github.com/open-code-review/open-code-review/internal/llm"
	"github.com/open-code-review/open-code-review/internal/model"
	"github.com/open-code-review/open-code-review/internal/session"
	"github.com/open-code-review/open-code-review/internal/stdout"
	"github.com/open-code-review/open-code-review/internal/telemetry"
	"github.com/open-code-review/open-code-review/internal/tool"
)

// Args holds all dependencies and configuration needed to run a review session.
type Args struct {
	// RepoDir is the root of the git repository.
	RepoDir string

	// From and To define the diff range (e.g., "main..feature-branch").
	From string
	To   string

	// Commit is a single commit hash to review (vs its parent).
	Commit string

	// Template loaded from YAML config file.
	Template template.Template

	// SystemRule holds path-based review rules loaded from a JSON config.
	SystemRule *rules.SystemRule

	// LLM client for model inference.
	LLMClient llm.LLMClient

	// Tool registry mapping tool aliases to implementations.
	Tools tool.Registry

	// PlanToolDefs holds llm.ToolDef entries enabled in plan_task, built once at startup.
	// When nil, plan phase sends no tool definitions (same as Java behavior when plan_task is false).
	PlanToolDefs []llm.ToolDef

	// MainToolDefs holds llm.ToolDef entries enabled in main_task, built once at startup.
	MainToolDefs []llm.ToolDef

	// CommentWorkerPool — separate goroutine pool for running asynchronous
	// comment post-processing tasks (tracking, re-tracking, reflection,
	// suggestion validation). This mirrors the Java side's subtaskExecutor
	// which executes the CODE_COMMENT tool off the critical path so that the
	// main LLM tool-use loop can continue issuing requests while comments are
	// being processed in the background.
	//
	// When nil (the default), comment processing happens synchronously inside
	// executeToolCall instead of via a separate worker pool.
	CommentWorkerPool *CommentWorkerPool

	// Concurrency limit for per-file subtasks. Defaults to number of CPUs.
	MaxConcurrency int

	// Concurrent task timeout in minutes. 0 means no timeout.
	ConcurrentTaskTimeout int

	// CommentCollector collects review comments generated by the code_comment tool.
	CommentCollector *tool.CommentCollector

	// DiffMap is a shared map that gets populated with parsed diff text.
	// When provided, loadDiffs writes into this map so that FileReadDiffProvider
	// can serve diff content to the LLM without duplicating git parsing.
	DiffMap map[string]string

	// Background is an optional requirement/business context string
	// injected into plan and main_task prompts via {{requirement_background}}.
	Background string

	// Model is the user-configured model name used as fallback when
	// template phases (plan/memory_compression) don't specify one.
	Model string

	// Session is an optional session history instance for collecting conversation records.
	// When nil, a default one is created automatically with git branch auto-detected from repoDir.
	Session *session.SessionHistory
}

// AgentWarning describes a non-fatal warning recorded during review.
type AgentWarning struct {
	File    string `json:"file"`
	Message string `json:"message"`
	Type    string `json:"type"`
}

// compression thresholds as fractions of MaxTokens.
const (
	tokenSoftThreshold    = 0.60 // async background compression
	tokenWarningThreshold = 0.80 // immediate sync compression
)

// round groups consecutive messages starting with an assistant message
// followed by zero or more tool result messages.
type round struct {
	assistantIdx int
	toolIdxs     []int
}

// partitionResult describes how messages should be split for compression.
type partitionResult struct {
	frozenEnd   int
	compressEnd int
	rounds      []round
	activeCount int
}

// compressionJob tracks an in-flight background compression operation.
type compressionJob struct {
	done    chan struct{}
	rebuilt []llm.Message
}

// Agent orchestrates the AI-powered code review.
type Agent struct {
	args              Args
	diffs             []model.Diff // parsed diffs
	totalInsertions   int64
	totalDeletions    int64
	currentDate       string
	session           *session.SessionHistory
	totalTokensUsed   int64 // accumulated total tokens from all LLM calls, accessed atomically
	totalInputTokens  int64 // accumulated input/prompt tokens, accessed atomically
	totalOutputTokens int64 // accumulated completion tokens, accessed atomically
	warningsMu        sync.Mutex
	warnings          []AgentWarning
	compressionMu     sync.Mutex
	pendingJob        *compressionJob
}

// CommentWorkerPool manages a fixed-size pool of workers dedicated to
// processing code-review comment post-steps (line-range tracking,
// re-tracking, reflection, suggestion validation) asynchronously.
//
// These steps can be time-consuming (network calls to LLM, external APIs,
// heavy computation). By offloading them to a worker pool the main LLM
// tool-use loop stays unblocked, reducing overall latency - just like the
// Java implementation uses a dedicated subtaskExecutor for the CODE_COMMENT
// tool (see CodeReviewNativeAgent.executeToolCall ~L640-642).
type CommentWorkerPool struct {
	semaphore chan struct{}
	wg        sync.WaitGroup
	resultsMu sync.Mutex
	results   []model.LlmComment
}

// NewCommentWorkerPool creates a pool with the given concurrency limit.
func NewCommentWorkerPool(workerCount int) *CommentWorkerPool {
	if workerCount <= 0 {
		workerCount = 8
	}
	return &CommentWorkerPool{
		semaphore: make(chan struct{}, workerCount),
	}
}

// Submit runs f in a background goroutine bounded by the semaphore.
// When f completes its return value is collected internally.
func (p *CommentWorkerPool) Submit(f func() ([]model.LlmComment, error)) {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()
		p.semaphore <- struct{}{}        // acquire
		defer func() { <-p.semaphore }() // release

		comments, err := f()
		if err != nil {
			fmt.Fprintf(stdout.Writer(), "[ocr] CommentWorkerPool error: %v\n", err)
		}
		p.resultsMu.Lock()
		p.results = append(p.results, comments...)
		p.resultsMu.Unlock()
	}()
}

// Await blocks until all submitted work has completed and returns aggregated results.
func (p *CommentWorkerPool) Await() []model.LlmComment {
	p.wg.Wait()
	return p.results
}

// New creates a new Agent from the given arguments.
func New(args Args) *Agent {
	if args.Tools == nil {
		args.Tools = make(tool.Registry)
	}
	if args.CommentCollector == nil {
		args.CommentCollector = tool.NewCommentCollector()
	}
	if args.Session == nil {
		gitBranch := detectGitBranch(args.RepoDir)
		args.Session = session.New(args.RepoDir, gitBranch, args.Model)
	}
	return &Agent{
		args:    args,
		session: args.Session,
	}
}

// Run executes the full review pipeline: parse diffs -> plan per file -> LLM tool-loop -> collect comments.
func (a *Agent) Run(ctx context.Context) ([]model.LlmComment, error) {
	// Step 1: Parse diffs
	ctx, diffSpan := telemetry.StartSpan(ctx, "diff.parse")
	if err := a.loadDiffs(); err != nil {
		diffSpan.End()
		return nil, fmt.Errorf("load diffs: %w", err)
	}
	telemetry.SetAttr(diffSpan, "files.changed", len(a.diffs))
	telemetry.SetAttr(diffSpan, "lines.inserted", int64(a.totalInsertions))
	telemetry.SetAttr(diffSpan, "lines.deleted", int64(a.totalDeletions))
	diffSpan.End()

	a.diffs = a.filterUnsupportedExts(a.diffs)

	if len(a.diffs) == 0 {
		fmt.Fprintln(stdout.Writer(), "[ocr] No supported files changed. Skipping review.")
		telemetry.Event(ctx, "no.files.changed")
		a.session.Finalize()
		return []model.LlmComment{}, nil
	}

	a.currentDate = time.Now().Format("2006-01-02 15:04")

	fmt.Fprintf(stdout.Writer(), "[ocr] Reviewing %d file(s) in %s\n", len(a.diffs), a.args.RepoDir)
	telemetry.Event(ctx, "review.started",
		telemetry.AnyToAttr("file.count", len(a.diffs)),
		telemetry.AnyToAttr("repo.dir", a.args.RepoDir))

	// Record file count metric.
	telemetry.RecordFilesReviewed(ctx, int64(len(a.diffs)))

	// Step 2: Dispatch per-file subtasks concurrently
	comments, err := a.dispatchSubtasks(ctx)
	if len(comments) > 0 {
		telemetry.RecordCommentsGenerated(ctx, int64(len(comments)))
	}
	a.session.Finalize()
	return comments, err
}

// Session returns the session history associated with this Agent.
func (a *Agent) Session() *session.SessionHistory {
	return a.session
}

// FilesReviewed returns the number of changed files included in this review.
func (a *Agent) FilesReviewed() int64 {
	return int64(len(a.diffs))
}

// Diffs returns the parsed diffs loaded by the agent.
func (a *Agent) Diffs() []model.Diff {
	return a.diffs
}

// TotalTokensUsed returns the accumulated total tokens from all LLM calls.
func (a *Agent) TotalTokensUsed() int64 {
	return atomic.LoadInt64(&a.totalTokensUsed)
}

// TotalInputTokens returns the accumulated input/prompt tokens from all LLM calls.
func (a *Agent) TotalInputTokens() int64 {
	return atomic.LoadInt64(&a.totalInputTokens)
}

// TotalOutputTokens returns the accumulated completion tokens from all LLM calls.
func (a *Agent) TotalOutputTokens() int64 {
	return atomic.LoadInt64(&a.totalOutputTokens)
}

// Warnings returns a copy of non-fatal warnings recorded during review.
func (a *Agent) Warnings() []AgentWarning {
	a.warningsMu.Lock()
	defer a.warningsMu.Unlock()
	out := make([]AgentWarning, len(a.warnings))
	copy(out, a.warnings)
	return out
}

// recordWarning adds a non-fatal warning to the agent's warning list.
func (a *Agent) recordWarning(warningType, file, message string) {
	a.warningsMu.Lock()
	a.warnings = append(a.warnings, AgentWarning{
		File:    file,
		Message: message,
		Type:    warningType,
	})
	a.warningsMu.Unlock()
}

// loadDiffs populates the diff-related fields.
func (a *Agent) loadDiffs() error {
	var provider *diff.Provider

	switch {
	case a.args.Commit != "":
		provider = diff.NewCommitProvider(a.args.RepoDir, a.args.Commit)
	case a.args.From != "" && a.args.To != "":
		provider = diff.NewProvider(a.args.RepoDir, a.args.From, a.args.To)
	default:
		provider = diff.NewWorkspaceProvider(a.args.RepoDir)
	}

	parsed, err := provider.GetDiff()
	if err != nil {
		return fmt.Errorf("get diffs: %w", err)
	}

	a.diffs = parsed

	for i := range parsed {
		d := &parsed[i]
		if d.NewPath != "/dev/null" && a.args.DiffMap != nil {
			a.args.DiffMap[d.NewPath] = d.Diff
		}
		a.totalInsertions += d.Insertions
		a.totalDeletions += d.Deletions
	}

	return nil
}

// lookupTool returns the provider for a given tool from the registry,
// or nil if the tool is not registered.
func lookupTool(reg tool.Registry, t tool.Tool) tool.Provider {
	p, ok := reg[t.Name()]
	if !ok {
		return nil
	}
	return p
}

// dispatchSubtasks runs the Plan + Main phases for each changed file concurrently.
func (a *Agent) dispatchSubtasks(ctx context.Context) ([]model.LlmComment, error) {
	startTime := time.Now()
	defer func() {
		telemetry.RecordReviewDuration(ctx, time.Since(startTime))
	}()

	// Pre-filter: discard diffs whose diff content alone exceeds 80% of the token threshold.
	a.diffs = a.filterLargeDiffs(a.diffs)
	if len(a.diffs) == 0 {
		return nil, fmt.Errorf("all diffs filtered out by token size")
	}

	var wg sync.WaitGroup

	concurrency := a.args.MaxConcurrency
	if concurrency <= 0 {
		concurrency = 8
	}

	sem := make(chan struct{}, concurrency)
	timeout := time.Duration(a.args.ConcurrentTaskTimeout) * time.Minute

	for i := range a.diffs {
		wg.Add(1)
		sem <- struct{}{} // acquire semaphore

		go func(d model.Diff) {
			defer wg.Done()
			defer func() { <-sem }() // release

			var fileCtx context.Context
			var cancel context.CancelFunc
			if timeout > 0 {
				fileCtx, cancel = context.WithTimeout(ctx, timeout)
				defer cancel()
			} else {
				fileCtx = ctx
			}

			if err := a.executeSubtask(fileCtx, d); err != nil {
				fmt.Fprintf(stdout.Writer(), "[ocr] Subtask error for %s: %v\n", d.NewPath, err)
				telemetry.ErrorEvent(fileCtx, "subtask.error", err,
					telemetry.AnyToAttr("file.path", d.NewPath))
			}
		}(a.diffs[i])
	}

	wg.Wait()

	// All subtasks finished — collect comments from the global collector once.
	if a.args.CommentWorkerPool != nil {
		a.args.CommentWorkerPool.Await()
	}
	return a.args.CommentCollector.Comments(), nil
}

// executeSubtask performs the Plan Phase + Main Loop for a single file.
func (a *Agent) executeSubtask(ctx context.Context, d model.Diff) error {
	ctx, span := telemetry.StartSpan(ctx, "subtask.execute."+d.NewPath)
	defer span.End()
	telemetry.SetAttr(span, "file.path", d.NewPath)
	telemetry.SetAttr(span, "lines.changed", d.Insertions+d.Deletions)
	telemetry.SetAttr(span, "lines.inserted", d.Insertions)
	telemetry.SetAttr(span, "lines.deleted", d.Deletions)

	if ctx.Err() != nil {
		return ctx.Err()
	}

	newPath := d.NewPath

	// Build change-files list excluding current file
	changeFilesExcludingCurrent := a.buildChangeFilesExcept(newPath)

	rule := a.resolveSystemRule(strings.ToLower(newPath))

	threshold := a.args.Template.PlanModeLineThreshold
	changeLines := d.Insertions + d.Deletions

	// Phase 1: Plan (skip when changes are below threshold)
	var planResult string
	if a.args.Template.PlanTask != nil && len(a.args.Template.PlanTask.Messages) > 0 && threshold > 0 && changeLines < int64(threshold) {
		fmt.Fprintf(stdout.Writer(), "[ocr] Skipping plan phase for %s (%d lines < threshold %d)\n", newPath, changeLines, threshold)
		telemetry.Event(ctx, "plan.skipped",
			telemetry.AnyToAttr("file.path", newPath),
			telemetry.AnyToAttr("lines.changed", changeLines),
			telemetry.AnyToAttr("threshold", threshold))
	} else if a.args.Template.PlanTask != nil && len(a.args.Template.PlanTask.Messages) > 0 {
		var err error
		planResult, err = a.executePlanPhase(ctx, newPath, d.Diff, changeFilesExcludingCurrent, rule)
		if err != nil {
			fmt.Fprintf(stdout.Writer(), "[ocr] Plan phase failed for %s: %v (continuing without plan)\n", newPath, err)
			telemetry.Eventf(ctx, "plan.failed", err.Error(),
				telemetry.AnyToAttr("file.path", newPath))
			planResult = ""
		}
	}

	// Phase 2: Main task loop
	if len(a.args.Template.MainTask.Messages) == 0 {
		return fmt.Errorf("main_task.messages is empty in template")
	}

	rawMsgs := a.args.Template.MainTask.Messages
	messages := make([]llm.Message, 0, len(rawMsgs))
	for _, m := range rawMsgs {
		content := m.Content
		content = strings.ReplaceAll(content, "{{current_system_date_time}}", a.currentDate)
		content = strings.ReplaceAll(content, "{{current_file_path}}", newPath)
		content = strings.ReplaceAll(content, "{{system_rule}}", rule)
		content = strings.ReplaceAll(content, "{{change_files}}", changeFilesExcludingCurrent)
		content = strings.ReplaceAll(content, "{{diff}}", d.Diff)
		content = strings.ReplaceAll(content, "{{requirement_background}}", a.args.Background)
		if planResult == "" {
			content = strings.ReplaceAll(content, "### 审查计划\n{{plan_guidance}}\n\n", "")
		} else {
			content = strings.ReplaceAll(content, "{{plan_guidance}}", planResult)
		}
		messages = append(messages, llm.NewTextMessage(m.Role, content))
	}

	tokenCount := countMessagesTokens(messages)
	maxAllowed := a.args.Template.MaxTokens
	tokenLimit := maxAllowed * 4 / 5 // 80% of MaxTokens
	if tokenCount > tokenLimit {
		msg := fmt.Sprintf("prompt tokens (%d) exceed %d%% of max_tokens(%d)", tokenCount, 80, maxAllowed)
		fmt.Fprintf(stdout.Writer(), "[ocr] WARNING: %s for %s\n", msg, newPath)
		a.recordWarning("token_threshold_exceeded", newPath, msg)
		telemetry.Event(ctx, "token.threshold.exceeded",
			telemetry.AnyToAttr("file.path", newPath),
			telemetry.AnyToAttr("tokens", tokenCount),
			telemetry.AnyToAttr("max_tokens", maxAllowed))
		return nil
	}

	return a.performLlmCodeReview(ctx, messages, newPath)
}

// buildChangeFilesExcept returns a formatted list of changed files except the given path.
func (a *Agent) buildChangeFilesExcept(excludePath string) string {
	var sb strings.Builder
	for i, d := range a.diffs {
		if d.IsBinary {
			continue
		}
		if d.NewPath == excludePath || d.OldPath == excludePath {
			continue
		}
		status := "MODIFIED"
		switch {
		case d.IsNew:
			status = "ADDED"
		case d.IsDeleted:
			status = "DELETED"
		case d.OldPath != d.NewPath:
			status = "RENAMED"
		}
		sb.WriteString(status + "   " + d.NewPath)
		if i < len(a.diffs)-1 {
			sb.WriteString("\n")
		}
	}
	return sb.String()
}

// resolveSystemRule returns the rule text for a given file path,
// matching against PathRuleMap glob patterns, falling back to DefaultRule.
func (a *Agent) resolveSystemRule(path string) string {
	if a.args.SystemRule == nil {
		return ""
	}
	return a.args.SystemRule.Resolve(path)
}

// filterLargeDiffs drops diffs whose diff content alone consumes more than 80% of MaxTokens.
func (a *Agent) filterLargeDiffs(diffs []model.Diff) []model.Diff {
	limit := a.args.Template.MaxTokens * 4 / 5
	if limit <= 0 {
		return diffs
	}
	var kept []model.Diff
	skipped := 0

	for _, d := range diffs {
		tokens := llm.CountTokens(d.Diff)
		if tokens > limit {
			fmt.Fprintf(stdout.Writer(), "[ocr] Skipping %s (~%d tokens exceeds 80%% of max_tokens(%d))\n",
				d.NewPath, tokens, a.args.Template.MaxTokens)
			skipped++
			continue
		}
		kept = append(kept, d)
	}

	if skipped > 0 {
		fmt.Fprintf(stdout.Writer(), "[ocr] Pre-filtered %d file(s) exceeding 80%% of max_tokens\n", skipped)
	}
	return kept
}

// filterUnsupportedExts drops diffs whose file extensions are not in the supported types allowlist.
func (a *Agent) filterUnsupportedExts(diffs []model.Diff) []model.Diff {
	var kept []model.Diff
	skipped := 0

	for _, d := range diffs {
		path := d.NewPath
		if path == "/dev/null" {
			path = d.OldPath
		}
		ext := a.extFromPath(path)
		if ext != "" && !allowedext.IsAllowedExt(ext) {
			fmt.Fprintf(stdout.Writer(), "[ocr] Skipping %s — extension %q not in supported file types\n", d.NewPath, ext)
			skipped++
			continue
		}
		kept = append(kept, d)
	}

	if skipped > 0 {
		fmt.Fprintf(stdout.Writer(), "[ocr] Skipped %d file(s) with unsupported extensions\n", skipped)
	}
	return kept
}

// extFromPath returns the file extension with leading dot, lowercased.
func (a *Agent) extFromPath(path string) string {
	basename := path
	if idx := strings.LastIndex(path, "/"); idx >= 0 {
		basename = path[idx+1:]
	}
	dot := strings.LastIndex(basename, ".")
	if dot <= 0 {
		return ""
	}
	return strings.ToLower(basename[dot:])
}

// executePlanPhase runs the plan task for a single file, sending template messages
// with resolved placeholders and collecting the LLM response as plan guidance.
func (a *Agent) executePlanPhase(ctx context.Context, newPath, rawDiff, changeFiles, rule string) (string, error) {
	pt := a.args.Template.PlanTask
	messages := make([]llm.Message, 0, len(pt.Messages))
	for _, m := range pt.Messages {
		content := m.Content
		content = strings.ReplaceAll(content, "{{current_system_date_time}}", a.currentDate)
		content = strings.ReplaceAll(content, "{{current_file_path}}", newPath)
		content = strings.ReplaceAll(content, "{{system_rule}}", rule)
		content = strings.ReplaceAll(content, "{{change_files}}", changeFiles)
		content = strings.ReplaceAll(content, "{{diff}}", rawDiff)
		content = strings.ReplaceAll(content, "{{requirement_background}}", a.args.Background)
		content = strings.ReplaceAll(content, "{{plan_tools}}", formatToolDefs(a.args.PlanToolDefs))
		messages = append(messages, llm.NewTextMessage(m.Role, content))
	}

	fs := a.session.GetOrCreateFileSession(newPath)
	rec := fs.AppendTaskRecord(session.PlanTask, messages)
	startTime := time.Now()

	resp, err := a.args.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{
		Model:     a.args.Model,
		Messages:  messages,
		MaxTokens: a.args.Template.MaxTokens,
	})
	if err != nil {
		rec.SetError(err, time.Since(startTime))
		return "", fmt.Errorf("plan request: %w", err)
	}
	rec.SetResponse(resp, time.Since(startTime))
	if resp.Usage != nil {
		atomic.AddInt64(&a.totalTokensUsed, int64(resp.Usage.TotalTokens))
		atomic.AddInt64(&a.totalInputTokens, int64(resp.Usage.PromptTokens+resp.Usage.CacheReadTokens))
		atomic.AddInt64(&a.totalOutputTokens, int64(resp.Usage.CompletionTokens+resp.Usage.CacheWriteTokens))
	}
	fmt.Fprintf(stdout.Writer(), "[ocr] Plan completed for %s\n", newPath)
	return resp.Content(), nil
}

// formatToolDefs renders tool definitions as human-readable text for embedding in prompts.
func formatToolDefs(toolDefs []llm.ToolDef) string {
	if len(toolDefs) == 0 {
		return ""
	}

	var sb strings.Builder
	sb.WriteString("### Available Tools (reference only — do not call)\n")
	for _, td := range toolDefs {
		fn := &td.Function
		sb.WriteString(fmt.Sprintf("- **%s**: %s\n", fn.Name, fn.Description))
		if params, ok := fn.Parameters["properties"].(map[string]any); ok && len(params) > 0 {
			sb.WriteString("  Parameters:\n")
			required := make(map[string]bool)
			if reqList, ok := fn.Parameters["required"].([]any); ok {
				for _, r := range reqList {
					if s, ok := r.(string); ok {
						required[s] = true
					}
				}
			}
			for name, p := range params {
				suffix := ""
				if required[name] {
					suffix = " (required)"
				}
				if pm, ok := p.(map[string]any); ok {
					desc, _ := pm["description"].(string)
					sb.WriteString(fmt.Sprintf("  - %s: %s%s\n", name, desc, suffix))
				} else {
					sb.WriteString(fmt.Sprintf("  - %s%s\n", name, suffix))
				}
			}
		}
	}
	return sb.String()
}

// performLlmCodeReview drives the main LLM conversation loop for a single file.
// It sends messages with tool definitions, handles tool calls returned by the model,
// and collects review comments until task_done is called or limits are reached.
func (a *Agent) performLlmCodeReview(ctx context.Context, messages []llm.Message, newPath string) error {
	toolReqCount := a.args.Template.MaxToolRequestTimes

	for toolReqCount > 0 {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		toolReqCount--

		fs := a.session.GetOrCreateFileSession(newPath)
		rec := fs.AppendTaskRecord(session.MainTask, append([]llm.Message(nil), messages...))
		startTime := time.Now()

		resp, err := a.args.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{
			Model:     a.args.Model,
			Messages:  messages,
			Tools:     a.args.MainToolDefs,
			MaxTokens: a.args.Template.MaxTokens,
		})
		duration := time.Since(startTime)
		if err != nil {
			rec.SetError(err, duration)
			telemetry.RecordLLMRequest(ctx, a.args.Model, duration, 0, "error")
			return fmt.Errorf("LLM completion error: %w", err)
		}
		rec.SetResponse(resp, duration)
		// Record LLM metrics with token info from API response usage field.
		totalTokens := int64(0)
		if resp.Usage != nil {
			totalTokens = resp.Usage.TotalTokens
			atomic.AddInt64(&a.totalInputTokens, int64(resp.Usage.PromptTokens+resp.Usage.CacheReadTokens))
			atomic.AddInt64(&a.totalOutputTokens, int64(resp.Usage.CompletionTokens+resp.Usage.CacheWriteTokens))
		}
		telemetry.RecordLLMRequest(ctx, a.args.Model, duration, totalTokens, "ok")
		atomic.AddInt64(&a.totalTokensUsed, totalTokens)

		content := resp.Content()
		calls := resp.ToolCalls()

		if len(calls) == 0 {
			// No tool calls - remind the model
			fmt.Fprintf(stdout.Writer(), "[ocr] No tool calls parsed for %s, retrying...\n", newPath)
			messages = append(messages, llm.NewTextMessage("user", "You did not successfully call any tools. Please try again or use task_done if finished."))
			if content != "" {
				messages = append(messages[:len(messages)-1], llm.NewTextMessage("assistant", content), messages[len(messages)-1])
			}
			continue
		}

		var results []tool.ToolCallResult
		taskCompleted := false
		hasValidResult := false

		for _, call := range calls {
			cp := a.executeToolCall(ctx, newPath, call, rec)
			if cp.Completed {
				results = append(results, tool.ToolCallResult{
					ToolCallID: call.ID,
					Name:       call.Function.Name,
					Result:     "Task completed successfully.",
				})
				taskCompleted = true
			} else if cp.Data != "" {
				results = append(results, tool.ToolCallResult{
					ToolCallID: call.ID,
					Name:       call.Function.Name,
					Result:     cp.Data,
				})
				hasValidResult = true
			} else {
				results = append(results, tool.ToolCallResult{
					ToolCallID: call.ID,
					Name:       call.Function.Name,
					Result:     "Error: Tool execution returned no result.",
				})
			}
		}

		if taskCompleted {
			break
		}
		if !hasValidResult {
			fmt.Fprintf(stdout.Writer(), "[ocr] No valid tool results for %s, stopping.\n", newPath)
			break
		}

		succeed := a.addNextMessage(content, calls, results, &messages, newPath)
		if !succeed {
			fmt.Fprintf(stdout.Writer(), "[ocr] Context compression exceeded threshold for %s, stopping.\n", newPath)
			break
		}
	}

	if toolReqCount <= 0 {
		fmt.Fprintf(stdout.Writer(), "[ocr] Max tool requests reached for %s.\n", newPath)
	}

	return nil
}

// executeToolCall executes a single tool call from the LLM response and records
// the result in session history. It handles async dispatch for code_comment when
// a worker pool is configured, otherwise runs synchronously.
func (a *Agent) executeToolCall(ctx context.Context, newPath string, call llm.ToolCall, rec *session.TaskRecord) tool.TaskCheckpoint {
	t := tool.OfName(call.Function.Name)
	if !t.IsKnown() {
		return tool.Of(tool.NotAvailableMsg)
	}

	if t == tool.TaskDone {
		return tool.Complete()
	}

	p := lookupTool(a.args.Tools, t)
	if p == nil {
		return tool.Of(tool.NotAvailableMsg)
	}

	var args map[string]any
	if err := json.Unmarshal([]byte(call.Function.Arguments), &args); err != nil {
		return tool.Of(fmt.Sprintf("Error parsing tool arguments for %s: %v", t.Name(), err))
	}

	// Inject current file path as default for code_comment when not provided.
	// The model already knows which file it's reviewing, so it omits path.
	if t == tool.CodeComment && newPath != "" {
		if _, ok := args["path"]; !ok {
			args["path"] = newPath
		}
	}

	startTime := time.Now()

	// Async path for code_comment when worker pool is configured
	// Mirrors Java: pendingCommentFutures.add(subtaskExecutor.submit(() -> getCodeComments(...)))
	if t == tool.CodeComment && a.args.CommentWorkerPool != nil {
		if rec != nil {
			rec.AddToolResult(t.Name(), call.Function.Arguments, "(async)")
		}
		pool := a.args.CommentWorkerPool
		telemetry.PrintToolCallStarted(t.Name(), args)
		pool.Submit(func() ([]model.LlmComment, error) {
			_, _ = p.Execute(args) // provider parses args and adds to collector
			telemetry.RecordToolCall(ctx, t.Name(), time.Since(startTime), true)
			return []model.LlmComment{}, nil
		})
		// Return immediate success - actual comment processing continues off
		// the critical path, exactly like Java's subtaskExecutor.submit for CODE_COMMENT.
		telemetry.RecordToolCall(ctx, t.Name(), time.Since(startTime), true)
		return tool.Of(tool.CommentSucceed)
	}

	// Synchronous path for all other tools
	telemetry.PrintToolCallStarted(t.Name(), args)
	result, err := p.Execute(args)
	dur := time.Since(startTime)
	ok := err == nil
	telemetry.RecordToolCall(ctx, t.Name(), dur, ok)

	if err != nil {
		telemetry.PrintToolCallError(t.Name(), err)
		return tool.Of(fmt.Sprintf("Error executing tool %s: %v", t.Name(), err))
	}
	telemetry.PrintToolCallFinished(t.Name(), dur)
	if rec != nil {
		rec.AddToolResult(t.Name(), call.Function.Arguments, result)
	}
	return tool.Of(result)
}

// collectPendingComments waits for any async workers then returns aggregated comments from the collector.
func (a *Agent) collectPendingComments() []model.LlmComment {
	if a.args.CommentWorkerPool != nil {
		a.args.CommentWorkerPool.Await()
	}
	return a.args.CommentCollector.Comments()
}

// addNextMessage adds assistant + tool response messages to the conversation history.
// Implements dual-threshold compression:
//   - 60% of MaxTokens: trigger async background compression (non-blocking)
//   - 80% of MaxTokens: perform synchronous compression immediately
func (a *Agent) addNextMessage(assistantContent string, toolCalls []llm.ToolCall, results []tool.ToolCallResult, messages *[]llm.Message, filePath string) bool {
	maxAllowed := a.args.Template.MaxTokens
	softLimit := int(float64(maxAllowed) * tokenSoftThreshold)
	warnLimit := int(float64(maxAllowed) * tokenWarningThreshold)

	// Try to apply any completed async compression from a previous iteration.
	a.tryApplyPendingCompression(messages)

	tokenCount := countMessagesTokens(*messages)

	// Hard threshold: synchronous compression.
	if tokenCount > warnLimit {
		a.cancelPendingCompression()
		*messages, _ = a.runCompression(*messages, filePath)
		tokenCount = countMessagesTokens(*messages)
	}

	// Soft threshold: async compression.
	if tokenCount > softLimit && a.pendingJob == nil {
		a.triggerAsyncCompression(*messages, filePath)
	}

	// Add assistant message with tool_calls when present.
	if len(toolCalls) > 0 {
		*messages = append(*messages, llm.NewToolCallMessage(assistantContent, toolCalls))
	} else if assistantContent != "" {
		*messages = append(*messages, llm.NewTextMessage("assistant", assistantContent))
	}

	// Add tool response messages using Claude's tool_result format.
	for _, r := range results {
		*messages = append(*messages, llm.NewToolResultMessage(r.ToolCallID, r.Result))
	}

	// Final check: compress synchronously if still over warning limit.
	finalCount := countMessagesTokens(*messages)
	if finalCount > warnLimit {
		a.cancelPendingCompression()
		*messages, _ = a.runCompression(*messages, filePath)
	}

	return countMessagesTokens(*messages) < warnLimit
}

func countMessagesTokens(msgs []llm.Message) int {
	var total int
	for _, m := range msgs {
		total += llm.CountTokens(m.ExtractText())
	}
	return total
}

// groupIntoRounds parses messages[start:] into logical (assistant + tool_results) pairs.
func groupIntoRounds(messages []llm.Message, start int) []round {
	var rounds []round
	i := start
	for i < len(messages) {
		if messages[i].Role == "assistant" {
			r := round{assistantIdx: i}
			i++
			for i < len(messages) && messages[i].Role == "tool" {
				r.toolIdxs = append(r.toolIdxs, i)
				i++
			}
			rounds = append(rounds, r)
		} else {
			i++
		}
	}
	return rounds
}

// computeActiveZoneSize returns how many trailing rounds fit within the remaining
// token budget after accounting for frozen zone and the compressed summary.
func computeActiveZoneSize(rounds []round, messages []llm.Message, maxTokens int, reservedTokens int) int {
	budget := int(float64(maxTokens)*tokenWarningThreshold) - reservedTokens
	if budget <= 0 {
		return 0
	}

	count := 0
	tokensUsed := 0
	for i := len(rounds) - 1; i >= 0; i-- {
		roundTokens := llm.CountTokens(messages[rounds[i].assistantIdx].ExtractText())
		for _, ti := range rounds[i].toolIdxs {
			roundTokens += llm.CountTokens(messages[ti].ExtractText())
		}
		if tokensUsed+roundTokens > budget {
			break
		}
		tokensUsed += roundTokens
		count++
	}
	return count
}

// partitionMessages divides messages into frozen, compress, and active zones.
// Frozen zone is always messages[0:2]. Active zone preserves the K most recent
// complete rounds based on available token budget.
func partitionMessages(messages []llm.Message, maxTokens int, prevSummaryTokenEstimate int) partitionResult {
	result := partitionResult{frozenEnd: 2}
	if len(messages) <= 2 {
		result.compressEnd = len(messages)
		return result
	}

	result.rounds = groupIntoRounds(messages, 2)
	if len(result.rounds) == 0 {
		result.compressEnd = len(messages)
		return result
	}

	result.activeCount = computeActiveZoneSize(result.rounds, messages, maxTokens, prevSummaryTokenEstimate)
	if result.activeCount >= len(result.rounds) {
		// Everything fits — no compression needed.
		result.compressEnd = len(messages)
		result.activeCount = 0
		return result
	}

	// compressEnd = index after the last round NOT in active zone.
	activeStartIdx := len(result.rounds) - result.activeCount
	lastCompressRound := result.rounds[activeStartIdx-1]
	if len(lastCompressRound.toolIdxs) > 0 {
		result.compressEnd = lastCompressRound.toolIdxs[len(lastCompressRound.toolIdxs)-1] + 1
	} else {
		result.compressEnd = lastCompressRound.assistantIdx + 1
	}

	return result
}

// stripMarkdownFences removes ```json and ``` wrappers that some models
// add around structured outputs.
func stripMarkdownFences(s string) string {
	s = strings.TrimSpace(s)
	if strings.HasPrefix(s, "```") {
		if nl := strings.IndexByte(s, '\n'); nl >= 0 {
			s = s[nl+1:]
		} else {
			s = strings.TrimPrefix(s, "```json")
			s = strings.TrimPrefix(s, "```")
		}
	}
	s = strings.TrimSpace(s)
	if strings.HasSuffix(s, "```") {
		s = strings.TrimSuffix(s, "```")
		s = strings.TrimSpace(s)
	}
	return s
}

// runCompression performs three-zone memory compression on the given messages.
// It summarizes the compress zone while preserving the active zone intact.
// Returns the rebuilt messages slice: [frozen] + [compressed_summary] + [active].
func (a *Agent) runCompression(msgs []llm.Message, filePath string) ([]llm.Message, error) {
	if len(a.args.Template.MemoryCompressionTask.Messages) == 0 || len(msgs) <= 2 {
		return msgs[:min(len(msgs), 2)], nil
	}

	part := partitionMessages(msgs, a.args.Template.MaxTokens, 0)
	if part.compressEnd <= part.frozenEnd {
		return msgs, nil
	}

	contextXML := buildMessageXML(msgs[part.frozenEnd:part.compressEnd])

	compressionMsgs := make([]llm.Message, 0, len(a.args.Template.MemoryCompressionTask.Messages))
	for _, m := range a.args.Template.MemoryCompressionTask.Messages {
		content := strings.ReplaceAll(m.Content, "{{context}}", contextXML)
		compressionMsgs = append(compressionMsgs, llm.NewTextMessage(m.Role, content))
	}

	startTime := time.Now()
	resp, err := a.args.LLMClient.Completions(llm.ChatRequest{
		Model:     a.args.Model,
		Messages:  compressionMsgs,
		MaxTokens: a.args.Template.MaxTokens,
	})
	duration := time.Since(startTime)

	fs := a.session.GetOrCreateFileSession(filePath)
	rec := fs.AppendTaskRecord(session.MemoryCompressionTask, compressionMsgs)
	if err != nil {
		rec.SetError(err, duration)
		fmt.Fprintf(stdout.Writer(), "[ocr] Memory compression failed: %v\n", err)
		return msgs[:part.frozenEnd], fmt.Errorf("memory compression: %w", err)
	}
	rec.SetResponse(resp, duration)
	if resp.Usage != nil {
		atomic.AddInt64(&a.totalTokensUsed, int64(resp.Usage.TotalTokens))
		atomic.AddInt64(&a.totalInputTokens, int64(resp.Usage.PromptTokens+resp.Usage.CacheReadTokens))
		atomic.AddInt64(&a.totalOutputTokens, int64(resp.Usage.CompletionTokens+resp.Usage.CacheWriteTokens))
	}

	rawSummary := stripMarkdownFences(resp.Content())
	if rawSummary == "" {
		return msgs[:part.frozenEnd], nil
	}

	rebuilt := make([]llm.Message, 2)
	copy(rebuilt, msgs[:2])

	userMsg := rebuilt[1]
	currentText := userMsg.ExtractText()
	rebuilt[1] = llm.NewTextMessage(userMsg.Role, currentText+"\n\n<previous_review_summary>\n"+rawSummary+"\n</previous_review_summary>")

	for i := part.compressEnd; i < len(msgs); i++ {
		rebuilt = append(rebuilt, msgs[i])
	}

	return rebuilt, nil
}

// triggerAsyncCompression kicks off a background compression job.
func (a *Agent) triggerAsyncCompression(messages []llm.Message, filePath string) {
	msgSnapshot := copyMessages(messages)

	job := &compressionJob{done: make(chan struct{})}
	a.compressionMu.Lock()
	a.pendingJob = job
	a.compressionMu.Unlock()

	go func() {
		rebuilt, _ := a.runCompression(msgSnapshot, filePath)

		a.compressionMu.Lock()
		defer a.compressionMu.Unlock()

		if a.pendingJob != job {
			return // cancelled or superseded
		}
		job.rebuilt = rebuilt
		close(job.done)
	}()
}

// tryApplyPendingCompression checks if a background compression completed
// and swaps the rebuilt messages into place. Returns true if applied.
func (a *Agent) tryApplyPendingCompression(messages *[]llm.Message) bool {
	a.compressionMu.Lock()
	job := a.pendingJob
	a.compressionMu.Unlock()

	if job == nil {
		return false
	}

	select {
	case <-job.done:
		a.compressionMu.Lock()
		if a.pendingJob == job && job.rebuilt != nil {
			*messages = job.rebuilt
			a.pendingJob = nil
		}
		a.compressionMu.Unlock()
		return true
	default:
		return false
	}
}

// cancelPendingCompression aborts any in-flight background compression.
func (a *Agent) cancelPendingCompression() {
	a.compressionMu.Lock()
	defer a.compressionMu.Unlock()

	if a.pendingJob != nil {
		a.pendingJob = nil
	}
}

// copyMessages creates a shallow copy of a message slice.
func copyMessages(msgs []llm.Message) []llm.Message {
	out := make([]llm.Message, len(msgs))
	copy(out, msgs)
	return out
}

// BuildToolDefs converts toolsconfig.ToolConfigEntry slice into []llm.ToolDef,
// filtering by phase (planOnly=true for plan_task, false for main_task).
func BuildToolDefs(entries []toolsconfig.ToolConfigEntry, planOnly bool) []llm.ToolDef {
	var defs []llm.ToolDef
	for _, e := range entries {
		defRaw, ok := e.ToolDefsByPhase(planOnly)
		if !ok {
			continue
		}
		var fn llm.FunctionDef
		if err := json.Unmarshal(defRaw, &fn); err != nil {
			fmt.Fprintf(stdout.Writer(), "[ocr] WARNING: failed to parse tool definition %q: %v\n", e.Name, err)
			continue
		}
		defs = append(defs, llm.ToolDef{
			Type:     "function",
			Function: fn,
		})
	}
	return defs
}

func buildMessageXML(msgs []llm.Message) string {
	var sb strings.Builder
	for i, m := range msgs {
		sb.WriteString(fmt.Sprintf("<message id=\"%d\" role=\"%s\">\n", i, m.Role))
		sb.WriteString("    <content>\n")
		sb.WriteString(fmt.Sprintf("      %s\n", m.ExtractText()))
		sb.WriteString("    </content>\n")
		sb.WriteString("</message>")
		if i < len(msgs)-1 {
			sb.WriteString("\n")
		}
	}
	return sb.String()
}

// detectGitBranch returns the current git branch name for the given repo, or empty string on failure.
func detectGitBranch(repoDir string) string {
	cmd := exec.Command("git", "-C", repoDir, "rev-parse", "--abbrev-ref", "HEAD")
	out, err := cmd.Output()
	if err != nil || len(out) == 0 {
		return ""
	}
	return strings.TrimSpace(string(out))
}