terminal

import { tool } from "@lmstudio/sdk";
import { z } from "zod";
import { readFile, writeFileSync } from "../fileUtilities";
import { sanitizeFilePath, validateWritePath } from "../securityEnhanced";
import { logger } from "../logger";

// ── Regex Utilities ──
function escapeRegExp(str: string): string {
  return str.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}

/**
 * Count all matches in a line using an existing global regex.
 * O(n) where n = number of matches. Resets lastIndex for each call.
 */
function countMatchesInLine(line: string, regex: RegExp): number {
  let count = 0;
  regex.lastIndex = 0; // Reset for this independent line
  let match: RegExpExecArray | null;
  while ((match = regex.exec(line)) !== null) {
    count++;
    if (match[0].length === 0) regex.lastIndex++; // Avoid infinite loop on zero-length matches
  }
  return count;
}

/**
 * Find and replace up to `max` occurrences in a line.
 * O(k * m) where k = max replacements, m = position of k-th match.
 * Stops immediately after reaching max (more efficient than .replace() which scans entire string).
 */
function findAndReplace(line: string, regex: RegExp, replacement: string, max: number): { replaced: string; count: number } {
  let count = 0;
  const parts: string[] = [];
  let lastIndex = 0;

  while (count < max) {
    regex.lastIndex = lastIndex; // Reset to continue from where we left off
    const match = regex.exec(line);
    if (!match || match.index === undefined) break;

    parts.push(line.slice(lastIndex, match.index));
    parts.push(replacement);

    count++;
    lastIndex = match.index + match[0].length;
    if (match[0].length === 0) lastIndex++; // Avoid infinite loop on zero-length matches
  }

  parts.push(line.slice(lastIndex));
  return { replaced: parts.join(""), count };
}

/**
 * Process a range of lines for replacement, applying pattern up to maxReplacements.
 */
function processReplacementLines(
  lines: string[],
  fromLine: number,
  toLine: number,
  pattern: RegExp,
  replacement: string,
  maxReplacements?: number
): { count: number; result: string[] } {
  let replacementsMade = 0;
  const result = [...lines]; // Shallow copy — we only mutate specific indices

  for (let i = fromLine; i < Math.min(toLine, lines.length); i++) {
    if (maxReplacements !== undefined && replacementsMade >= maxReplacements) break;

    const remaining = maxReplacements !== undefined ? maxReplacements - replacementsMade : Infinity;
    const { count } = findAndReplace(result[i], pattern, replacement, remaining);

    if (count > 0) {
      result[i] = findAndReplace(result[i], pattern, replacement, remaining).replaced;
      replacementsMade += count;
    }
  }

  return { count: replacementsMade, result };
}

/**
 * Build a context snippet for display after changes.
 */
function buildChangeSnippet(lines: string[], fromLine: number, toLine: number): string {
  const snippets: string[] = [];
  const start = Math.max(0, fromLine - 2);
  const end = Math.min(lines.length, toLine + 3);

  for (let i = start; i < end; i++) {
    const lineNum = i + 1;
    const prefix = lineNum < 10 ? " " : "";
    snippets.push(`${prefix}${lineNum}: ${lines[i]}`);
  }
  return snippets.join("\n");
}

