api.ts

import { loadAuth } from "./auth.js";

const ANTIGRAVITY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com";

// Headers mimicking the Cloud Code extension
const HEADERS = {
  "User-Agent": "antigravity/1.11.5 darwin/arm64",
  "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
  "Client-Metadata": JSON.stringify({
    ideType: "IDE_UNSPECIFIED",
    platform: "PLATFORM_UNSPECIFIED",
    pluginType: "GEMINI",
  }),
};

export async function enhancePrompt(draft: string): Promise<string> {
  const auth = loadAuth();
  if (!auth?.access || !auth?.projectId) throw new Error("No Antigravity credentials found.");

  const model = "gemini-3-flash";
  const url = `${ANTIGRAVITY_ENDPOINT}/v1internal:streamGenerateContent?alt=sse`;
  
  const systemPrompt = "You are an expert AI art prompter. Take the user's rough idea and transform it into a highly detailed, descriptive prompt suitable for a high-end image generation model. Focus on lighting, style, composition, and texture. Output ONLY the enhanced prompt, no preamble.";

  const body = {
    project: auth.projectId,
    model,
    request: {
      contents: [{ role: "user", parts: [{ text: draft }] }],
      systemInstruction: { parts: [{ text: systemPrompt }] },
    },
    requestType: "agent",
    userAgent: "antigravity",
  };

  const response = await fetch(url, {
    method: "POST",
    headers: {
      "Authorization": `Bearer ${auth.access}`,
      "Content-Type": "application/json",
      "Accept": "text/event-stream",
      ...HEADERS,
    },
    body: JSON.stringify(body),
  });

  if (!response.ok) throw new Error(`Prompt enhancement failed: ${response.status}`);

  // Simple SSE parser for text
  const reader = response.body?.getReader();
  if (!reader) throw new Error("No response body");
  
  const decoder = new TextDecoder();
  let fullText = "";
  let buffer = "";

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    
    buffer += decoder.decode(value, { stream: true });
    const lines = buffer.split("\n");
    buffer = lines.pop() || "";

    for (const line of lines) {
      if (!line.startsWith("data:")) continue;
      try {
        const json = JSON.parse(line.slice(5).trim());
        const text = json.response?.candidates?.[0]?.content?.parts?.[0]?.text;
        if (text) fullText += text;
      } catch {}
    }
  }
  
  return fullText.trim();
}

export async function generateImage(prompt: string, aspectRatio: string = "1:1", stylePrompt: string = ""): Promise<{ data: string, mimeType: string }> {
  const auth = loadAuth();
  if (!auth?.access || !auth?.projectId) throw new Error("No Antigravity credentials found.");

  const model = "gemini-3-pro-image";
  const url = `${ANTIGRAVITY_ENDPOINT}/v1internal:streamGenerateContent?alt=sse`;

  const request: any = {
    contents: [{ role: "user", parts: [{ text: prompt }] }],
    generationConfig: {
      imageConfig: { aspectRatio },
      candidateCount: 1,
    },
  };

  // Add system instruction if style prompt is present
  if (stylePrompt.trim()) {
    request.systemInstruction = {
      parts: [{ text: stylePrompt.trim() }]
    };
  }

  const body = {
    project: auth.projectId,
    model,
    request,
    requestType: "agent",
    userAgent: "antigravity",
  };

  const response = await fetch(url, {
    method: "POST",
    headers: {
      "Authorization": `Bearer ${auth.access}`,
      "Content-Type": "application/json",
      "Accept": "text/event-stream",
      ...HEADERS,
    },
    body: JSON.stringify(body),
  });

  if (!response.ok) throw new Error(`Image generation failed: ${response.status}`);

  const reader = response.body?.getReader();
  if (!reader) throw new Error("No response body");
  
  const decoder = new TextDecoder();
  let buffer = "";

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    
    buffer += decoder.decode(value, { stream: true });
    const lines = buffer.split("\n");
    buffer = lines.pop() || "";

    for (const line of lines) {
      if (!line.startsWith("data:")) continue;
      try {
        const json = JSON.parse(line.slice(5).trim());
        const part = json.response?.candidates?.[0]?.content?.parts?.[0];
        if (part?.inlineData?.data) {
          return {
            data: part.inlineData.data,
            mimeType: part.inlineData.mimeType || "image/png"
          };
        }
      } catch {}
    }
  }
  
  throw new Error("No image data received from API");
}