makelore/electron/image-workspace/local-workspace.ts

import { createHash, randomUUID } from 'node:crypto';
import { mkdir, readFile, rename, rm, writeFile } from 'node:fs/promises';
import { join } from 'node:path';
import type {
  ImageWorkspaceCapabilities,
  ImageWorkspaceGenerationSettings,
  ImageWorkspaceImage,
  ImageWorkspaceProject,
  ImageWorkspaceReferenceUploadInput,
  ImageWorkspaceReferenceUploadResult,
  ImageWorkspaceSendMessageInput,
  ImageWorkspaceSnapshot,
} from '../../shared/image-workspace';

const LOCAL_WORKSPACE_SCHEMA_VERSION = 1;
const LOCAL_WORKSPACE_DIRECTORY = 'image-workspace-development';
const LOCAL_WORKSPACE_FILE = 'workspace.json';
const MAX_PROJECT_NAME_LENGTH = 80;
const MAX_PROMPT_LENGTH = 4_000;
const MAX_AGENT_COUNT = 50;
const MAX_REFERENCE_BYTES = 5 * 1024 * 1024;

const LOCAL_CAPABILITIES: ImageWorkspaceCapabilities = {
  modes: [
    { id: 'generate', label: '图片生成', description: '从文字或参考图生成新的图片' },
    { id: 'edit', label: '参考图编辑', description: '使用参考图继续调整画面' },
  ],
  models: [
    { id: 'local-placeholder-v1', label: '标准创作模型', description: '适用于图片生成与参考图编辑' },
  ],
  aspectRatios: [
    { id: '1:1', label: '1:1' },
    { id: '3:4', label: '3:4' },
    { id: '4:3', label: '4:3' },
    { id: '16:9', label: '16:9' },
  ],
  resolutions: [
    { id: '1024', label: '1K' },
    { id: '2048', label: '2K' },
  ],
  outputCounts: [
    { id: '1', label: '1 张' },
    { id: '2', label: '2 张' },
    { id: '4', label: '4 张' },
  ],
  defaultModeId: 'generate',
  defaultModelId: 'local-placeholder-v1',
  defaultAspectRatioId: '1:1',
  defaultResolutionId: '1024',
  defaultOutputCountId: '1',
  maxReferenceImages: 4,
  referenceUpload: {
    enabled: true,
    acceptedMimeTypes: ['image/png', 'image/jpeg', 'image/webp'],
    maxBytes: MAX_REFERENCE_BYTES,
  },
};

type PersistedImageWorkspace = {
  schemaVersion: typeof LOCAL_WORKSPACE_SCHEMA_VERSION;
  activeProjectId: string | null;
  projects: ImageWorkspaceProject[];
  referencesByProjectId: Record<string, ImageWorkspaceImage[]>;
};

type LocalImageWorkspaceOptions = {
  userDataDir: string;
  now?: () => Date;
  createId?: () => string;
};

export class LocalImageWorkspaceError extends Error {
  readonly status: number;
  readonly code: string;

  constructor(status: number, code: string, message: string) {
    super(message);
    this.name = 'LocalImageWorkspaceError';
    this.status = status;
    this.code = code;
  }
}

function createEmptyState(): PersistedImageWorkspace {
  return {
    schemaVersion: LOCAL_WORKSPACE_SCHEMA_VERSION,
    activeProjectId: null,
    projects: [],
    referencesByProjectId: {},
  };
}

function clone<T>(value: T): T {
  return JSON.parse(JSON.stringify(value)) as T;
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null && !Array.isArray(value);
}

function isPersistedWorkspace(value: unknown): value is PersistedImageWorkspace {
  return isRecord(value)
    && value.schemaVersion === LOCAL_WORKSPACE_SCHEMA_VERSION
    && (typeof value.activeProjectId === 'string' || value.activeProjectId === null)
    && Array.isArray(value.projects)
    && isRecord(value.referencesByProjectId);
}

