ai-terminal-assistant/packages/core/src/utils/image.ts

/**
 * 图片处理工具
 *
 * 提供图片文件读取、格式检测、base64 编码等功能
 */

import * as fs from 'fs/promises';
import * as path from 'path';

/** 支持的图片扩展名 */
export const IMAGE_EXTENSIONS = ['.png', '.jpg', '.jpeg', '.gif', '.webp'];

/** 图片 MIME 类型映射 */
const MIME_TYPES: Record<string, string> = {
  '.png': 'image/png',
  '.jpg': 'image/jpeg',
  '.jpeg': 'image/jpeg',
  '.gif': 'image/gif',
  '.webp': 'image/webp',
};

/** 图片信息 */
export interface ImageInfo {
  /** 原始文件路径 */
  path: string;
  /** 文件名 */
  filename: string;
  /** 扩展名 */
  extension: string;
  /** MIME 类型 */
  mimeType: string;
  /** 文件大小（字节） */
  size: number;
  /** base64 编码的数据 */
  base64: string;
  /** 完整的 data URL */
  dataUrl: string;
}

/** 图片加载结果 */
export interface ImageLoadResult {
  success: boolean;
  image?: ImageInfo;
  error?: string;
}

/**
 * 判断文件路径是否为图片
 */
export function isImagePath(filePath: string): boolean {
  const ext = path.extname(filePath).toLowerCase();
  return IMAGE_EXTENSIONS.includes(ext);
}

/**
 * 从输入文本中提取图片引用
 * 支持多种格式:
 * 1. @path/to/image.png（不带空格的路径）
 * 2. @"path/to/image with spaces.png"（带空格的路径用引号包裹）
 * 3. @/path/to/image.png（绝对路径，自动匹配到图片扩展名结束）
 *
 * @param input 用户输入
 * @returns 图片路径列表和去除图片引用后的文本
 */
export function extractImageReferences(input: string): {
  imagePaths: string[];
  textContent: string;
} {
  const imagePaths: string[] = [];
  let textContent = input;

  // 模式1: 带引号的路径 @"path/to/image.png" 或 @'path/to/image.png'
  const quotedMatches = [...input.matchAll(/@["']([^"']+\.(?:png|jpg|jpeg|gif|webp))["']/gi)];
  for (const match of quotedMatches) {
    imagePaths.push(match[1]);
    textContent = textContent.replace(match[0], ' ');
  }

  // 模式2: 绝对路径（以 / 或 ~ 开头，匹配到图片扩展名结束）
  // 支持路径中包含空格
  const absoluteMatches = [...textContent.matchAll(/@([/~][^\n]*?\.(?:png|jpg|jpeg|gif|webp))(?=\s|$)/gi)];
  for (const match of absoluteMatches) {
    if (!imagePaths.includes(match[1])) {
      imagePaths.push(match[1]);
      textContent = textContent.replace(match[0], ' ');
    }
  }

  // 模式3: 相对路径（不以 / 开头，不包含空格）
  const relativeMatches = [...textContent.matchAll(/@((?:\.\/|\.\.\/)?[^\s@"'/][^\s@"']*\.(?:png|jpg|jpeg|gif|webp))/gi)];
  for (const match of relativeMatches) {
    if (!imagePaths.includes(match[1])) {
      imagePaths.push(match[1]);
      textContent = textContent.replace(match[0], ' ');
    }
  }

  // 清理多余空格
  textContent = textContent.replace(/\s+/g, ' ').trim();

  return { imagePaths, textContent };
}

/**
 * 加载图片文件
 * @param filePath 图片路径（相对或绝对）
 * @param workdir 工作目录（用于解析相对路径）
 */
export async function loadImage(
  filePath: string,
  workdir: string = process.cwd()
): Promise<ImageLoadResult> {
  try {
    // 解析路径
    const absolutePath = path.isAbsolute(filePath)
      ? filePath
      : path.resolve(workdir, filePath);

    // 检查扩展名
    const ext = path.extname(absolutePath).toLowerCase();
    if (!IMAGE_EXTENSIONS.includes(ext)) {
      return {
        success: false,
        error: `不支持的图片格式: ${ext}。支持的格式: ${IMAGE_EXTENSIONS.join(', ')}`,
      };
    }

    // 读取文件
    const buffer = await fs.readFile(absolutePath);
    const stats = await fs.stat(absolutePath);

    // 转换为 base64
    const base64 = buffer.toString('base64');
    const mimeType = MIME_TYPES[ext] || 'application/octet-stream';
    const dataUrl = `data:${mimeType};base64,${base64}`;

    return {
      success: true,
      image: {
        path: absolutePath,
        filename: path.basename(absolutePath),
        extension: ext,
        mimeType,
        size: stats.size,
        base64,
        dataUrl,
      },
    };
  } catch (error) {
    if ((error as NodeJS.ErrnoException).code === 'ENOENT') {
      return {
        success: false,
        error: `图片文件不存在: ${filePath}`,
      };
    }
    return {
      success: false,
      error: `加载图片失败: ${error instanceof Error ? error.message : String(error)}`,
    };
  }
}

/**
 * 批量加载图片
 * @param filePaths 图片路径列表
 * @param workdir 工作目录
 */
export async function loadImages(
  filePaths: string[],
  workdir: string = process.cwd()
): Promise<{
  images: ImageInfo[];
  errors: Array<{ path: string; error: string }>;
}> {
  const images: ImageInfo[] = [];
  const errors: Array<{ path: string; error: string }> = [];

  for (const filePath of filePaths) {
    const result = await loadImage(filePath, workdir);
    if (result.success && result.image) {
      images.push(result.image);
    } else {
      errors.push({ path: filePath, error: result.error || '未知错误' });
    }
  }

  return { images, errors };
}

/**
 * 格式化文件大小
 */
export function formatFileSize(bytes: number): string {
  if (bytes < 1024) {
    return `${bytes}B`;
  } else if (bytes < 1024 * 1024) {
    return `${(bytes / 1024).toFixed(1)}KB`;
  } else {
    return `${(bytes / (1024 * 1024)).toFixed(1)}MB`;
  }
}