artifacts/api-server/src/routes/images.ts

import { Router } from "express";
import { db, imagesTable, usersTable, configTable, creditTransactionsTable } from "@workspace/db";
import {
  GenerateImageBody,
  GetImageHistoryQueryParams,
  DeleteImageParams,
} from "@workspace/api-zod";
import { desc, eq, count, and, or, isNull, sql } from "drizzle-orm";
import {
  getValidBearerToken, refreshAccessToken,
  getPoolToken, markAccountUsed, tryRefreshPoolAccount, disablePoolAccount,
} from "./config";
import { optionalJwtAuth } from "./auth";
import { generateGuardId } from "../guardId";

const router = Router();

const GEMINIGEN_BASE = "https://api.geminigen.ai/api";

const STYLE_PROMPTS: Record<string, string> = {
  realistic: "photorealistic, high quality, detailed, 8k resolution",
  anime: "anime style, manga art style, japanese animation",
  artistic: "artistic, fine art, expressive brushwork",
  cartoon: "cartoon style, colorful, fun illustration",
  sketch: "pencil sketch, hand drawn, black and white drawing",
  oil_painting: "oil painting, classical art style, textured canvas",
  watercolor: "watercolor painting, soft colors, fluid brushstrokes",
  digital_art: "digital art, concept art, highly detailed digital illustration",
};

const ORIENTATION_MAP: Record<string, string> = {
  "1:1": "square",
  "16:9": "landscape",
  "9:16": "portrait",
  "4:3": "landscape",
  "3:4": "portrait",
  "2:3": "portrait",
  "3:2": "landscape",
};

const USER_AGENT = "Mozilla/5.0 (iPhone; CPU iPhone OS 18_3 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) CriOS/124.0.6367.111 Mobile/15E148 Safari/604.1";

function base64ToBlob(base64: string, mime: string): Blob {
  const binary = Buffer.from(base64, "base64");
  return new Blob([binary], { type: mime });
}

async function pollForImage(uuid: string, token: string, maxWaitMs = 120000): Promise<string | null> {
  const interval = 3000;
  const start = Date.now();

  while (Date.now() - start < maxWaitMs) {
    await new Promise((r) => setTimeout(r, interval));

    const resp = await fetch(`${GEMINIGEN_BASE}/history/${uuid}`, {
      headers: {
        Authorization: `Bearer ${token}`,
        "x-guard-id": generateGuardId("/api/history/" + uuid, "get"),
        "User-Agent": USER_AGENT,
        Accept: "application/json",
      },
    });

    if (!resp.ok) break;

    const data = await resp.json() as {
      status?: number;
      generated_image?: Array<{ image_url?: string }>;
    };

    if (data.status === 2 || data.status === 3) {
      return data.generated_image?.[0]?.image_url || null;
    }

    if (typeof data.status === "number" && data.status > 3) break;
  }

  return null;
}

async function callGrokEndpoint(prompt: string, orientation: string, token: string, refImageBase64?: string, refImageMime?: string) {
  const form = new FormData();
  form.append("prompt", prompt);
  form.append("orientation", orientation);
  form.append("num_result", "1");

  if (refImageBase64 && refImageMime) {
    const blob = base64ToBlob(refImageBase64, refImageMime);
    form.append("files", blob, "reference.jpg");
  }

  const resp = await fetch(`${GEMINIGEN_BASE}/imagen/grok`, {
    method: "POST",
    headers: { Authorization: `Bearer ${token}`, "x-guard-id": generateGuardId("/api/imagen/grok", "post"), "User-Agent": USER_AGENT, Accept: "application/json" },
    body: form,
  });

  const body = await resp.json().catch(async () => ({ raw: await resp.text().catch(() => "") }));
  return { status: resp.status, body };
}

async function callMetaEndpoint(prompt: string, orientation: string, token: string, refImageBase64?: string, refImageMime?: string) {
  const form = new FormData();
  form.append("prompt", prompt);
  form.append("orientation", orientation);
  form.append("num_result", "1");

  if (refImageBase64 && refImageMime) {
    const blob = base64ToBlob(refImageBase64, refImageMime);
    form.append("files", blob, "reference.jpg");
  }

  const resp = await fetch(`${GEMINIGEN_BASE}/meta_ai/generate`, {
    method: "POST",
    headers: { Authorization: `Bearer ${token}`, "x-guard-id": generateGuardId("/api/meta_ai/generate", "post"), "User-Agent": USER_AGENT, Accept: "application/json" },
    body: form,
  });

  const body = await resp.json().catch(async () => ({ raw: await resp.text().catch(() => "") }));
  return { status: resp.status, body };
}