function escapeXml(value: string): string {
  return value
    .replaceAll('&', '&amp;')
    .replaceAll('<', '&lt;')
    .replaceAll('>', '&gt;')
    .replaceAll('"', '&quot;')
    .replaceAll("'", '&apos;');
}

function wrapPrompt(prompt: string, maxCharacters = 24, maxLines = 3): string[] {
  const characters = Array.from(prompt.replace(/\s+/g, ' ').trim());
  const lines: string[] = [];
  for (let index = 0; index < characters.length && lines.length < maxLines; index += maxCharacters) {
    const line = characters.slice(index, index + maxCharacters).join('');
    const truncated = index + maxCharacters < characters.length && lines.length === maxLines - 1;
    lines.push(truncated ? `${line.slice(0, Math.max(0, line.length - 1))}…` : line);
  }
  return lines.length > 0 ? lines : ['未填写提示词'];
}

function resolveDimensions(settings: Required<ImageWorkspaceGenerationSettings>): { width: number; height: number } {
  const longEdge = settings.resolutionId === '2048' ? 2048 : 1024;
  const [widthRatio, heightRatio] = (settings.aspectRatioId ?? '1:1')
    .split(':')
    .map((value) => Number(value));
  if (!widthRatio || !heightRatio) return { width: longEdge, height: longEdge };
  if (widthRatio >= heightRatio) {
    return { width: longEdge, height: Math.round(longEdge * heightRatio / widthRatio) };
  }
  return { width: Math.round(longEdge * widthRatio / heightRatio), height: longEdge };
}

function optionId(
  value: string | undefined,
  defaultValue: string,
  options: Array<{ id: string }>,
  label: string,
): string {
  const resolved = value || defaultValue;
  if (!options.some((option) => option.id === resolved)) {
    throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_SETTING', `当前创作空间不支持这个${label}`);
  }
  return resolved;
}

function normalizeSettings(settings: ImageWorkspaceGenerationSettings): Required<ImageWorkspaceGenerationSettings> {
  return {
    modeId: optionId(settings.modeId, LOCAL_CAPABILITIES.defaultModeId!, LOCAL_CAPABILITIES.modes, '创作模式'),
    modelId: optionId(settings.modelId, LOCAL_CAPABILITIES.defaultModelId!, LOCAL_CAPABILITIES.models, '模型'),
    aspectRatioId: optionId(
      settings.aspectRatioId,
      LOCAL_CAPABILITIES.defaultAspectRatioId!,
      LOCAL_CAPABILITIES.aspectRatios,
      '画幅',
    ),
    resolutionId: optionId(
      settings.resolutionId,
      LOCAL_CAPABILITIES.defaultResolutionId!,
      LOCAL_CAPABILITIES.resolutions,
      '分辨率',
    ),
    outputCountId: optionId(
      settings.outputCountId,
      LOCAL_CAPABILITIES.defaultOutputCountId!,
      LOCAL_CAPABILITIES.outputCounts,
      '生成数量',
    ),
  };
}

