/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import type { Content } from '@google/genai';
import { createHash } from 'node:crypto';
import { GeminiEventType, type ServerGeminiStreamEvent } from '../core/turn.js';
import {
  logLoopDetected,
  logLoopDetectionDisabled,
  logLlmLoopCheck,
} from '../telemetry/loggers.js';
import {
  LoopDetectedEvent,
  LoopDetectionDisabledEvent,
  LoopType,
  LlmLoopCheckEvent,
  LlmRole,
} from '../telemetry/types.js';
import {
  isFunctionCall,
  isFunctionResponse,
} from '../utils/messageInspectors.js';
import { debugLogger } from '../utils/debugLogger.js';
import type { AgentLoopContext } from '../config/agent-loop-context.js';

const TOOL_CALL_LOOP_THRESHOLD = 5;
const CONTENT_LOOP_THRESHOLD = 10;
const CONTENT_CHUNK_SIZE = 50;
const MAX_HISTORY_LENGTH = 5000;

/**
 * The number of recent conversation turns to include in the history when asking the LLM to check for a loop.
 */
const LLM_LOOP_CHECK_HISTORY_COUNT = 20;

/**
 * The number of turns that must pass in a single prompt before the LLM-based loop check is activated.
 */
const LLM_CHECK_AFTER_TURNS = 30;

/**
 * The default interval, in number of turns, at which the LLM-based loop check is performed.
 * This value is adjusted dynamically based on the LLM's confidence.
 */
const DEFAULT_LLM_CHECK_INTERVAL = 10;

/**
 * The minimum interval for LLM-based loop checks.
 * This is used when the confidence of a loop is high, to check more frequently.
 */
const MIN_LLM_CHECK_INTERVAL = 5;

/**
 * The maximum interval for LLM-based loop checks.
 * This is used when the confidence of a loop is low, to check less frequently.
 */
const MAX_LLM_CHECK_INTERVAL = 15;

/**
 * The confidence threshold above which the LLM is considered to have detected a loop.
 */
const LLM_CONFIDENCE_THRESHOLD = 0.9;
const DOUBLE_CHECK_MODEL_ALIAS = 'loop-detection-double-check';

const LOOP_DETECTION_SYSTEM_PROMPT = `You are a diagnostic agent that determines whether a conversational AI assistant is stuck in an unproductive loop. Analyze the conversation history (and, if provided, the original user request) to make this determination.

## What constitutes an unproductive state

An unproductive state requires BOTH of the following to be true:
1. The assistant has exhibited a repetitive pattern over at least 5 consecutive model actions (tool calls or text responses, counting only model-role turns).
2. The repetition produces NO net change or forward progress toward the user's goal.

Specific patterns to look for:
- **Alternating cycles with no net effect:** The assistant cycles between the same actions (e.g., edit_file → run_build → edit_file → run_build) where each iteration applies the same edit and encounters the same error, making zero progress. Note: alternating between actions is only a loop if the arguments and outcomes are substantively identical each cycle. If the assistant is modifying different code or getting different errors, that is debugging progress, not a loop.
- **Semantic repetition with identical outcomes:** The assistant calls the same tool with semantically equivalent arguments (same file, same line range, same content) multiple times consecutively, and each call produces the same outcome. This does NOT include build/test commands that are re-run after making code changes between invocations — re-running a build to verify a fix is normal workflow.
- **Stuck reasoning:** The assistant produces multiple consecutive text responses that restate the same plan, question, or analysis without taking any new action or making a decision. This does NOT include command output that happens to contain repeated status lines or warnings.

## What is NOT an unproductive state

You MUST distinguish repetitive-looking but productive work from true loops. The following are examples of forward progress and must NOT be flagged:

- **Cross-file batch operations:** A series of tool calls with the same tool name but targeting different files (different file paths in the arguments). For example, adding license headers to 20 files, or running the same refactoring across multiple modules.
- **Incremental same-file edits:** Multiple edits to the same file that target different line ranges, different functions, or different text content (e.g., adding docstrings to functions one by one).
- **Sequential processing:** A series of read or search operations on different files/paths to gather information.
- **Retry with variation:** Re-attempting a failed operation with modified arguments or a different approach.

## Argument analysis (critical)

When evaluating tool calls, you MUST compare the **arguments** of each call, not just the tool name. Pay close attention to:
- **File paths:** Different file paths mean different targets — this is distinct work, not repetition.
- **Line numbers and text content:** Different line ranges or different old_string/new_string values indicate distinct edits.
- **Search queries and patterns:** Different search terms indicate information gathering, not looping.

A loop exists only when the same tool is called with semantically equivalent arguments repeatedly, indicating no forward progress.

## Using the original user request

If the original user request is provided, use it to contextualize the assistant's behavior. If the request implies a batch or multi-step operation (e.g., "update all files", "refactor every module", "add tests for each function"), then repetitive tool calls with varying arguments are expected and should weigh heavily against flagging a loop.`;