async function callImagenEndpoint(model: string, prompt: string, aspectRatio: string, style: string, token: string, refImageBase64?: string, refImageMime?: string, resolution?: string) {
  const form = new FormData();
  form.append("prompt", prompt);
  form.append("model", model);
  form.append("aspect_ratio", aspectRatio);
  form.append("output_format", "jpg");
  if (resolution) form.append("resolution", resolution);

  if (refImageBase64 && refImageMime) {
    const blob = base64ToBlob(refImageBase64, refImageMime);
    form.append("files", blob, "reference.jpg");
  }

  const resp = await fetch(`${GEMINIGEN_BASE}/generate_image`, {
    method: "POST",
    headers: { Authorization: `Bearer ${token}`, "x-guard-id": generateGuardId("/api/generate_image", "post"), "User-Agent": USER_AGENT, Accept: "application/json" },
    body: form,
  });

  const body = await resp.json().catch(async () => ({ raw: await resp.text().catch(() => "") }));
  return { status: resp.status, body };
}

// ── Credit helpers ─────────────────────────────────────────────────────────────
async function getConfigVal(key: string): Promise<string | null> {
  const rows = await db.select({ value: configTable.value }).from(configTable).where(eq(configTable.key, key)).limit(1);
  return rows[0]?.value ?? null;
}

async function checkAndDeductCredits(userId: number, cost: number, description: string): Promise<{ ok: boolean; balance?: number }> {
  const enabled = await getConfigVal("enable_credits");
  if (enabled !== "true") return { ok: true };

  const [user] = await db.select({ credits: usersTable.credits }).from(usersTable).where(eq(usersTable.id, userId)).limit(1);
  if (!user) return { ok: false };
  if (user.credits < cost) return { ok: false, balance: user.credits };

  const [updated] = await db
    .update(usersTable)
    .set({ credits: sql`${usersTable.credits} - ${cost}` })
    .where(eq(usersTable.id, userId))
    .returning({ credits: usersTable.credits });

  await db.insert(creditTransactionsTable).values({
    userId,
    amount: -cost,
    type: "spend",
    description,
  });

  return { ok: true, balance: updated.credits };
}