function createPlaceholderImage(
  id: string,
  prompt: string,
  settings: Required<ImageWorkspaceGenerationSettings>,
  referenceImageIds: string[],
  outputIndex: number,
): ImageWorkspaceImage {
  const seed = createHash('sha256')
    .update(JSON.stringify({ prompt, settings, referenceImageIds, outputIndex }))
    .digest('hex');
  const hue = Number.parseInt(seed.slice(0, 4), 16) % 360;
  const accentHue = (hue + 58 + Number.parseInt(seed.slice(4, 6), 16) % 90) % 360;
  const { width, height } = resolveDimensions(settings);
  const promptLines = wrapPrompt(prompt);
  const promptText = promptLines.map((line, index) => (
    `<tspan x="${Math.round(width * 0.08)}" dy="${index === 0 ? 0 : Math.round(height * 0.075)}">${escapeXml(line)}</tspan>`
  )).join('');
  const metadata = [
    LOCAL_CAPABILITIES.models.find((model) => model.id === settings.modelId)?.label ?? '标准创作模型',
    settings.aspectRatioId,
    `${width}×${height}`,
    referenceImageIds.length > 0 ? `${referenceImageIds.length} 张参考图` : '纯文字创作',
  ].join(' · ');
  const svg = `
    <svg xmlns="http://www.w3.org/2000/svg" width="${width}" height="${height}" viewBox="0 0 ${width} ${height}">
      <rect width="${width}" height="${height}" fill="hsl(${hue} 72% 88%)"/>
      <circle cx="${Math.round(width * 0.82)}" cy="${Math.round(height * 0.2)}" r="${Math.round(Math.min(width, height) * 0.28)}" fill="hsl(${accentHue} 78% 68%)" opacity="0.72"/>
      <rect x="${Math.round(width * 0.06)}" y="${Math.round(height * 0.08)}" width="${Math.round(width * 0.88)}" height="${Math.round(height * 0.84)}" rx="${Math.round(Math.min(width, height) * 0.04)}" fill="white" fill-opacity="0.78" stroke="#202638" stroke-width="${Math.max(4, Math.round(Math.min(width, height) * 0.008))}"/>
      <text x="${Math.round(width * 0.08)}" y="${Math.round(height * 0.19)}" fill="#202638" font-family="Arial, sans-serif" font-size="${Math.max(24, Math.round(Math.min(width, height) * 0.055))}" font-weight="800">CREATIVE · ${outputIndex + 1}</text>
      <text x="${Math.round(width * 0.08)}" y="${Math.round(height * 0.39)}" fill="#202638" font-family="Arial, sans-serif" font-size="${Math.max(22, Math.round(Math.min(width, height) * 0.047))}" font-weight="700">${promptText}</text>
      <text x="${Math.round(width * 0.08)}" y="${Math.round(height * 0.82)}" fill="#5f4638" font-family="Arial, sans-serif" font-size="${Math.max(16, Math.round(Math.min(width, height) * 0.027))}" font-weight="700">${escapeXml(metadata)}</text>
    </svg>
  `.trim();
  const url = `data:image/svg+xml;base64,${Buffer.from(svg, 'utf8').toString('base64')}`;
  return {
    id,
    url,
    thumbnailUrl: url,
    alt: `生成图片：${prompt}`,
  };
}

function validateBase64(contentBase64: string): Buffer {
  const normalized = contentBase64.trim();
  if (!normalized || normalized.length % 4 !== 0 || !/^[A-Za-z0-9+/]*={0,2}$/.test(normalized)) {
    throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_REFERENCE', '参考图内容不是有效的 Base64');
  }
  const bytes = Buffer.from(normalized, 'base64');
  if (bytes.length === 0 || bytes.length > MAX_REFERENCE_BYTES) {
    throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_REFERENCE', '参考图大小必须在 5 MB 以内');
  }
  return bytes;
}

export function isLocalImageWorkspaceDevelopmentEnabled(options: {
  isPackaged: boolean;
  configuredMode?: string;
}): boolean {
  return !options.isPackaged && options.configuredMode?.trim().toLowerCase() === 'local';
}

export function getLocalImageWorkspaceDirectory(userDataDir: string): string {
  return join(userDataDir, LOCAL_WORKSPACE_DIRECTORY);
}

export class LocalImageWorkspace {
  private readonly rootDirectory: string;
  private readonly stateFile: string;
  private readonly now: () => Date;
  private readonly createId: () => string;
  private state: PersistedImageWorkspace | null = null;
  private operation: Promise<void> = Promise.resolve();

  constructor(options: LocalImageWorkspaceOptions) {
    this.rootDirectory = getLocalImageWorkspaceDirectory(options.userDataDir);
    this.stateFile = join(this.rootDirectory, LOCAL_WORKSPACE_FILE);
    this.now = options.now ?? (() => new Date());
    this.createId = options.createId ?? randomUUID;
  }

