// Async task polling infrastructure
// All z-image.ai generate endpoints return a taskId, then you poll status until done

import { getClient } from "./client.js";
import type { ApiResponse, StatusData } from "./types.js";
import { extractOutputUrl } from "./types.js";
import * as fs from "fs";
import * as path from "path";
import * as os from "os";

/** Save a base64 data URI to a temp file; return the file path. */
function saveDataUri(dataUri: string, taskId: string): string {
  const match = dataUri.match(/^data:([^;]+);base64,(.+)$/);
  if (!match) return dataUri;
  const [, mimeType, base64] = match;
  const ext = mimeType.split("/")[1]?.split("+")[0] || "bin";
  const tmpDir = process.env.TMPDIR || os.tmpdir();
  const outPath = path.join(tmpDir, `zimage-${taskId}-${Date.now()}.${ext}`);
  fs.writeFileSync(outPath, Buffer.from(base64, "base64"));
  return outPath;
}

const POLL_INTERVAL_MS = 2500;
const MAX_POLLS = 180; // 180 * 2.5s = 7.5 minutes max

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

export interface PollResult {
  taskId: string;
  output: string;
  raw?: StatusData;
}

/**
 * Submit a generation task and poll until completion.
 *
 * @param endpoint - The POST endpoint (e.g., "/api/ai/generate")
 * @param scenario - The scenario string for status polling
 * @param body - The request body
 * @param maxPolls - Max poll attempts (default 180)
 * @returns The output URL(s) from the completed task
 */
export async function submitAndPoll(
  endpoint: string,
  scenario: string,
  body: Record<string, unknown>,
  maxPolls: number = MAX_POLLS,
): Promise<PollResult> {
  const client = getClient();

  // Submit the generation request
  const genResponse = await client.post<ApiResponse<{ taskId: string }>>(endpoint, body);

  if (!genResponse.data?.taskId) {
    // Check for insufficient credits
    const d = genResponse.data as Record<string, unknown> | undefined;
    if (d && d.hasEnoughCredits === false) {
      throw new Error(`Insufficient credits: need ${d.requiredCredits}, have ${d.remainingCredits}`);
    }
    // Some endpoints may return result directly
    const directOutput = extractOutputUrl(genResponse.data as unknown as StatusData);
    if (directOutput) {
      return { taskId: "sync", output: directOutput };
    }
    throw new Error(`No taskId in response: ${JSON.stringify(genResponse).slice(0, 500)}`);
  }

  const taskId = genResponse.data.taskId;

  // Poll for completion
  for (let i = 0; i < maxPolls; i++) {
    await sleep(POLL_INTERVAL_MS);

    const statusResponse = await client.get<ApiResponse<StatusData>>(
      `/api/ai/status?taskId=${encodeURIComponent(taskId)}&scenario=${encodeURIComponent(scenario)}`,
    );

    const data = statusResponse.data;
    if (!data) continue;

    const status = data.status?.toLowerCase?.() || "";

    if (status === "failed" || status === "error") {
      throw new Error(`Task ${taskId} failed: ${JSON.stringify(data).slice(0, 500)}`);
    }

    const outputUrl = extractOutputUrl(data);
    if (outputUrl || status === "completed" || status === "succeeded") {
      const resolved = outputUrl?.startsWith("data:") ? saveDataUri(outputUrl, taskId) : (outputUrl || "");
      return {
        taskId,
        output: resolved,
        raw: data,
      };
    }
  }

  throw new Error(`Task ${taskId} polling timeout after ${maxPolls * POLL_INTERVAL_MS / 1000}s`);
}