router.post("/generate", optionalJwtAuth, async (req, res) => {
  const bodyResult = GenerateImageBody.safeParse(req.body);
  if (!bodyResult.success) {
    return res.status(400).json({ error: "VALIDATION_ERROR", message: "Invalid request body" });
  }

  const {
    prompt,
    style = "realistic",
    aspectRatio = "1:1",
    model = "grok",
    resolution,
    referenceImageBase64,
    referenceImageMime,
    isPrivate = false,
  } = bodyResult.data as any;

  const userId: number | null = (req as any).jwtUserId ?? null;

  // ── Credits check ────────────────────────────────────────────────────────────
  if (userId !== null) {
    const costStr = await getConfigVal("image_gen_cost");
    const cost = Number(costStr) || 0;
    if (cost > 0) {
      const creditResult = await checkAndDeductCredits(userId, cost, `圖片生成（${model}）`);
      if (!creditResult.ok) {
        return res.status(402).json({
          error: "INSUFFICIENT_CREDITS",
          message: `點數不足，此操作需要 ${cost} 點`,
          balance: creditResult.balance ?? 0,
        });
      }
    }
  }

  const stylePrompt = style === "none" ? "" : (STYLE_PROMPTS[style] || "");
  const fullPrompt = stylePrompt ? `${prompt}, ${stylePrompt}` : prompt;
  const orientation = ORIENTATION_MAP[aspectRatio] || "square";

  const isImagenModel = model === "imagen-pro" || model === "imagen-4" || model === "imagen-flash" || model === "nano-banana-pro" || model === "nano-banana-2";
  const apiModelId = isImagenModel ? model : model;

  let imageUrl = "";
  let usedFallback = false;
  let fallbackReason = "";
  let responseStatus = 0;
  let responseBody: unknown = {};
  let pollResult: Record<string, unknown> = {};
  const startTime = Date.now();

  // ── Pool-aware token selection ──────────────────────────────────────────────
  const failedPoolIds: number[] = [];
  let currentAccountId: number | null = null;

  async function pickToken(): Promise<string | null> {
    const poolEntry = await getPoolToken(failedPoolIds);
    if (poolEntry) {
      currentAccountId = poolEntry.accountId;
      return poolEntry.token;
    }
    currentAccountId = null;
    return getValidBearerToken();
  }

  async function handleTokenExpiry(): Promise<string | null> {
    if (currentAccountId !== null) {
      const refreshed = await tryRefreshPoolAccount(currentAccountId);
      if (refreshed) return refreshed;
      failedPoolIds.push(currentAccountId);
      const next = await getPoolToken(failedPoolIds);
      if (next) {
        currentAccountId = next.accountId;
        return next.token;
      }
    }
    return refreshAccessToken();
  }

  let token = await pickToken();

  const requestInfo = {
    url: isImagenModel
      ? `${GEMINIGEN_BASE}/generate_image`
      : model === "meta"
      ? `${GEMINIGEN_BASE}/meta_ai/generate`
      : `${GEMINIGEN_BASE}/imagen/grok`,
    model: isImagenModel ? apiModelId : model,
    fields: isImagenModel
      ? { prompt: fullPrompt, model: apiModelId, aspect_ratio: aspectRatio, output_format: "jpg", ...(resolution ? { resolution } : {}), hasReferenceImage: !!referenceImageBase64 }
      : { prompt: fullPrompt, orientation, num_result: "1", hasReferenceImage: !!referenceImageBase64 },
  };

  try {
    if (!token) throw new Error("未設定 API Token，請到設定頁面輸入 token");

    let result: { status: number; body: unknown };

    if (model === "grok") {
      result = await callGrokEndpoint(fullPrompt, orientation, token, referenceImageBase64, referenceImageMime);
    } else if (model === "meta") {
      result = await callMetaEndpoint(fullPrompt, orientation, token, referenceImageBase64, referenceImageMime);
    } else {
      result = await callImagenEndpoint(apiModelId, fullPrompt, aspectRatio, style, token, referenceImageBase64, referenceImageMime, resolution);
    }

    if (result.status === 401) {
      const newToken = await handleTokenExpiry();
      if (!newToken) throw new Error("Token 已過期且無法自動刷新");
      token = newToken;

      if (model === "grok") result = await callGrokEndpoint(fullPrompt, orientation, token, referenceImageBase64, referenceImageMime);
      else if (model === "meta") result = await callMetaEndpoint(fullPrompt, orientation, token, referenceImageBase64, referenceImageMime);
      else result = await callImagenEndpoint(apiModelId, fullPrompt, aspectRatio, style, token, referenceImageBase64, referenceImageMime, resolution);
    }

    responseStatus = result.status;
    responseBody = result.body;

    const data = result.body as { uuid?: string; base64_images?: string; generated_image?: Array<{ image_url?: string }>; detail?: { error_code?: string; error_message?: string } };

    const errMsg = (data?.detail?.error_message || "").toLowerCase();
    const isTokenExpired = result.status === 401 || result.status === 403
      || data?.detail?.error_code === "TOKEN_EXPIRED"
      || errMsg.includes("expired") || errMsg.includes("token");

    if (isTokenExpired) {
      const newToken = await handleTokenExpiry();
      if (!newToken) throw new Error("Token 已過期且無法自動刷新，請重新取得");
      token = newToken;

      let retryResult: { status: number; body: unknown };
      if (model === "grok") retryResult = await callGrokEndpoint(fullPrompt, orientation, token, referenceImageBase64, referenceImageMime);
      else if (model === "meta") retryResult = await callMetaEndpoint(fullPrompt, orientation, token, referenceImageBase64, referenceImageMime);
      else retryResult = await callImagenEndpoint(apiModelId, fullPrompt, aspectRatio, style, token, referenceImageBase64, referenceImageMime, resolution);

      responseStatus = retryResult.status;
      responseBody = retryResult.body;
      Object.assign(data, retryResult.body);
    }

    if (!result.status.toString().startsWith("2") && !isTokenExpired) {
      const msg = (data as any)?.detail?.error_message || (data as any)?.detail?.error_code || `HTTP ${result.status}`;
      throw new Error(`API 錯誤：${msg}`);
    }

    const finalData = responseBody as typeof data;

    if ((finalData as any)?.detail?.error_code && (finalData as any)?.detail?.error_code !== "TOKEN_EXPIRED") {
      const msg = (finalData as any)?.detail?.error_message || (finalData as any)?.detail?.error_code;
      throw new Error(`API 錯誤：${msg}`);
    }

    if (data.base64_images) {
      imageUrl = `data:image/png;base64,${data.base64_images}`;
      pollResult = { type: "immediate_base64" };
    } else if (data.generated_image?.[0]?.image_url) {
      imageUrl = data.generated_image[0].image_url;
      pollResult = { type: "immediate_url" };
    } else if (data.uuid) {
      pollResult.uuid = data.uuid;
      const polledUrl = await pollForImage(data.uuid, token);
      if (!polledUrl) throw new Error("圖片生成逾時或未返回結果");
      imageUrl = polledUrl;
      pollResult.imageUrl = imageUrl;
      pollResult.status = "completed";
    } else {
      throw new Error("API 未返回圖片或任務 UUID");
    }
  } catch (err: unknown) {
    const errMsg = err instanceof Error ? err.message : String(err);
    req.log.warn({ err }, "Image generation failed, using fallback");
    usedFallback = true;
    fallbackReason = errMsg;

    const fallbackSizes: Record<string, string> = {
      "1:1": "1024/1024", "16:9": "1344/768", "9:16": "768/1344",
      "4:3": "1152/896", "3:4": "896/1152", "2:3": "768/1152", "3:2": "1152/768",
    };
    const seed = Math.floor(Math.random() * 1000000);
    imageUrl = `https://picsum.photos/seed/${seed}/${fallbackSizes[aspectRatio] || "1024/1024"}`;
  }

  const durationMs = Date.now() - startTime;
  const isTokenError = usedFallback && (
    fallbackReason?.includes("Token 已過期") ||
    fallbackReason?.includes("無法自動刷新") ||
    fallbackReason?.includes("未設定 API Token") ||
    fallbackReason?.includes("REFRESH_TOKEN_EXPIRED")
  );

  // Mark the pool account as used (after successful generation)
  if (currentAccountId !== null) {
    markAccountUsed(currentAccountId).catch(() => {});
  }

  const [inserted] = await db
    .insert(imagesTable)
    .values({ imageUrl, prompt, style, aspectRatio, model, isPrivate: !!isPrivate, userId })
    .returning();

  res.json({
    id: inserted.id,
    imageUrl: inserted.imageUrl,
    prompt: inserted.prompt,
    style: inserted.style,
    aspectRatio: inserted.aspectRatio,
    model: inserted.model,
    createdAt: inserted.createdAt.toISOString(),
    ...(usedFallback ? { error: fallbackReason, tokenExpired: isTokenError } : {}),
    apiDebug: {
      requestUrl: requestInfo.url,
      requestMethod: "POST",
      requestContentType: "multipart/form-data",
      requestHeaders: {
        Authorization: token ? "Bearer ****" : "(無 Token)",
        "User-Agent": USER_AGENT,
        Accept: "application/json",
      },
      requestBody: requestInfo.fields,
      responseStatus,
      responseBody,
      pollResult,
      durationMs,
      usedFallback,
      ...(fallbackReason ? { fallbackReason } : {}),
    },
  });
});

