import ZAI from "z-ai-web-dev-sdk";

let _zai: any = null;

export async function getAI() {
  if (_zai) return _zai;
  _zai = await ZAI.create();
  return _zai;
}

/**
 * Generic chat completion helper.
 * Returns the assistant message text.
 */
export async function aiChat(
  systemPrompt: string,
  userPrompt: string,
  options?: { temperature?: number; maxTokens?: number },
): Promise<string> {
  const ai = await getAI();
  const res = await ai.chat.completions.create({
    messages: [
      { role: "system", content: systemPrompt },
      { role: "user", content: userPrompt },
    ],
    temperature: options?.temperature ?? 0.7,
    max_tokens: options?.maxTokens ?? 1200,
  });
  // z-ai sdk returns choices array
  const text =
    res?.choices?.[0]?.message?.content ??
    res?.choices?.[0]?.delta?.content ??
    (typeof res === "string" ? res : "");
  return typeof text === "string" ? text : String(text ?? "");
}

/**
 * Returns a JSON object from the model. Falls back to raw text on parse failure.
 */
export async function aiJSON<T = any>(
  systemPrompt: string,
  userPrompt: string,
): Promise<T> {
  const raw = await aiChat(
    systemPrompt +
      "\n\nمهم: پاسخ را فقط به صورت یک JSON معتبر و بدون متن اضافه برگردان. هیچ markdown یا توضیحی ننویس.",
    userPrompt,
    { temperature: 0.4 },
  );
  // strip code fences if present
  const cleaned = raw
    .replace(/```json\s*/gi, "")
    .replace(/```\s*$/g, "")
    .trim();
  try {
    return JSON.parse(cleaned) as T;
  } catch {
    // try to extract first {...} block
    const match = cleaned.match(/\{[\s\S]*\}/);
    if (match) {
      try {
        return JSON.parse(match[0]) as T;
      } catch {
        /* ignore */
      }
    }
    throw new Error("AI did not return valid JSON: " + cleaned.slice(0, 200));
  }
}