import type { Token } from './types';
import { findAllSequences, formatAnchor } from './tokens';
import {
  SPECIAL_PATCH_BEGIN_FILE_MARKER,
  SPECIAL_PATCH_END_FILE_MARKER,
} from '../../src/llm-orchestration/parser/parsing.constants';

interface PatchResult {
  replaceStart: number;
  replaceEnd: number;
  patchInsertStart: number;
  patchInsertEnd: number;
}

interface SimplePatchResult {
  replaceStart: number;
  replaceEnd: number;
}

// Helper to prioritize "Ambiguous" errors over "Not found" errors
const updateLastError = (
  currentError: Error | null,
  newError: Error,
): Error => {
  if (!currentError) return newError;
  // If we already have an Ambiguous error, keep it unless the new one is also Ambiguous
  // (Assuming Ambiguous is more useful/specific than "Not found")
  const currentIsAmbiguous = currentError.message.includes('Ambiguous');
  const newIsAmbiguous = newError.message.includes('Ambiguous');

  if (currentIsAmbiguous && !newIsAmbiguous) {
    return currentError;
  }
  return newError;
};

// Internal helper for the original matching logic
function _findBeginOfFilePatchLocation(
  sourceTokens: Token[],
  patchTokens: Token[],
): SimplePatchResult {
  let replaceStart: number | null = null;
  let replaceEnd: number | null = null;
  let lastError = '';

  for (let anchorSize = 1; anchorSize <= patchTokens.length; anchorSize++) {
    const suffixAnchor = patchTokens.slice(patchTokens.length - anchorSize);
    const indices = findAllSequences(sourceTokens, suffixAnchor);

    if (indices.length === 1) {
      replaceStart = 0;
      const suffixIndex = indices[0];
      replaceEnd = sourceTokens[suffixIndex + anchorSize - 1].endIndex;
      break; // Found unique anchor
    }

    if (indices.length > 1) {
      const formattedAnchor = formatAnchor(suffixAnchor);
      const locations = indices
        .map(
          (i) =>
            `line ${sourceTokens[i].startPosition?.row ? sourceTokens[i].startPosition.row + 1 : '?'}`,
        )
        .join(', ');
      lastError = `Ambiguous suffix anchor. The sequence "${formattedAnchor}" was found at ${indices.length} locations: ${locations}.`;
    }
  }

  if (replaceStart === null || replaceEnd === null) {
    const smallestAnchor = patchTokens.slice(patchTokens.length - 1);
    throw new Error(
      lastError ||
        `Suffix anchor not found. The sequence "${formatAnchor(
          smallestAnchor,
        )}" could not be located.`,
    );
  }

  return { replaceStart, replaceEnd };
}

export function handleBeginOfFilePatch(
  sourceTokens: Token[],
  originalPatchTokens: Token[],
): PatchResult {
  if (originalPatchTokens.length === 0) {
    throw new Error(
      `Patch is empty after removing ${SPECIAL_PATCH_BEGIN_FILE_MARKER} marker.`,
    );
  }

  let patchAttempt = [...originalPatchTokens];
  let lastError: Error | null = null;

  // Outer loop for trimming tokens from the start
  while (patchAttempt.length >= 1) {
    try {
      const { replaceStart, replaceEnd } = _findBeginOfFilePatchLocation(
        sourceTokens,
        patchAttempt,
      );
      const patchInsertStart = patchAttempt[0].startIndex;
      const patchInsertEnd = patchAttempt[patchAttempt.length - 1].endIndex;
      return { replaceStart, replaceEnd, patchInsertStart, patchInsertEnd };
    } catch (e) {
      lastError = updateLastError(lastError, e as Error);
      patchAttempt = patchAttempt.slice(1); // Trim one token from the beginning
    }
  }

  throw new Error(
    `Failed to apply ${SPECIAL_PATCH_BEGIN_FILE_MARKER} patch. Could not find a unique anchor, even after trimming tokens. Last known error: ${lastError?.message}`,
  );
}