  getSnapshot(): Promise<ImageWorkspaceSnapshot> {
    return this.enqueue(async () => this.snapshot(await this.load()));
  }

  createProject(name: string): Promise<ImageWorkspaceSnapshot> {
    return this.mutate(async (state) => {
      const projectName = name.trim();
      if (!projectName || projectName.length > MAX_PROJECT_NAME_LENGTH) {
        throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_PROJECT_NAME', '项目名称需为 1–80 个字符');
      }
      const timestamp = this.now().toISOString();
      const projectId = `local-project-${this.createId()}`;
      const agentId = `local-agent-${this.createId()}`;
      state.projects.push({
        id: projectId,
        name: projectName,
        agents: [{ id: agentId, name: '视觉创作 Agent' }],
        activeAgentId: agentId,
        messages: [],
        updatedAt: timestamp,
      });
      state.referencesByProjectId[projectId] = [];
      state.activeProjectId = projectId;
    });
  }

  addAgent(projectId: string): Promise<ImageWorkspaceSnapshot> {
    return this.mutate(async (state) => {
      const project = this.requireProject(state, projectId);
      if (project.agents.length >= MAX_AGENT_COUNT) {
        throw new LocalImageWorkspaceError(409, 'IMAGE_WORKSPACE_AGENT_LIMIT', '每个项目最多添加 50 个 Agent');
      }
      const agentId = `local-agent-${this.createId()}`;
      project.agents.push({ id: agentId, name: `创作 Agent ${project.agents.length + 1}` });
      project.updatedAt = this.now().toISOString();
    });
  }

  sendMessage(input: ImageWorkspaceSendMessageInput): Promise<ImageWorkspaceSnapshot> {
    return this.mutate(async (state) => {
      const project = this.requireProject(state, input.projectId);
      if (!project.agents.some((agent) => agent.id === input.agentId)) {
        throw new LocalImageWorkspaceError(404, 'IMAGE_WORKSPACE_AGENT_NOT_FOUND', '当前项目中不存在这个 Agent');
      }
      const prompt = input.prompt.trim();
      if (!prompt || prompt.length > MAX_PROMPT_LENGTH) {
        throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_PROMPT', '创作描述需为 1–4000 个字符');
      }
      const referenceIds = [...new Set(input.referenceImageIds)];
      if (referenceIds.length > LOCAL_CAPABILITIES.maxReferenceImages) {
        throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_REFERENCE_LIMIT', '参考图数量超过创作空间限制');
      }
      const availableImages = new Map<string, ImageWorkspaceImage>();
      for (const image of state.referencesByProjectId[project.id] ?? []) availableImages.set(image.id, image);
      for (const message of project.messages) {
        for (const image of message.images) availableImages.set(image.id, image);
      }
      const references = referenceIds.map((id) => {
        const image = availableImages.get(id);
        if (!image) {
          throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_REFERENCE_NOT_FOUND', '所选参考图不存在');
        }
        return image;
      });
      const settings = normalizeSettings(input.settings);
      const outputCount = Number(settings.outputCountId);
      const timestamp = this.now().toISOString();
      const images = Array.from({ length: outputCount }, (_, index) => createPlaceholderImage(
        `local-image-${this.createId()}`,
        prompt,
        settings,
        referenceIds,
        index,
      ));
      project.messages.push(
        {
          id: `local-message-${this.createId()}`,
          role: 'user',
          agentId: input.agentId,
          text: prompt,
          images: clone(references),
          status: 'succeeded',
          createdAt: timestamp,
        },
        {
          id: `local-message-${this.createId()}`,
          role: 'assistant',
          agentId: input.agentId,
          text: `已生成 ${outputCount} 张图片。`,
          images,
          status: 'succeeded',
          createdAt: timestamp,
        },
      );
      project.activeAgentId = input.agentId;
      project.updatedAt = timestamp;
      state.activeProjectId = project.id;
    });
  }