const LOOP_DETECTION_SCHEMA: Record<string, unknown> = {
  type: 'object',
  properties: {
    unproductive_state_analysis: {
      type: 'string',
      description:
        'Your reasoning on if the conversation is looping without forward progress.',
    },
    unproductive_state_confidence: {
      type: 'number',
      description:
        'A number between 0.0 and 1.0 representing your confidence that the conversation is in an unproductive state.',
    },
  },
  required: ['unproductive_state_analysis', 'unproductive_state_confidence'],
};

/**
 * Result of a loop detection check.
 */
export interface LoopDetectionResult {
  count: number;
  type?: LoopType;
  detail?: string;
  confirmedByModel?: string;
}
/**
 * Service for detecting and preventing infinite loops in AI responses.
 * Monitors tool call repetitions and content sentence repetitions.
 */
export class LoopDetectionService {
  private readonly context: AgentLoopContext;
  private promptId = '';
  private userPrompt = '';

  // Tool call tracking
  private lastToolCallKey: string | null = null;
  private toolCallRepetitionCount: number = 0;

  // Content streaming tracking
  private streamContentHistory = '';
  private contentStats = new Map<string, number[]>();
  private lastContentIndex = 0;
  private loopDetected = false;
  private detectedCount = 0;
  private lastLoopDetail?: string;
  private inCodeBlock = false;

  private lastLoopType?: LoopType;
  // LLM loop track tracking
  private turnsInCurrentPrompt = 0;
  private llmCheckInterval = DEFAULT_LLM_CHECK_INTERVAL;
  private lastCheckTurn = 0;

  // Session-level disable flag
  private disabledForSession = false;

  constructor(context: AgentLoopContext) {
    this.context = context;
  }

  /**
   * Disables loop detection for the current session.
   */
  disableForSession(): void {
    this.disabledForSession = true;
    logLoopDetectionDisabled(
      this.context.config,
      new LoopDetectionDisabledEvent(this.promptId),
    );
  }

  private getToolCallKey(toolCall: { name: string; args: object }): string {
    const argsString = JSON.stringify(toolCall.args);
    const keyString = `${toolCall.name}:${argsString}`;
    return createHash('sha256').update(keyString).digest('hex');
  }

  /**
   * Processes a stream event and checks for loop conditions.
   * @param event - The stream event to process
   * @returns A LoopDetectionResult
   */
  addAndCheck(event: ServerGeminiStreamEvent): LoopDetectionResult {
    if (
      this.disabledForSession ||
      this.context.config.getDisableLoopDetection()
    ) {
      return { count: 0 };
    }
    if (this.loopDetected) {
      return {
        count: this.detectedCount,
        type: this.lastLoopType,
        detail: this.lastLoopDetail,
      };
    }

    let isLoop = false;
    let detail: string | undefined;

    switch (event.type) {
      case GeminiEventType.ToolCallRequest:
        // content chanting only happens in one single stream, reset if there
        // is a tool call in between
        this.resetContentTracking();
        isLoop = this.checkToolCallLoop(event.value);
        if (isLoop) {
          detail = `Repeated tool call: ${event.value.name} with arguments ${JSON.stringify(event.value.args)}`;
        }
        break;
      case GeminiEventType.Content:
        isLoop = this.checkContentLoop(event.value);
        if (isLoop) {
          detail = `Repeating content detected: "${this.streamContentHistory.substring(Math.max(0, this.lastContentIndex - 20), this.lastContentIndex + CONTENT_CHUNK_SIZE).trim()}..."`;
        }
        break;
      default:
        break;
    }

    if (isLoop) {
      this.loopDetected = true;
      this.detectedCount++;
      this.lastLoopDetail = detail;
      this.lastLoopType =
        event.type === GeminiEventType.ToolCallRequest
          ? LoopType.CONSECUTIVE_IDENTICAL_TOOL_CALLS
          : LoopType.CONTENT_CHANTING_LOOP;

      logLoopDetected(
        this.context.config,
        new LoopDetectedEvent(
          this.lastLoopType,
          this.promptId,
          this.detectedCount,
        ),
      );
    }
    return isLoop
      ? {
          count: this.detectedCount,
          type: this.lastLoopType,
          detail: this.lastLoopDetail,
        }
      : { count: 0 };
  }

