import { detectLatexLineEnding, type LatexLineEnding } from './line-endings';
import type { SourceTextEdit } from './text-edit';

const INLINE_GRAPHICS_ENVIRONMENTS = new Set([
  'figure',
  'figure*',
  'minipage',
  'subfigure',
]);
const BLOCKED_FLOAT_ENVIRONMENTS = new Set([
  'array',
  'table',
  'table*',
  'tabular',
  'tabular*',
  'align',
  'align*',
  'displaymath',
  'equation',
  'equation*',
  'gather',
  'gather*',
  'math',
  'multline',
  'multline*',
]);

export type LatexImageInsertionOptions = {
  source: string;
  cursor: number;
  projectPath: string;
  width?: number;
  lineEnding?: LatexLineEnding;
};

export type LatexImageInsertion = {
  edit: SourceTextEdit;
  kind:
    | 'inside-figure'
    | 'after-blocked-environment'
    | 'before-blocked-environment'
    | 'figure';
};

function environmentStack(source: string): string[] {
  const stack: string[] = [];
  const tokens = /\\(begin|end)\{([^{}]+)\}/g;
  for (const token of source.matchAll(tokens)) {
    const [, action, environment] = token;
    if (!environment) {
      continue;
    }
    if (action === 'begin') {
      stack.push(environment);
      continue;
    }
    const index = stack.lastIndexOf(environment);
    if (index >= 0) {
      stack.splice(index, 1);
    }
  }
  return stack;
}

function humanizeFileName(projectPath: string): string {
  const leaf = projectPath.slice(projectPath.lastIndexOf('/') + 1);
  const base = leaf.includes('.') ? leaf.slice(0, leaf.lastIndexOf('.')) : leaf;
  const words = base
    .replace(/[_-]+/g, ' ')
    .replace(/\s+/g, ' ')
    .trim();
  return words ? words[0]?.toUpperCase() + words.slice(1) : 'Figure';
}

function labelForFile(projectPath: string): string {
  const leaf = projectPath.slice(projectPath.lastIndexOf('/') + 1);
  const base = leaf.includes('.') ? leaf.slice(0, leaf.lastIndexOf('.')) : leaf;
  const slug = base
    .toLowerCase()
    .replace(/[^a-z0-9]+/g, '-')
    .replace(/^-+|-+$/g, '');
  return `fig:${slug || 'image'}`;
}

function lineIndent(source: string, cursor: number): string {
  const lineStart = source.lastIndexOf('\n', Math.max(0, cursor - 1)) + 1;
  return source.slice(lineStart, cursor).match(/^\s*/)?.[0] ?? '';
}

function figureSnippet(projectPath: string, width: number, lineEnding: LatexLineEnding): string {
  const caption = humanizeFileName(projectPath);
  const label = labelForFile(projectPath);
  return [
    '\\begin{figure}[htbp]',
    '  \\centering',
    `  \\includegraphics[width=${width.toFixed(2)}\\linewidth]{${projectPath}}`,
    `  \\caption{${caption}}`,
    `  \\label{${label}}`,
    '\\end{figure}',
  ].join(lineEnding);
}

function surroundBlock(source: string, cursor: number, block: string, lineEnding: LatexLineEnding): string {
  const before = source.slice(0, cursor);
  const after = source.slice(cursor);
  const doubleLineEnding = lineEnding + lineEnding;
  const prefix =
    before.length === 0
      ? ''
      : before.endsWith(doubleLineEnding)
        ? ''
        : before.endsWith(lineEnding)
          ? lineEnding
          : doubleLineEnding;
  const suffix =
    after.length === 0
      ? ''
      : after.startsWith(doubleLineEnding)
        ? ''
        : after.startsWith(lineEnding)
          ? lineEnding
          : doubleLineEnding;
  return `${prefix}${block}${suffix}`;
}

function findEnvironmentEnd(
  source: string,
  cursor: number,
  environment: string,
): number | undefined {
  const escaped = environment.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
  const close = new RegExp(`\\\\end\\{${escaped}\\}`, 'g');
  close.lastIndex = cursor;
  const match = close.exec(source);
  return match ? match.index + match[0].length : undefined;
}

function findEnvironmentStart(
  source: string,
  cursor: number,
  environment: string,
): number | undefined {
  const escaped = environment.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
  const open = new RegExp(`\\\\begin\\{${escaped}\\}`, 'g');
  let found: number | undefined;
  for (const match of source.slice(0, cursor).matchAll(open)) {
    found = match.index;
  }
  return found;
}

export function createLatexImageInsertion(
  options: LatexImageInsertionOptions,
): LatexImageInsertion {
  const cursor = Math.max(0, Math.min(options.cursor, options.source.length));
  const width = Math.min(1, Math.max(0.05, options.width ?? 0.8));
  const lineEnding = options.lineEnding ?? detectLatexLineEnding(options.source);
  const stack = environmentStack(options.source.slice(0, cursor));
  const current = stack.at(-1);

  if (current && INLINE_GRAPHICS_ENVIRONMENTS.has(current)) {
    const indent = lineIndent(options.source, cursor);
    const text = `${indent}\\includegraphics[width=\\linewidth]{${options.projectPath}}`;
    return {
      kind: 'inside-figure',
      edit: { from: cursor, to: cursor, text },
    };
  }

  if (current && BLOCKED_FLOAT_ENVIRONMENTS.has(current)) {
    const block = figureSnippet(options.projectPath, width, lineEnding);
    const end = findEnvironmentEnd(options.source, cursor, current);
    if (end !== undefined) {
      return {
        kind: 'after-blocked-environment',
        edit: {
          from: end,
          to: end,
          text: surroundBlock(options.source, end, block, lineEnding),
        },
      };
    }

    const start = findEnvironmentStart(options.source, cursor, current);
    if (start !== undefined) {
      return {
        kind: 'before-blocked-environment',
        edit: {
          from: start,
          to: start,
          text: surroundBlock(options.source, start, block, lineEnding),
        },
      };
    }
  }

  return {
    kind: 'figure',
    edit: {
      from: cursor,
      to: cursor,
      text: surroundBlock(
        options.source,
        cursor,
        figureSnippet(options.projectPath, width, lineEnding),
        lineEnding,
      ),
    },
  };
}