src/api/generate.ts

import { Hono } from 'hono';
import { eq } from 'drizzle-orm';
import { prompts, images } from '../db/schema';
import { getFullSettings } from './settings';
import type { Env, Variables } from '../index';
import { db as globalDb } from '../db';

export const generateApi = new Hono<{ Bindings: Env; Variables: Variables }>();

// 全局 jobs 表（进程内）—— 任务从客户端连接解耦，刷新页面不影响
type JobState = {
  jobId: string;
  promptId: string;
  status: 'running' | 'done' | 'error' | 'cancelled';
  size: string;
  quality: string;
  startedAt: number;
  finishedAt?: number;
  errorMsg?: string;
  imageId?: string;
  generationDuration?: number;
  abort?: AbortController;
};

const jobs = new Map<string, JobState>();

// 已完成 job 保留 5 分钟供前端拉取，之后清理
const KEEP_DONE_MS = 5 * 60 * 1000;

function gcJobs() {
  const now = Date.now();
  for (const [k, v] of jobs) {
    if (v.status !== 'running' && v.finishedAt && now - v.finishedAt > KEEP_DONE_MS) {
      jobs.delete(k);
    }
  }
}

// 列出当前所有 job（运行中 + 最近完成）
generateApi.get('/jobs', (c) => {
  gcJobs();
  const list = Array.from(jobs.values()).map(({ abort, ...rest }) => rest);
  return c.json(list);
});

// 取消运行中的 job
generateApi.delete('/jobs/:jobId', (c) => {
  const j = jobs.get(c.req.param('jobId'));
  if (!j) return c.json({ error: 'not_found' }, 404);
  if (j.status !== 'running') return c.json({ ok: true, status: j.status });
  j.abort?.abort();
  j.status = 'cancelled';
  j.finishedAt = Date.now();
  return c.json({ ok: true });
});

// 启动生成；立即返回 jobId，任务后台跑
generateApi.post('/:promptId', async (c) => {
  const db = c.var.db;
  const promptId = c.req.param('promptId');
  const body = await c.req.json<{ size?: string; quality?: string }>().catch(() => ({} as { size?: string; quality?: string }));
  const size = body.size || '1024x1024';
  const quality = body.quality || 'hd';

  const p = (await db.select().from(prompts).where(eq(prompts.id, promptId))).at(0);
  if (!p) return c.json({ error: 'prompt_not_found' }, 404);
  const promptText = (p.contentImage || '').trim();
  if (!promptText) return c.json({ error: 'no_image_prompt' }, 400);

  const s = await getFullSettings(db);
  if (!s.openaiApiKey) return c.json({ error: 'api_key_not_set' }, 400);

  const jobId = crypto.randomUUID();
  const ctrl = new AbortController();
  jobs.set(jobId, {
    jobId,
    promptId,
    status: 'running',
    size,
    quality,
    startedAt: Date.now(),
    abort: ctrl,
  });

  // 真后台 —— 不 await，立即返回
  runGenerationJob(jobId, promptId, promptText, size, quality, s, ctrl.signal).catch((err) => {
    const j = jobs.get(jobId);
    if (!j || j.status !== 'running') return;
    j.status = err?.name === 'AbortError' ? 'cancelled' : 'error';
    j.errorMsg = String(err?.message || err);
    j.finishedAt = Date.now();
  });

  return c.json({ jobId, status: 'running' });
});

async function runGenerationJob(
  jobId: string,
  promptId: string,
  promptText: string,
  size: string,
  quality: string,
  s: any,
  signal: AbortSignal
) {
  const t0 = Date.now();
  const url = (s.openaiBaseUrl || 'https://api.openai.com/v1').replace(/\/+$/, '') + '/images/generations';
  const res = await fetch(url, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      Authorization: `Bearer ${s.openaiApiKey}`,
    },
    body: JSON.stringify({
      model: s.imageModel || 'gpt-image-2',
      prompt: promptText,
      size,
      quality,
      n: 1,
      response_format: 'b64_json',
    }),
    signal,
  });
  if (!res.ok) {
    const text = await res.text();
    throw new Error(`HTTP ${res.status}: ${text.slice(0, 300)}`);
  }
  const data: any = await res.json();
  const item = data?.data?.[0];
  let b64: string | null = null;
  if (item?.b64_json) b64 = `data:image/png;base64,${item.b64_json}`;
  else if (item?.url) {
    try {
      const r2 = await fetch(item.url, { signal });
      if (r2.ok) {
        const buf = new Uint8Array(await r2.arrayBuffer());
        const mime = r2.headers.get('content-type') || 'image/png';
        b64 = `data:${mime};base64,${arrayBufferToBase64(buf)}`;
      }
    } catch {}
  }
  if (!b64) throw new Error('响应无 b64_json 也无可下载 url');

  const duration = Date.now() - t0;
  const imageId = crypto.randomUUID();
  await globalDb.insert(images).values({
    id: imageId,
    promptId,
    b64,
    model: s.imageModel || 'gpt-image-2',
    size,
    quality,
    generationDuration: duration,
    promptUsed: promptText,
    createdAt: Date.now(),
  });
  await globalDb.update(prompts).set({ updatedAt: Date.now() }).where(eq(prompts.id, promptId));

  const j = jobs.get(jobId);
  if (j) {
    j.status = 'done';
    j.finishedAt = Date.now();
    j.imageId = imageId;
    j.generationDuration = duration;
  }
}

function arrayBufferToBase64(buf: Uint8Array): string {
  let binary = '';
  const chunk = 0x8000;
  for (let i = 0; i < buf.length; i += chunk) {
    binary += String.fromCharCode.apply(null, Array.from(buf.subarray(i, i + chunk)));
  }
  return btoa(binary);
}