  /**
   * Signals the start of a new turn in the conversation.
   *
   * This method increments the turn counter and, if specific conditions are met,
   * triggers an LLM-based check to detect potential conversation loops. The check
   * is performed periodically based on the `llmCheckInterval`.
   *
   * @param signal - An AbortSignal to allow for cancellation of the asynchronous LLM check.
   * @returns A promise that resolves to a LoopDetectionResult.
   */
  async turnStarted(signal: AbortSignal): Promise<LoopDetectionResult> {
    if (
      this.disabledForSession ||
      this.context.config.getDisableLoopDetection()
    ) {
      return { count: 0 };
    }
    if (this.loopDetected) {
      return {
        count: this.detectedCount,
        type: this.lastLoopType,
        detail: this.lastLoopDetail,
      };
    }

    this.turnsInCurrentPrompt++;

    if (
      this.turnsInCurrentPrompt >= LLM_CHECK_AFTER_TURNS &&
      this.turnsInCurrentPrompt - this.lastCheckTurn >= this.llmCheckInterval
    ) {
      this.lastCheckTurn = this.turnsInCurrentPrompt;
      const { isLoop, analysis, confirmedByModel } =
        await this.checkForLoopWithLLM(signal);
      if (isLoop) {
        this.loopDetected = true;
        this.detectedCount++;
        this.lastLoopDetail = analysis;
        this.lastLoopType = LoopType.LLM_DETECTED_LOOP;

        logLoopDetected(
          this.context.config,
          new LoopDetectedEvent(
            this.lastLoopType,
            this.promptId,
            this.detectedCount,
            confirmedByModel,
            analysis,
            LLM_CONFIDENCE_THRESHOLD,
          ),
        );

        return {
          count: this.detectedCount,
          type: this.lastLoopType,
          detail: this.lastLoopDetail,
          confirmedByModel,
        };
      }
    }
    return { count: 0 };
  }

  private checkToolCallLoop(toolCall: { name: string; args: object }): boolean {
    const key = this.getToolCallKey(toolCall);
    if (this.lastToolCallKey === key) {
      this.toolCallRepetitionCount++;
    } else {
      this.lastToolCallKey = key;
      this.toolCallRepetitionCount = 1;
    }
    if (this.toolCallRepetitionCount >= TOOL_CALL_LOOP_THRESHOLD) {
      return true;
    }
    return false;
  }