router.get("/history", optionalJwtAuth, async (req, res) => {
  const paramsResult = GetImageHistoryQueryParams.safeParse({
    limit: req.query.limit ? Number(req.query.limit) : 20,
    offset: req.query.offset ? Number(req.query.offset) : 0,
  });

  const { limit = 20, offset = 0 } = paramsResult.success ? paramsResult.data : {};
  const currentUserId: number | null = (req as any).jwtUserId ?? null;

  // Visibility filter:
  // - Public images are always visible
  // - Private images are only visible to their owner
  const visibilityFilter = currentUserId
    ? or(eq(imagesTable.isPrivate, false), and(eq(imagesTable.isPrivate, true), eq(imagesTable.userId, currentUserId)))
    : eq(imagesTable.isPrivate, false);

  const [images, [{ value: total }]] = await Promise.all([
    db.select().from(imagesTable).where(visibilityFilter).orderBy(desc(imagesTable.createdAt)).limit(limit).offset(offset),
    db.select({ value: count() }).from(imagesTable).where(visibilityFilter),
  ]);

  res.json({
    images: images.map((img) => ({
      id: img.id,
      imageUrl: img.imageUrl,
      prompt: img.prompt,
      style: img.style,
      aspectRatio: img.aspectRatio,
      model: img.model,
      isPrivate: img.isPrivate,
      userId: img.userId,
      createdAt: img.createdAt.toISOString(),
    })),
    total: Number(total),
  });
});

router.delete("/:id", async (req, res) => {
  const paramsResult = DeleteImageParams.safeParse({ id: Number(req.params.id) });
  if (!paramsResult.success) {
    return res.status(400).json({ error: "INVALID_ID", message: "Invalid image ID" });
  }

  const deleted = await db.delete(imagesTable).where(eq(imagesTable.id, paramsResult.data.id)).returning();

  if (deleted.length === 0) {
    return res.status(404).json({ error: "NOT_FOUND", message: "Image not found" });
  }

  res.json({ success: true, message: "Image deleted successfully" });
});

export default router;