/**
 * JSON parsing with fallback strategies for AI-generated responses.
 */

import { jsonrepair } from 'jsonrepair';
import { noopGenerationLogger, type GenerationLogger } from './logger.js';

export interface JsonParsingOptions {
  logger?: GenerationLogger;
}

function repairQuotedPropertyFragments(jsonStr: string): string {
  return jsonStr.replace(
    /([,{]\s*)"([A-Za-z_][A-Za-z0-9_]*)\s*:\s*(true|false|null|[+-]?\d+(?:\.\d+)?)"(?=\s*[,}])/g,
    (_match, prefix, key, value) => `${prefix}"${key}": ${value}`,
  );
}

function logJsonParseError(
  stage: string,
  jsonStr: string,
  error: unknown,
  logger: GenerationLogger,
): void {
  const message = error instanceof Error ? error.message : String(error);
  const positionMatch = message.match(/position\s+(\d+)/i);
  const position = positionMatch ? Number(positionMatch[1]) : undefined;

  if (typeof position === 'number' && Number.isFinite(position)) {
    const start = Math.max(0, position - 120);
    const end = Math.min(jsonStr.length, position + 120);
    logger.warn(
      `${stage} parse error at position ${position}: ${message}. Context: ${jsonStr
        .slice(start, end)
        .replace(/\n/g, '\\n')}`,
    );
    return;
  }

  logger.warn(`${stage} parse error: ${message}`);
}

export function parseJsonResponse<T>(response: string, options: JsonParsingOptions = {}): T | null {
  const logger = options.logger ?? noopGenerationLogger;
  const exactParsed = tryParseExactJson<T>(response);
  if (exactParsed !== null) return exactParsed;

  const cleanedResponse = stripReasoningPrefix(response);
  if (cleanedResponse !== response.trim()) {
    const parsedCleaned = parseJsonResponseCandidate<T>(cleanedResponse, logger);
    if (parsedCleaned !== null) return parsedCleaned;
  }

  const parsed = parseJsonResponseCandidate<T>(response, logger);
  if (parsed !== null) return parsed;

  logger.error('Failed to parse JSON from response');
  logger.error('Raw response (first 500 chars):', cleanedResponse.substring(0, 500));
  logger.error(
    'Raw response (last 500 chars):',
    cleanedResponse.substring(Math.max(0, cleanedResponse.length - 500)),
  );

  return null;
}

function tryParseExactJson<T>(response: string): T | null {
  try {
    return JSON.parse(response.trim()) as T;
  } catch {
    return null;
  }
}

function stripReasoningPrefix(response: string): string {
  const trimmed = response.trim();
  const matches = [...trimmed.matchAll(/<\/(?:think|thinking|reasoning)>\s*/gi)];
  const lastMatch = matches.at(-1);

  if (!lastMatch || lastMatch.index === undefined) return trimmed;

  return trimmed.slice(lastMatch.index + lastMatch[0].length).trim();
}

function parseJsonResponseCandidate<T>(response: string, logger: GenerationLogger): T | null {
  const cleanedResponse = response.trim();

  // Strategy 1: Try to extract JSON from markdown code blocks (may have multiple)
  const codeBlockMatches = cleanedResponse.matchAll(/```(?:json)?\s*([\s\S]*?)```/g);
  for (const match of codeBlockMatches) {
    const extracted = match[1].trim();
    // Only try if it looks like JSON (starts with { or [)
    if (extracted.startsWith('{') || extracted.startsWith('[')) {
      const result = tryParseJson<T>(extracted, { logger });
      if (result !== null) {
        logger.debug('Successfully parsed JSON from code block');
        return result;
      }
    }
  }

  // Strategy 2: Try to find JSON structure directly in response (no code block)
  // Look for array or object start
  const jsonStartArray = cleanedResponse.indexOf('[');
  const jsonStartObject = cleanedResponse.indexOf('{');

  if (jsonStartArray !== -1 || jsonStartObject !== -1) {
    // Prefer the structure that appears first
    const startIndex =
      jsonStartArray === -1
        ? jsonStartObject
        : jsonStartObject === -1
          ? jsonStartArray
          : Math.min(jsonStartArray, jsonStartObject);

    // Find the matching close bracket
    let depth = 0;
    let endIndex = -1;
    let inString = false;
    let escapeNext = false;

    for (let i = startIndex; i < cleanedResponse.length; i++) {
      const char = cleanedResponse[i];

      if (escapeNext) {
        escapeNext = false;
        continue;
      }

      if (char === '\\' && inString) {
        escapeNext = true;
        continue;
      }

      // escapeNext is always false here: the first branch of the loop resets
      // and skips the iteration whenever it was set.
      if (char === '"') {
        inString = !inString;
        continue;
      }

      if (!inString) {
        if (char === '[' || char === '{') depth++;
        else if (char === ']' || char === '}') {
          depth--;
          if (depth === 0) {
            endIndex = i;
            break;
          }
        }
      }
    }

    if (endIndex !== -1) {
      const jsonStr = cleanedResponse.substring(startIndex, endIndex + 1);
      const result = tryParseJson<T>(jsonStr, { logger });
      if (result !== null) {
        logger.debug('Successfully parsed JSON from response body');
        return result;
      }
    }
  }

  // Strategy 3: Last resort - try the whole response
  const result = tryParseJson<T>(cleanedResponse.trim(), { logger });
  if (result !== null) {
    logger.debug('Successfully parsed raw response as JSON');
    return result;
  }

  return null;
}

