gai/src/ai.ts

import type { Config, StreamCallbacks } from "./types";

interface ChatMessage {
  role: "system" | "user" | "assistant";
  content: string;
}

const MAX_RETRIES = 3;
const RETRY_DELAY_MS = 1000;

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

function cleanMessage(raw: string): string {
  let msg = raw.trim();
  // Strip code fences if the whole response is wrapped
  if (msg.startsWith("```") && msg.endsWith("```")) {
    const lines = msg.split("\n");
    if (lines.length > 2) {
      lines.shift();
      lines.pop();
      msg = lines.join("\n").trim();
    }
  }
  return msg;
}

export async function callAI(
  config: Config,
  systemPrompt: string,
  userPrompt: string,
  callbacks?: StreamCallbacks,
): Promise<string> {
  const url = `${config.apiBase.replace(/\/$/, "")}/chat/completions`;
  const stream = callbacks != null;

  const messages: ChatMessage[] = [
    { role: "system", content: systemPrompt },
    { role: "user", content: userPrompt },
  ];

  for (let attempt = 1; attempt <= MAX_RETRIES; attempt++) {
    try {
      const response = await fetch(url, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          Authorization: `Bearer ${config.apiKey}`,
        },
        body: JSON.stringify({
          model: config.model,
          max_tokens: config.maxTokens,
          temperature: config.temperature,
          messages,
          stream,
        }),
      });

      if (!response.ok) {
        const text = await response.text();
        if (response.status === 429 && attempt < MAX_RETRIES) {
          await sleep(RETRY_DELAY_MS * attempt);
          continue;
        }
        throw new Error(`API request failed (${response.status}): ${text}`);
      }

      if (stream && response.body) {
        return await readStream(response.body, callbacks!);
      }

      const data = (await response.json()) as {
        choices?: Array<{ message?: { content?: string | null }; finish_reason?: string }>;
        error?: { message?: string; type?: string; code?: string };
      };

      if (data.error) {
        throw new Error(`API error: ${data.error.message ?? JSON.stringify(data.error)}`);
      }

      const raw = data.choices?.[0]?.message?.content;
      const finishReason = data.choices?.[0]?.finish_reason;

      if (raw && raw.trim()) return raw;

      if (finishReason === "length") {
        throw new Error("Response truncated (max_tokens too low). Try increasing GAI_MAX_TOKENS.");
      }
      if (finishReason === "content_filter") {
        throw new Error("Response blocked by content filter.");
      }
      if (attempt < MAX_RETRIES) {
        await sleep(RETRY_DELAY_MS * attempt);
        continue;
      }

      throw new Error(
        `Empty response from AI after ${MAX_RETRIES} attempts. finish_reason: ${finishReason ?? "unknown"}`,
      );
    } catch (err) {
      if (attempt >= MAX_RETRIES) throw err;
      if (err instanceof Error && err.message.startsWith("API error")) throw err;
      if (err instanceof Error && err.message.includes("max_tokens")) throw err;
      if (err instanceof Error && err.message.includes("content filter")) throw err;
      await sleep(RETRY_DELAY_MS * attempt);
    }
  }

  throw new Error("Failed to generate response");
}

async function readStream(body: ReadableStream<Uint8Array>, callbacks: StreamCallbacks): Promise<string> {
  const reader = body.getReader();
  const decoder = new TextDecoder();
  let fullText = "";
  let buffer = "";

  try {
    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) {
        const trimmed = line.trim();
        if (!trimmed || !trimmed.startsWith("data:")) continue;

        const data = trimmed.slice(5).trim();
        if (data === "[DONE]") continue;

        try {
          const parsed = JSON.parse(data) as {
            choices?: Array<{ delta?: { content?: string }; finish_reason?: string }>;
            error?: { message?: string };
          };
          if (parsed.error) {
            callbacks.onError?.(new Error(`Stream error: ${parsed.error.message ?? "unknown"}`));
            continue;
          }
          const token = parsed.choices?.[0]?.delta?.content;
          if (token) {
            fullText += token;
            callbacks.onToken?.(token);
          }
          const finishReason = parsed.choices?.[0]?.finish_reason;
          if (finishReason === "length") {
            callbacks.onError?.(new Error("Response truncated (max_tokens too low)."));
          }
        } catch {
          // Skip unparseable SSE lines
        }
      }
    }
  } finally {
    try { await reader.cancel(); } catch {}
    // releaseLock is not needed after cancel
  }

  callbacks.onDone?.(fullText);
  return fullText;
}

export async function generateCommitMessage(
  config: Config,
  systemPrompt: string,
  userPrompt: string,
  callbacks?: StreamCallbacks,
): Promise<string> {
  const raw = await callAI(config, systemPrompt, userPrompt, callbacks);
  return cleanMessage(raw);
}

export async function generatePRMessage(
  config: Config,
  systemPrompt: string,
  userPrompt: string,
  callbacks?: StreamCallbacks,
): Promise<{ title: string; body: string }> {
  const raw = await callAI(config, systemPrompt, userPrompt, callbacks);
  const cleaned = cleanMessage(raw);

  const lines = cleaned.split("\n");
  const title = lines[0]?.trim() || "Update";
  let bodyStart = 1;

  while (bodyStart < lines.length && lines[bodyStart]?.trim() === "") {
    bodyStart++;
  }

  const body = lines.slice(bodyStart).join("\n").trim();
  return { title, body };
}