  /**
   * Detects content loops by analyzing streaming text for repetitive patterns.
   *
   * The algorithm works by:
   * 1. Appending new content to the streaming history
   * 2. Truncating history if it exceeds the maximum length
   * 3. Analyzing content chunks for repetitive patterns using hashing
   * 4. Detecting loops when identical chunks appear frequently within a short distance
   * 5. Disabling loop detection within code blocks to prevent false positives,
   *    as repetitive code structures are common and not necessarily loops.
   */
  private checkContentLoop(content: string): boolean {
    // Different content elements can often contain repetitive syntax that is not indicative of a loop.
    // To avoid false positives, we detect when we encounter different content types and
    // reset tracking to avoid analyzing content that spans across different element boundaries.
    const numFences = (content.match(/```/g) ?? []).length;
    const hasTable = /(^|\n)\s*(\|.*\||[|+-]{3,})/.test(content);
    const hasListItem =
      /(^|\n)\s*[*-+]\s/.test(content) || /(^|\n)\s*\d+\.\s/.test(content);
    const hasHeading = /(^|\n)#+\s/.test(content);
    const hasBlockquote = /(^|\n)>\s/.test(content);
    const isDivider = /^[+-_=*\u2500-\u257F]+$/.test(content);

    if (
      numFences ||
      hasTable ||
      hasListItem ||
      hasHeading ||
      hasBlockquote ||
      isDivider
    ) {
      // Reset tracking when different content elements are detected to avoid analyzing content
      // that spans across different element boundaries.
      this.resetContentTracking();
    }

    const wasInCodeBlock = this.inCodeBlock;
    this.inCodeBlock =
      numFences % 2 === 0 ? this.inCodeBlock : !this.inCodeBlock;
    if (wasInCodeBlock || this.inCodeBlock || isDivider) {
      return false;
    }

    this.streamContentHistory += content;

    this.truncateAndUpdate();
    return this.analyzeContentChunksForLoop();
  }

  /**
   * Truncates the content history to prevent unbounded memory growth.
   * When truncating, adjusts all stored indices to maintain their relative positions.
   */
  private truncateAndUpdate(): void {
    if (this.streamContentHistory.length <= MAX_HISTORY_LENGTH) {
      return;
    }

    // Calculate how much content to remove from the beginning
    const truncationAmount =
      this.streamContentHistory.length - MAX_HISTORY_LENGTH;
    this.streamContentHistory =
      this.streamContentHistory.slice(truncationAmount);
    this.lastContentIndex = Math.max(
      0,
      this.lastContentIndex - truncationAmount,
    );

    // Update all stored chunk indices to account for the truncation
    for (const [hash, oldIndices] of this.contentStats.entries()) {
      const adjustedIndices = oldIndices
        .map((index) => index - truncationAmount)
        .filter((index) => index >= 0);

      if (adjustedIndices.length > 0) {
        this.contentStats.set(hash, adjustedIndices);
      } else {
        this.contentStats.delete(hash);
      }
    }
  }

  /**
   * Analyzes content in fixed-size chunks to detect repetitive patterns.
   *
   * Uses a sliding window approach:
   * 1. Extract chunks of fixed size (CONTENT_CHUNK_SIZE)
   * 2. Hash each chunk for efficient comparison
   * 3. Track positions where identical chunks appear
   * 4. Detect loops when chunks repeat frequently within a short distance
   */
  private analyzeContentChunksForLoop(): boolean {
    while (this.hasMoreChunksToProcess()) {
      // Extract current chunk of text
      const currentChunk = this.streamContentHistory.substring(
        this.lastContentIndex,
        this.lastContentIndex + CONTENT_CHUNK_SIZE,
      );
      const chunkHash = createHash('sha256').update(currentChunk).digest('hex');

      if (this.isLoopDetectedForChunk(currentChunk, chunkHash)) {
        return true;
      }

      // Move to next position in the sliding window
      this.lastContentIndex++;
    }

    return false;
  }

  private hasMoreChunksToProcess(): boolean {
    return (
      this.lastContentIndex + CONTENT_CHUNK_SIZE <=
      this.streamContentHistory.length
    );
  }

  /**
   * Determines if a content chunk indicates a loop pattern.
   *
   * Loop detection logic:
   * 1. Check if we've seen this hash before (new chunks are stored for future comparison)
   * 2. Verify actual content matches to prevent hash collisions
   * 3. Track all positions where this chunk appears
   * 4. A loop is detected when the same chunk appears CONTENT_LOOP_THRESHOLD times
   *    within a small average distance (≤ 5 * chunk size)
   */
  private isLoopDetectedForChunk(chunk: string, hash: string): boolean {
    const existingIndices = this.contentStats.get(hash);

    if (!existingIndices) {
      this.contentStats.set(hash, [this.lastContentIndex]);
      return false;
    }

    if (!this.isActualContentMatch(chunk, existingIndices[0])) {
      return false;
    }

    existingIndices.push(this.lastContentIndex);

    if (existingIndices.length < CONTENT_LOOP_THRESHOLD) {
      return false;
    }

    // Analyze the most recent occurrences to see if they're clustered closely together
    const recentIndices = existingIndices.slice(-CONTENT_LOOP_THRESHOLD);
    const totalDistance =
      recentIndices[recentIndices.length - 1] - recentIndices[0];
    const averageDistance = totalDistance / (CONTENT_LOOP_THRESHOLD - 1);
    const maxAllowedDistance = CONTENT_CHUNK_SIZE * 5;

    if (averageDistance > maxAllowedDistance) {
      return false;
    }

    // Verify that the sequence is actually repeating, not just sharing a common prefix.
    // For a true loop, the text between occurrences of the chunk (the period) should be highly repetitive.
    const periods = new Set<string>();
    for (let i = 0; i < recentIndices.length - 1; i++) {
      periods.add(
        this.streamContentHistory.substring(
          recentIndices[i],
          recentIndices[i + 1],
        ),
      );
    }

    // If the periods are mostly unique, it's a list of distinct items with a shared prefix.
    // A true loop will have a small number of unique periods (usually 1, sometimes 2 or 3).
    // We use Math.floor(CONTENT_LOOP_THRESHOLD / 2) as a safe threshold.
    if (periods.size > Math.floor(CONTENT_LOOP_THRESHOLD / 2)) {
      return false;
    }

    return true;
  }

  /**
   * Verifies that two chunks with the same hash actually contain identical content.
   * This prevents false positives from hash collisions.
   */
  private isActualContentMatch(
    currentChunk: string,
    originalIndex: number,
  ): boolean {
    const originalChunk = this.streamContentHistory.substring(
      originalIndex,
      originalIndex + CONTENT_CHUNK_SIZE,
    );
    return originalChunk === currentChunk;
  }

  private trimRecentHistory(history: Content[]): Content[] {
    // A function response must be preceded by a function call.
    // Continuously removes dangling function calls from the end of the history
    // until the last turn is not a function call.
    while (history.length > 0 && isFunctionCall(history[history.length - 1])) {
      history.pop();
    }

    // A function response should follow a function call.
    // Continuously removes leading function responses from the beginning of history
    // until the first turn is not a function response.
    while (history.length > 0 && isFunctionResponse(history[0])) {
      history.shift();
    }

    return history;
  }

  private async checkForLoopWithLLM(signal: AbortSignal): Promise<{
    isLoop: boolean;
    analysis?: string;
    confirmedByModel?: string;
  }> {
    const recentHistory = this.context.geminiClient
      .getHistory()
      .slice(-LLM_LOOP_CHECK_HISTORY_COUNT);

    const trimmedHistory = this.trimRecentHistory(recentHistory);

    const taskPrompt = `Please analyze the conversation history to determine the possibility that the conversation is stuck in a repetitive, non-productive state. Consider the original user request when evaluating whether repeated tool calls represent legitimate batch work or an actual loop. Provide your response in the requested JSON format.`;

    const contents = [
      ...(this.userPrompt
        ? [
            {
              role: 'user' as const,
              parts: [
                {
                  text: `<original_user_request>\n${this.userPrompt}\n</original_user_request>`,
                },
              ],
            },
          ]
        : []),
      ...trimmedHistory,
      { role: 'user', parts: [{ text: taskPrompt }] },
    ];
    if (contents.length > 0 && isFunctionCall(contents[0])) {
      contents.unshift({
        role: 'user',
        parts: [{ text: 'Recent conversation history:' }],
      });
    }

    const flashResult = await this.queryLoopDetectionModel(
      'loop-detection',
      contents,
      signal,
    );

    if (!flashResult) {
      return { isLoop: false };
    }

    const flashConfidence =
      // eslint-disable-next-line no-restricted-syntax
      typeof flashResult['unproductive_state_confidence'] === 'number'
        ? flashResult['unproductive_state_confidence']
        : 0;
    const flashAnalysis =
      // eslint-disable-next-line no-restricted-syntax
      typeof flashResult['unproductive_state_analysis'] === 'string'
        ? flashResult['unproductive_state_analysis']
        : '';

    const doubleCheckModelName =
      this.context.config.modelConfigService.getResolvedConfig({
        model: DOUBLE_CHECK_MODEL_ALIAS,
      }).model;

    if (flashConfidence < LLM_CONFIDENCE_THRESHOLD) {
      logLlmLoopCheck(
        this.context.config,
        new LlmLoopCheckEvent(
          this.promptId,
          flashConfidence,
          doubleCheckModelName,
          -1,
        ),
      );
      this.updateCheckInterval(flashConfidence);
      return { isLoop: false };
    }

    const availability = this.context.config.getModelAvailabilityService();

    if (!availability.snapshot(doubleCheckModelName).available) {
      const flashModelName =
        this.context.config.modelConfigService.getResolvedConfig({
          model: 'loop-detection',
        }).model;
      return {
        isLoop: true,
        analysis: flashAnalysis,
        confirmedByModel: flashModelName,
      };
    }

    // Double check with configured model
    const mainModelResult = await this.queryLoopDetectionModel(
      DOUBLE_CHECK_MODEL_ALIAS,
      contents,
      signal,
    );

    const mainModelConfidence =
      mainModelResult &&
      // eslint-disable-next-line no-restricted-syntax
      typeof mainModelResult['unproductive_state_confidence'] === 'number'
        ? mainModelResult['unproductive_state_confidence']
        : 0;
    const mainModelAnalysis =
      mainModelResult &&
      // eslint-disable-next-line no-restricted-syntax
      typeof mainModelResult['unproductive_state_analysis'] === 'string'
        ? mainModelResult['unproductive_state_analysis']
        : undefined;

    logLlmLoopCheck(
      this.context.config,
      new LlmLoopCheckEvent(
        this.promptId,
        flashConfidence,
        doubleCheckModelName,
        mainModelConfidence,
      ),
    );

    if (mainModelResult) {
      if (mainModelConfidence >= LLM_CONFIDENCE_THRESHOLD) {
        return {
          isLoop: true,
          analysis: mainModelAnalysis,
          confirmedByModel: doubleCheckModelName,
        };
      } else {
        this.updateCheckInterval(mainModelConfidence);
      }
    }

    return { isLoop: false };
  }

  private async queryLoopDetectionModel(
    model: string,
    contents: Content[],
    signal: AbortSignal,
  ): Promise<Record<string, unknown> | null> {
    try {
      const result = await this.context.config.getBaseLlmClient().generateJson({
        modelConfigKey: { model },
        contents,
        schema: LOOP_DETECTION_SCHEMA,
        systemInstruction: LOOP_DETECTION_SYSTEM_PROMPT,
        abortSignal: signal,
        promptId: this.promptId,
        maxAttempts: 2,
        role: LlmRole.UTILITY_LOOP_DETECTOR,
      });

      if (
        result &&
        // eslint-disable-next-line no-restricted-syntax
        typeof result['unproductive_state_confidence'] === 'number'
      ) {
        return result;
      }
      return null;
    } catch (error) {
      if (this.context.config.getDebugMode()) {
        debugLogger.warn(
          `Error querying loop detection model (${model}): ${String(error)}`,
        );
      }
      return null;
    }
  }

  private updateCheckInterval(unproductive_state_confidence: number): void {
    this.llmCheckInterval = Math.round(
      MIN_LLM_CHECK_INTERVAL +
        (MAX_LLM_CHECK_INTERVAL - MIN_LLM_CHECK_INTERVAL) *
          (1 - unproductive_state_confidence),
    );
  }

  /**
   * Resets all loop detection state.
   */
  reset(promptId: string, userPrompt?: string): void {
    this.promptId = promptId;
    this.userPrompt = userPrompt ?? '';
    this.resetToolCallCount();
    this.resetContentTracking();
    this.resetLlmCheckTracking();
    this.loopDetected = false;
    this.detectedCount = 0;
    this.lastLoopDetail = undefined;
    this.lastLoopType = undefined;
  }

  /**
   * Resets the loop detected flag to allow a recovery turn to proceed.
   * This preserves the detectedCount so that the next detection will be count 2.
   */
  clearDetection(): void {
    this.loopDetected = false;
  }

  private resetToolCallCount(): void {
    this.lastToolCallKey = null;
    this.toolCallRepetitionCount = 0;
  }

  private resetContentTracking(resetHistory = true): void {
    if (resetHistory) {
      this.streamContentHistory = '';
    }
    this.contentStats.clear();
    this.lastContentIndex = 0;
  }

  private resetLlmCheckTracking(): void {
    this.turnsInCurrentPrompt = 0;
    this.llmCheckInterval = DEFAULT_LLM_CHECK_INTERVAL;
    this.lastCheckTurn = 0;
  }
}