// Internal helper for the original matching logic
function _findEndOfFilePatchLocation(
  sourceTokens: Token[],
  patchTokens: Token[],
  sourceCode: string,
): SimplePatchResult {
  let replaceStart: number | null = null;
  let replaceEnd: number | null = null;
  let lastError = '';

  for (let anchorSize = 1; anchorSize <= patchTokens.length; anchorSize++) {
    const prefixAnchor = patchTokens.slice(0, anchorSize);
    const indices = findAllSequences(sourceTokens, prefixAnchor);

    if (indices.length === 1) {
      replaceStart = sourceTokens[indices[0]].startIndex;
      replaceEnd = sourceCode.length;
      break; // Found unique anchor
    }

    if (indices.length > 1) {
      const formattedAnchor = formatAnchor(prefixAnchor);
      const locations = indices
        .map(
          (i) =>
            `line ${sourceTokens[i].startPosition?.row ? sourceTokens[i].startPosition.row + 1 : '?'}`,
        )
        .join(', ');
      lastError = `Ambiguous prefix anchor. The sequence "${formattedAnchor}" was found at ${indices.length} locations: ${locations}.`;
    }
  }

  if (replaceStart === null || replaceEnd === null) {
    const smallestAnchor = patchTokens.slice(0, 1);
    throw new Error(
      lastError ||
        `Prefix anchor not found. The sequence "${formatAnchor(
          smallestAnchor,
        )}" could not be located.`,
    );
  }
  return { replaceStart, replaceEnd };
}

export function handleEndOfFilePatch(
  sourceTokens: Token[],
  originalPatchTokens: Token[],
  sourceCode: string,
): PatchResult {
  if (originalPatchTokens.length === 0) {
    throw new Error(
      `Patch is empty after removing ${SPECIAL_PATCH_END_FILE_MARKER} marker.`,
    );
  }

  let patchAttempt = [...originalPatchTokens];
  let lastError: Error | null = null;

  // Outer loop for trimming tokens from the end
  while (patchAttempt.length >= 1) {
    try {
      const { replaceStart, replaceEnd } = _findEndOfFilePatchLocation(
        sourceTokens,
        patchAttempt,
        sourceCode,
      );
      const patchInsertStart = patchAttempt[0].startIndex;
      const patchInsertEnd = patchAttempt[patchAttempt.length - 1].endIndex;
      return { replaceStart, replaceEnd, patchInsertStart, patchInsertEnd };
    } catch (e) {
      lastError = updateLastError(lastError, e as Error);
      patchAttempt = patchAttempt.slice(0, -1); // Trim one token from the end
    }
  }

  throw new Error(
    `Failed to apply ${SPECIAL_PATCH_END_FILE_MARKER} patch. Could not find a unique anchor, even after trimming tokens. Last known error: ${lastError?.message}`,
  );
}

interface PrefixResult {
  prefixAnchor: Token[];
  prefixIndex: number;
  replaceStart: number;
}

function _findPrefixLocation(
  sourceTokens: Token[],
  patchTokens: Token[],
): PrefixResult {
  let prefixAnchor: Token[] | null = null;
  let prefixIndex: number | null = null;
  let bestPrefixError: string | null = null;

  for (let prefixSize = 1; prefixSize < patchTokens.length; prefixSize++) {
    const currentPrefix = patchTokens.slice(0, prefixSize);
    const prefixIndices = findAllSequences(sourceTokens, currentPrefix);

    if (prefixIndices.length === 1) {
      prefixAnchor = currentPrefix;
      prefixIndex = prefixIndices[0];
      bestPrefixError = null; // Found it
      break;
    }
    if (prefixIndices.length > 1) {
      const formatted = formatAnchor(currentPrefix);
      const locations = prefixIndices
        .map(
          (i) =>
            `line ${sourceTokens[i].startPosition?.row ? sourceTokens[i].startPosition.row + 1 : '?'}`,
        )
        .join(', ');
      bestPrefixError = `Ambiguous prefix anchor. The sequence "${formatted}" was found at ${prefixIndices.length} locations: ${locations}.`;
    }
    if (prefixIndices.length === 0) {
      if (!bestPrefixError) {
        const formatted = formatAnchor(currentPrefix);
        bestPrefixError = `Prefix anchor not found: "${formatted}"`;
      }
      break; // A larger prefix will also not be found
    }
  }

  if (!prefixAnchor || prefixIndex === null) {
    throw new Error(bestPrefixError || 'Could not find a unique prefix anchor.');
  }

  const replaceStart = sourceTokens[prefixIndex].startIndex;

  return { prefixAnchor, prefixIndex, replaceStart };
}

interface SuffixResult {
  replaceEnd: number;
}