  uploadReference(input: ImageWorkspaceReferenceUploadInput): Promise<ImageWorkspaceReferenceUploadResult> {
    return this.enqueue(async () => {
      const state = await this.load();
      const project = this.requireProject(state, input.projectId);
      if (!LOCAL_CAPABILITIES.referenceUpload.acceptedMimeTypes.includes(input.mimeType)) {
        throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_REFERENCE_TYPE', '创作空间仅支持 PNG、JPEG 和 WebP');
      }
      const bytes = validateBase64(input.contentBase64);
      const fileName = input.fileName.trim().slice(0, 180) || '本地参考图';
      const reference: ImageWorkspaceImage = {
        id: `local-reference-${this.createId()}`,
        url: `data:${input.mimeType};base64,${bytes.toString('base64')}`,
        alt: fileName,
      };
      state.referencesByProjectId[project.id] ??= [];
      state.referencesByProjectId[project.id].push(reference);
      project.updatedAt = this.now().toISOString();
      await this.persist(state);
      return { workspace: this.snapshot(state), reference: clone(reference) };
    });
  }

  reset(): Promise<ImageWorkspaceSnapshot> {
    return this.enqueue(async () => {
      await rm(this.rootDirectory, { recursive: true, force: true });
      this.state = createEmptyState();
      return this.snapshot(this.state);
    });
  }

  private async mutate(
    operation: (state: PersistedImageWorkspace) => Promise<void>,
  ): Promise<ImageWorkspaceSnapshot> {
    return this.enqueue(async () => {
      const state = await this.load();
      await operation(state);
      await this.persist(state);
      return this.snapshot(state);
    });
  }

  private enqueue<T>(operation: () => Promise<T>): Promise<T> {
    const result = this.operation.then(operation, operation);
    this.operation = result.then(() => undefined, () => undefined);
    return result;
  }

  private async load(): Promise<PersistedImageWorkspace> {
    if (this.state) return this.state;
    try {
      const parsed = JSON.parse(await readFile(this.stateFile, 'utf8')) as unknown;
      if (!isPersistedWorkspace(parsed)) {
        throw new LocalImageWorkspaceError(
          500,
          'IMAGE_WORKSPACE_LOCAL_DATA_INVALID',
          '创作空间数据暂时无法读取',
        );
      }
      this.state = parsed;
    } catch (error) {
      if ((error as NodeJS.ErrnoException).code === 'ENOENT') {
        this.state = createEmptyState();
      } else {
        throw error;
      }
    }
    return this.state;
  }

  private async persist(state: PersistedImageWorkspace): Promise<void> {
    await mkdir(this.rootDirectory, { recursive: true });
    const temporaryFile = `${this.stateFile}.${process.pid}.tmp`;
    await writeFile(temporaryFile, `${JSON.stringify(state, null, 2)}\n`, 'utf8');
    try {
      await rename(temporaryFile, this.stateFile);
    } catch (error) {
      const code = (error as NodeJS.ErrnoException).code;
      if (code !== 'EEXIST' && code !== 'EPERM') throw error;
      await rm(this.stateFile, { force: true });
      await rename(temporaryFile, this.stateFile);
    }
    this.state = state;
  }

  private requireProject(state: PersistedImageWorkspace, projectId: string): ImageWorkspaceProject {
    const project = state.projects.find((item) => item.id === projectId);
    if (!project) {
      throw new LocalImageWorkspaceError(404, 'IMAGE_WORKSPACE_PROJECT_NOT_FOUND', '项目不存在');
    }
    return project;
  }

  private snapshot(state: PersistedImageWorkspace): ImageWorkspaceSnapshot {
    return {
      capabilities: clone(LOCAL_CAPABILITIES),
      projects: clone(state.projects),
      activeProjectId: state.activeProjectId,
    };
  }
}