// ── Tool Registration ──
export async function registerTextTools(): Promise<any[]> {
  const tools: any[] = [];

  // ── replace_text ──
  tools.push(tool({
    name: "replace_text",
    description: "Find and replace text within a file. Supports literal strings and regex patterns.",
    parameters: {
      file_path: z.string().describe("Path to the file to modify"),
      find: z.string().describe("Text or regex pattern to find"),
      replace_with: z.string().describe("Text to replace matches with"),
      regex: z.boolean().optional().describe("Treat 'find' as a regex pattern (default: false)"),
      case_sensitive: z.boolean().optional().describe("Case sensitivity (default: true for regex, true for literal)"),
      max_replacements: z.number().optional().describe("Maximum number of replacements to make (default: all)"),
      start_line: z.number().optional().describe("Line number to start replacements from (1-based, default: 1)"),
      end_line: z.number().optional().describe("Line number to stop replacements at (1-based, default: end of file)"),
      create_backup: z.boolean().optional().describe("Create a .bak backup before modifying (default: false)"),
    },
    implementation: async ({ file_path, find, replace_with, regex, case_sensitive, max_replacements, start_line, end_line, create_backup }) => {
      try {
        const validation = validateWritePath(file_path);
        if (!validation.valid) return { success: false, replacements_made: 0, error: validation.error };
        const safePath = validation.sanitized!;

        const readResult = readFile(safePath);
        if (!readResult) return { success: false, replacements_made: 0, error: `Failed to read file: ${safePath}` };

        let content = readResult;

        if (create_backup) writeFileSync(safePath + ".bak", content);

        const lines = content.split("\n");
        const fromLine = (start_line ?? 1) - 1;
        const effectiveToLine = end_line !== undefined ? Math.min(end_line, lines.length) : lines.length;

        if (fromLine < 0 || fromLine >= lines.length) {
          return { success: false, replacements_made: 0, error: `Invalid start_line: ${start_line}. File has ${lines.length} lines.` };
        }

        const caseFlag = (case_sensitive ?? true) ? "" : "i";
        let pattern: RegExp;

        if (regex) {
          try {
            pattern = new RegExp(find, caseFlag + "g");
          } catch (err: any) {
            return { success: false, replacements_made: 0, error: `Invalid regex pattern '${find}': ${err.message}` };
          }
        } else {
          const escaped = escapeRegExp(find);
          pattern = new RegExp(escaped, caseFlag + "g");
        }

        // Use findAndReplace directly instead of replaceInLine (more efficient)
        let replacementsMade = 0;
        for (let i = fromLine; i < effectiveToLine; i++) {
          if (max_replacements !== undefined && replacementsMade >= max_replacements) break;

          const remaining = max_replacements !== undefined ? Math.max(0, max_replacements - replacementsMade) : Infinity;
          const { replaced, count } = findAndReplace(lines[i], pattern, replace_with, remaining);
          
          if (count > 0) {
            lines[i] = replaced;
            replacementsMade += count;
          }
        }

        writeFileSync(safePath, lines.join("\n"));

        return {
          success: true,
          replacements_made: replacementsMade,
          file_path: safePath,
          file_size: content.length,
          changed: replacementsMade > 0,
          snippet: replacementsMade > 0 ? buildChangeSnippet(lines, fromLine, effectiveToLine) : undefined,
        };
      } catch (err: any) {
        logger.error('replace_text failed', { error: err.message });
        return { success: false, replacements_made: 0, error: err.message };
      }
    },
  }));

  // ── count_text ──
  tools.push(tool({
    name: "count_text",
    description: "Count occurrences of text or regex patterns in a file.",
    parameters: {
      file_path: z.string().describe("Path to the file to search"),
      pattern: z.string().describe("Text or regex pattern to count"),
      regex: z.boolean().optional().describe("Treat 'pattern' as a regex pattern (default: false)"),
      case_sensitive: z.boolean().optional().describe("Case sensitivity (default: true)"),
      start_line: z.number().optional().describe("Line number to start counting from (1-based, default: 1)"),
      end_line: z.number().optional().describe("Line number to stop counting at (1-based, default: end of file)"),
    },
    implementation: async ({ file_path, pattern, regex, case_sensitive, start_line, end_line }) => {
      try {
        const validation = sanitizeFilePath(file_path);
        if (!validation.valid) return { success: false, count: 0, error: validation.error };
        const safePath = validation.sanitized!;

        const readResult = readFile(safePath);
        if (!readResult) return { success: false, count: 0, error: `Failed to read file: ${safePath}` };

        const lines = readResult.split("\n");
        const fromLine = (start_line ?? 1) - 1;
        const effectiveToLine = end_line !== undefined ? Math.min(end_line, lines.length) : lines.length;

        if (fromLine < 0 || fromLine >= lines.length) {
          return { success: false, count: 0, error: `Invalid start_line: ${start_line}. File has ${lines.length} lines.` };
        }

        const caseFlag = (case_sensitive ?? true) ? "" : "i";
        let regExp: RegExp;

        if (regex) {
          try {
            regExp = new RegExp(pattern, caseFlag + "g");
          } catch (err: any) {
            return { success: false, count: 0, error: `Invalid regex pattern '${pattern}': ${err.message}` };
          }
        } else {
          const escaped = escapeRegExp(pattern);
          regExp = new RegExp(escaped, caseFlag + "g");
        }

        let totalCount = 0;
        const matchLines: number[] = [];

        for (let i = fromLine; i < effectiveToLine; i++) {
          const count = countMatchesInLine(lines[i], regExp);
          if (count > 0) {
            totalCount += count;
            if (matchLines.length < 20) matchLines.push(i + 1);
          }
        }

        return {
          success: true,
          count: totalCount,
          pattern,
          regex_used: regex ?? false,
          case_sensitive: case_sensitive ?? true,
          match_lines: matchLines,
          match_lines_truncated: matchLines.length >= 20,
          total_lines_scanned: effectiveToLine - fromLine,
        };
      } catch (err: any) {
        logger.error('count_text failed', { error: err.message });
        return { success: false, count: 0, error: err.message };
      }
    },
  }));

  // ── insert_text ──
  tools.push(tool({
    name: "insert_text",
    description: "Insert text at a specific line in a file.",
    parameters: {
      file_path: z.string().describe("Path to the file to modify"),
      insert_at_line: z.number().describe("Line number to insert before (1-based)"),
      content: z.string().describe("Text content to insert (may contain newlines)"),
    },
    implementation: async ({ file_path, insert_at_line, content }) => {
      try {
        const validation = validateWritePath(file_path);
        if (!validation.valid) return { success: false, error: validation.error };
        const safePath = validation.sanitized!;

        const readResult = readFile(safePath);
        if (!readResult) return { success: false, error: `Failed to read file: ${safePath}` };

        const lines = readResult.split("\n");
        let insertIndex: number;

        if (insert_at_line === 0) {
          insertIndex = 0;
        } else if (insert_at_line > lines.length) {
          insertIndex = lines.length;
        } else {
          insertIndex = insert_at_line - 1;
        }

        if (insertIndex < 0 || insertIndex > lines.length) {
          return { success: false, error: `Invalid insert_at_line: ${insert_at_line}. File has ${lines.length} lines.` };
        }

        const insertedLines = content.split("\n");
        lines.splice(insertIndex, 0, ...insertedLines);

        writeFileSync(safePath, lines.join("\n"));

        return {
          success: true,
          file_path: safePath,
          lines_inserted: insertedLines.length,
          new_total_lines: lines.length,
          insertion_point: insertIndex + 1,
        };
      } catch (err: any) {
        logger.error('insert_text failed', { error: err.message });
        return { success: false, error: err.message };
      }
    },
  }));

  // ── delete_lines ──
  tools.push(tool({
    name: "delete_lines",
    description: "Delete a range of lines from a file.",
    parameters: {
      file_path: z.string().describe("Path to the file to modify"),
      start_line: z.number().describe("First line to delete (1-based)"),
      end_line: z.number().describe("Last line to delete (1-based)"),
    },
    implementation: async ({ file_path, start_line, end_line }) => {
      try {
        const validation = validateWritePath(file_path);
        if (!validation.valid) return { success: false, lines_deleted: 0, error: validation.error };
        const safePath = validation.sanitized!;

        const readResult = readFile(safePath);
        if (!readResult) return { success: false, lines_deleted: 0, error: `Failed to read file: ${safePath}` };

        const lines = readResult.split("\n");
        const totalLines = lines.length;

        const from = Math.max(1, start_line);
        const to = Math.min(end_line, totalLines);

        if (from > totalLines || to < 1 || from > to) {
          return { success: false, lines_deleted: 0, error: `Invalid line range: ${start_line}-${end_line}. File has ${totalLines} lines.` };
        }

        const deletedCount = to - (from - 1);
        lines.splice(from - 1, deletedCount);

        writeFileSync(safePath, lines.join("\n"));

        // Build context snippet around deletion point
        const contextStart = Math.max(0, from - 4);
        const contextEnd = Math.min(lines.length, from + 3);
        const contextLines: string[] = [];
        for (let i = contextStart; i < contextEnd; i++) {
          const lineNum = i + 1;
          const prefix = lineNum < 10 ? " " : "";
          contextLines.push(`${prefix}${lineNum}: ${lines[i]}`);
        }

        return {
          success: true,
          lines_deleted: deletedCount,
          file_path: safePath,
          new_total_lines: lines.length,
          context: contextLines.join("\n"),
        };
      } catch (err: any) {
        logger.error('delete_lines failed', { error: err.message });
        return { success: false, lines_deleted: 0, error: err.message };
      }
    },
  }));

  return tools;
}