import { initializeParser } from './parser';
import {
  handleBeginOfFilePatch,
  handleEndOfFilePatch,
  handleStandardPatch,
} from './patcher';
import {
  SPECIAL_PATCH_BEGIN_FILE_MARKER,
  SPECIAL_PATCH_END_FILE_MARKER,
} from '../../src/llm-orchestration/parser/parsing.constants';
import { TokenizerStrategy } from './tokenizer.interface';
import { TreeSitterTokenizer } from './strategies/tree-sitter-tokenizer';
import { TiktokenTokenizer } from './strategies/tiktoken-tokenizer';

export interface ApplyPatchOptions {
  grammarPath?: string;
  useTiktoken?: boolean;
}

export async function applySnippetPatch(
  sourceCode: string,
  patchCode: string,
  optionsOrGrammarPath: string | ApplyPatchOptions,
): Promise<string> {
  let strategy: TokenizerStrategy;
  let options: ApplyPatchOptions;

  // Backward compatibility for when the 3rd argument was just grammarPath string
  if (typeof optionsOrGrammarPath === 'string') {
    options = { grammarPath: optionsOrGrammarPath };
  } else {
    options = optionsOrGrammarPath;
  }

  if (options.useTiktoken) {
    strategy = new TiktokenTokenizer();
  } else if (options.grammarPath) {
    const parser = await initializeParser(options.grammarPath);
    strategy = new TreeSitterTokenizer(parser);
  } else {
    // Default to Tiktoken if no grammar path is provided
    strategy = new TiktokenTokenizer();
  }

  const sourceTokens = strategy.tokenize(sourceCode);

  let patchResult: {
    replaceStart: number;
    replaceEnd: number;
    patchInsertStart: number;
    patchInsertEnd: number;
  };
  let processedPatchCode = patchCode;

  const beginOfFileRegex = new RegExp(
    `//\\s*${SPECIAL_PATCH_BEGIN_FILE_MARKER}.*`,
  );
  const endOfFileRegex = new RegExp(`//\\s*${SPECIAL_PATCH_END_FILE_MARKER}.*`);

  const hasBeginOfFile = beginOfFileRegex.test(patchCode);
  const hasEndOfFile = endOfFileRegex.test(patchCode);

  if (hasBeginOfFile) {
    processedPatchCode = patchCode.replace(beginOfFileRegex, '');
    // We need to tokenize the processed patch code
    // We filter out empty text tokens if any strategy produces them, though likely not needed for tiktoken
    const patchTokens = strategy
      .tokenize(processedPatchCode)
      .filter((t) => t.text !== '' && t.text !== '\n');
    patchResult = handleBeginOfFilePatch(sourceTokens, patchTokens);
  } else if (hasEndOfFile) {
    processedPatchCode = patchCode.replace(endOfFileRegex, '');
    const patchTokens = strategy
      .tokenize(processedPatchCode)
      .filter((t) => t.text !== '' && t.text !== '\n');
    patchResult = handleEndOfFilePatch(sourceTokens, patchTokens, sourceCode);
  } else {
    processedPatchCode = patchCode.trim();
    const patchTokens = strategy
      .tokenize(processedPatchCode)
      .filter((t) => t.text !== '');

    // Heuristic: If using Tiktoken, we discard the first and last 3 tokens to improve matching resilience.
    // This prevents issues where the LLM hallucinating extra delimiters or context at the edges of the snippet.
    const isTiktoken = options.useTiktoken || !options.grammarPath;
    if (
      isTiktoken &&
      patchTokens.length >= 6 &&
      patchTokens.length < sourceTokens.length
    ) {
      const innerTokens = patchTokens.slice(2, -2);
      patchResult = handleStandardPatch(sourceTokens, innerTokens);
    } else {
      patchResult = handleStandardPatch(sourceTokens, patchTokens);
    }
  }

  // NOTE: replaceStart/End are byte offsets.
  // If your code can contain non-ASCII characters, you need to
  // map byte offsets to JS string indices. For pure ASCII, this
  // works as-is.
  const prefix = sourceCode.slice(0, patchResult.replaceStart);
  const suffix = sourceCode.slice(patchResult.replaceEnd);
  const finalPatchContent = processedPatchCode.slice(
    patchResult.patchInsertStart,
    patchResult.patchInsertEnd,
  );

  return prefix + finalPatchContent + suffix;
}
