import { Router } from "express";
import { createHash } from "crypto";
import { db, apiKeysTable } from "@workspace/db";
import { eq } from "drizzle-orm";
import { getValidBearerToken, refreshAccessToken, getPoolToken, markAccountUsed, tryRefreshPoolAccount } from "./config";

const router = Router();

const GEMINIGEN_BASE = "https://api.geminigen.ai/api";
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";

const SUPPORTED_MODELS = [
  "grok", "meta", "imagen-pro", "imagen-4", "imagen-flash", "nano-banana-pro", "nano-banana-2",
];

const ALIAS_MAP: Record<string, string> = {
  "dall-e-3": "grok",
  "dall-e-2": "grok",
  "gpt-image-1": "grok",
};

const SIZE_TO_ASPECT: Record<string, string> = {
  "256x256": "1:1",
  "512x512": "1:1",
  "1024x1024": "1:1",
  "1792x1024": "16:9",
  "1024x1792": "9:16",
};

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",
};

async function validateApiKey(req: any, res: any, next: any) {
  const auth = req.headers.authorization as string | undefined;
  if (!auth?.startsWith("Bearer ")) {
    return res.status(401).json({ error: { message: "No API key provided.", type: "invalid_request_error", code: "invalid_api_key" } });
  }
  const raw = auth.slice(7);
  const hash = createHash("sha256").update(raw).digest("hex");
  const [apiKey] = await db.select().from(apiKeysTable).where(eq(apiKeysTable.keyHash, hash)).limit(1);
  if (!apiKey) {
    return res.status(401).json({ error: { message: "Invalid API key.", type: "invalid_request_error", code: "invalid_api_key" } });
  }
  db.update(apiKeysTable).set({ lastUsedAt: new Date() }).where(eq(apiKeysTable.id, apiKey.id)).catch(() => {});
  req.apiKeyRecord = apiKey;
  next();
}

async function callGrokEndpoint(prompt: string, orientation: string, token: string) {
  const form = new FormData();
  form.append("prompt", prompt);
  form.append("orientation", orientation);
  form.append("num_result", "1");
  const resp = await fetch(`${GEMINIGEN_BASE}/imagen/grok`, {
    method: "POST",
    headers: { Authorization: `Bearer ${token}`, "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) {
  const form = new FormData();
  form.append("prompt", prompt);
  form.append("orientation", orientation);
  form.append("num_result", "1");
  const resp = await fetch(`${GEMINIGEN_BASE}/meta_ai/generate`, {
    method: "POST",
    headers: { Authorization: `Bearer ${token}`, "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, token: string) {
  const form = new FormData();
  form.append("prompt", prompt);
  form.append("model", model);
  form.append("aspect_ratio", aspectRatio);
  form.append("output_format", "jpg");
  const resp = await fetch(`${GEMINIGEN_BASE}/generate_image`, {
    method: "POST",
    headers: { Authorization: `Bearer ${token}`, "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 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}`, "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;
}

router.get("/models", validateApiKey, (_req, res) => {
  const created = Math.floor(Date.now() / 1000);
  res.json({
    object: "list",
    data: SUPPORTED_MODELS.map((id) => ({
      id,
      object: "model",
      created,
      owned_by: "starforge",
    })),
  });
});

router.post("/images/generations", validateApiKey, async (req, res) => {
  const {
    model: rawModel = "grok",
    prompt,
    n = 1,
    size = "1024x1024",
    response_format = "url",
  } = req.body as {
    model?: string;
    prompt?: string;
    n?: number;
    size?: string;
    response_format?: string;
  };

  if (!prompt) {
    return res.status(400).json({ error: { message: "prompt is required", type: "invalid_request_error" } });
  }

  const model = ALIAS_MAP[rawModel] || (SUPPORTED_MODELS.includes(rawModel) ? rawModel : "grok");
  const aspectRatio = SIZE_TO_ASPECT[size] || "1:1";
  const orientation = ORIENTATION_MAP[aspectRatio] || "square";
  const isImagenModel = ["imagen-pro", "imagen-4", "imagen-flash", "nano-banana-pro", "nano-banana-2"].includes(model);

  // 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 handleExpiry(): 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();
  if (!token) return res.status(503).json({ error: { message: "Service not configured.", type: "server_error" } });

  const images: { url?: string; b64_json?: string }[] = [];
  const count = Math.min(Math.max(Number(n) || 1, 1), 4);

  try {
    for (let i = 0; i < count; i++) {
      let result: { status: number; body: unknown };

      if (model === "grok") result = await callGrokEndpoint(prompt, orientation, token);
      else if (model === "meta") result = await callMetaEndpoint(prompt, orientation, token);
      else result = await callImagenEndpoint(model, prompt, aspectRatio, token);

      if (result.status === 401) {
        const newToken = await handleExpiry();
        if (!newToken) throw new Error("Token expired");
        token = newToken;
        if (model === "grok") result = await callGrokEndpoint(prompt, orientation, token);
        else if (model === "meta") result = await callMetaEndpoint(prompt, orientation, token);
        else result = await callImagenEndpoint(model, prompt, aspectRatio, token);
      }

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

      let imageUrl = "";
      if (data.base64_images) {
        imageUrl = `data:image/png;base64,${data.base64_images}`;
      } else if (data.generated_image?.[0]?.image_url) {
        imageUrl = data.generated_image[0].image_url;
      } else if (data.uuid) {
        const polled = await pollForImage(data.uuid, token);
        if (!polled) throw new Error("Generation timed out");
        imageUrl = polled;
      }

      if (response_format === "b64_json" && imageUrl.startsWith("data:")) {
        images.push({ b64_json: imageUrl.split(",")[1] });
      } else {
        images.push({ url: imageUrl });
      }
    }
  } catch (err: unknown) {
    const msg = err instanceof Error ? err.message : String(err);
    return res.status(500).json({ error: { message: msg, type: "server_error" } });
  }

  if (currentAccountId !== null) markAccountUsed(currentAccountId).catch(() => {});
  res.json({ created: Math.floor(Date.now() / 1000), data: images });
});

export default router;