function _findSuffixLocation(
  sourceTokens: Token[],
  patchTokens: Token[],
  prefixAnchor: Token[],
  prefixIndex: number,
): SuffixResult {
  let suffixAnchor: Token[] | null = null;
  let suffixIndex: number | null = null;
  let bestSuffixError: string | null = null;
  const searchStartIndex = prefixIndex + prefixAnchor.length;
  const sourceAfterPrefix = sourceTokens.slice(searchStartIndex);

  for (
    let suffixSize = 1;
    suffixSize <= patchTokens.length - prefixAnchor.length;
    suffixSize++
  ) {
    const currentSuffix = patchTokens.slice(patchTokens.length - suffixSize);
    const suffixIndices = findAllSequences(sourceAfterPrefix, currentSuffix);

    if (suffixIndices.length === 1) {
      suffixAnchor = currentSuffix;
      suffixIndex = searchStartIndex + suffixIndices[0];
      bestSuffixError = null; // Found it
      break;
    }
    if (suffixIndices.length > 1) {
      const formatted = formatAnchor(currentSuffix);
      const locations = suffixIndices
        .map(
          (i) =>
            `line ${
              sourceTokens[searchStartIndex + i].startPosition?.row
                ? sourceTokens[searchStartIndex + i].startPosition.row + 1
                : '?'
            }`,
        )
        .join(', ');
      bestSuffixError = `Ambiguous suffix anchor. The sequence "${formatted}" was found ${suffixIndices.length} times after the prefix: ${locations}.`;
    }
  }
  if (!suffixAnchor || suffixIndex === null) {
    if (bestSuffixError) {
      throw new Error(bestSuffixError);
    }
    const prefixLocation = `line ${
      sourceTokens[prefixIndex].startPosition?.row
        ? sourceTokens[prefixIndex].startPosition.row + 1
        : '?'
    }`;
    const formattedPrefix = formatAnchor(prefixAnchor);
    const smallestSuffix = formatAnchor(
      patchTokens.slice(patchTokens.length - 1),
    );

    throw new Error(
      `Could not find a unique suffix anchor after the prefix anchor "${formattedPrefix}" (found at ${prefixLocation}). ` +
        `The smallest suffix searched for ("${smallestSuffix}") was not found after it.`,
    );
  }

  const replaceEnd =
    sourceTokens[suffixIndex + suffixAnchor.length - 1].endIndex;

  return { replaceEnd };
}

export function handleStandardPatch(
  sourceTokens: Token[],
  originalPatchTokens: Token[],
): PatchResult {
  if (originalPatchTokens.length < 2) {
    throw new Error(
      'Patch must contain at least two tokens to form a prefix and a suffix.',
    );
  }

  let startTrim = 0;
  const endTrim = 0;
  let lastError: Error | null = null;

  let prefixInfo: PrefixResult | null = null;

  // Loop 1: Find Prefix (trim from start)
  let patchAttempt = [...originalPatchTokens];
  while (patchAttempt.length >= 2) {
    try {
      prefixInfo = _findPrefixLocation(sourceTokens, patchAttempt);
      // Success finding prefix
      startTrim = originalPatchTokens.length - patchAttempt.length;
      break;
    } catch (e) {
      lastError = updateLastError(lastError, e as Error);
      // Trim one token from the start
      patchAttempt = patchAttempt.slice(1);
    }
  }

  if (!prefixInfo) {
    throw new Error(
      `Failed to apply patch. Could not find a unique prefix anchor, even after trimming tokens. Last known error: ${lastError?.message}`,
    );
  }

  // Loop 2: Find Suffix (trim from end)
  // Reset patchAttempt to start from the found prefix startTrim, but allow trimming end
  patchAttempt = originalPatchTokens.slice(startTrim);
  lastError = null; // Reset last error for suffix search phase

  while (patchAttempt.length >= prefixInfo.prefixAnchor.length + 1) {
    // Need at least prefix + 1 token? Or just prefix + suffix?
    try {
      const { replaceEnd } = _findSuffixLocation(
        sourceTokens,
        patchAttempt,
        prefixInfo.prefixAnchor,
        prefixInfo.prefixIndex,
      );
      // Success finding suffix
      const patchInsertStart = patchAttempt[0].startIndex;
      const patchInsertEnd = patchAttempt[patchAttempt.length - 1].endIndex;
      return {
        replaceStart: prefixInfo.replaceStart,
        replaceEnd,
        patchInsertStart,
        patchInsertEnd,
      };
    } catch (e) {
      lastError = updateLastError(lastError, e as Error);
      // Trim one token from the end
      patchAttempt = patchAttempt.slice(0, -1);
    }
  }

  throw new Error(
    `Failed to apply patch. Could not find a unique suffix anchor, even after trimming tokens. Last known error: ${lastError?.message}`,
  );
}