/**
 * Try to parse JSON with various fixes for common AI response issues
 */
export function tryParseJson<T>(jsonStr: string, options: JsonParsingOptions = {}): T | null {
  const logger = options.logger ?? noopGenerationLogger;
  // Attempt 1: Try parsing as-is
  try {
    return JSON.parse(jsonStr) as T;
  } catch (error) {
    logJsonParseError('Attempt 1', jsonStr, error, logger);
    // Continue to fix attempts
  }

  // Attempt 2: Fix common JSON issues from AI responses
  try {
    let fixed = jsonStr;

    // Fix 0: Recover malformed property fragments that were accidentally
    // emitted as standalone strings inside an object, such as:
    // `"height: 76"` -> `"height": 76`
    // `"fixedRatio: false"` -> `"fixedRatio": false`
    // The object-context prefix/suffix guards keep valid JSON strings intact.
    fixed = repairQuotedPropertyFragments(fixed);

    // Fix 1: Handle LaTeX-style escapes that break JSON (e.g., \frac, \left, \right, \times, etc.)
    // These are common in math content and need to be double-escaped
    // Match backslash followed by letters (LaTeX commands) inside strings,
    // but skip valid JSON escape sequences (\b, \f, \n, \r, \t, \u)
    fixed = fixed.replace(/"([^"\\]*(?:\\.[^"\\]*)*)"/g, (_match, content) => {
      // Double-escape backslash+letter ONLY for non-JSON-escape letters
      const fixedContent = content.replace(/\\([a-zA-Z])/g, (_m: string, ch: string) => {
        // Preserve valid JSON escape sequences
        if ('bfnrtu'.includes(ch)) return `\\${ch}`;
        return `\\\\${ch}`;
      });
      return `"${fixedContent}"`;
    });

    // Fix 2: Fix other invalid escape sequences (e.g., \S, \L, etc.)
    // Valid JSON escapes: \", \\, \/, \b, \f, \n, \r, \t, \uXXXX
    fixed = fixed.replace(/\\([^"\\\/bfnrtu\n\r])/g, (match, char) => {
      // If it's a letter, it's likely a LaTeX command
      if (/[a-zA-Z]/.test(char)) {
        return '\\\\' + char;
      }
      return match;
    });

    // Fix 3: Try to fix truncated JSON arrays/objects
    const trimmed = fixed.trim();
    if (trimmed.startsWith('[') && !trimmed.endsWith(']')) {
      const lastCompleteObj = fixed.lastIndexOf('}');
      if (lastCompleteObj > 0) {
        fixed = fixed.substring(0, lastCompleteObj + 1) + ']';
        logger.warn('Fixed truncated JSON array');
      }
    } else if (trimmed.startsWith('{') && !trimmed.endsWith('}')) {
      // Try to close incomplete object
      const openBraces = (fixed.match(/{/g) || []).length;
      const closeBraces = (fixed.match(/}/g) || []).length;
      if (openBraces > closeBraces) {
        fixed += '}'.repeat(openBraces - closeBraces);
        logger.warn('Fixed truncated JSON object');
      }
    }

    return JSON.parse(fixed) as T;
  } catch (error) {
    logJsonParseError('Attempt 2', jsonStr, error, logger);
    // Continue to next attempt
  }

  // Attempt 3: Use jsonrepair to fix malformed JSON (e.g. unescaped quotes in Chinese text)
  try {
    const repaired = jsonrepair(jsonStr);
    return JSON.parse(repaired) as T;
  } catch (error) {
    logJsonParseError('Attempt 3', jsonStr, error, logger);
    // Continue to next attempt
  }

  // Attempt 4: More aggressive fixing - remove control characters
  try {
    let fixed = jsonStr;

    // Remove or escape control characters
    fixed = fixed.replace(/[\x00-\x1F\x7F]/g, (char) => {
      switch (char) {
        case '\n':
          return '\\n';
        case '\r':
          return '\\r';
        case '\t':
          return '\\t';
        default:
          return '';
      }
    });

    return JSON.parse(fixed) as T;
  } catch (error) {
    logJsonParseError('Attempt 4', jsonStr, error, logger);
    return null;
  }
}