client.ts

// Z.AI API client — thin wrappers over the non-chat endpoints.
//
// These endpoints are NOT chat completions; they are one-shot media/parse
// utilities. See https://docs.z.ai/api-reference/introduction
//
//   /images/generations  — text → image (GLM-Image)
//   /layout_parsing      — image/PDF → markdown layout (GLM-OCR)
//
// Auth: HTTP Bearer. The key is resolved at call time — reuses the Pi `zai`
// provider key from /login, falling back to $ZAI_API_KEY.

import type { ExtensionContext } from "@earendil-works/pi-coding-agent";

const DEFAULT_BASE_URL = "https://api.z.ai/api/paas/v4";

/** Resolve a Z.AI API key. Prefers the Pi-configured `zai` key, falls back to env. */
export async function resolveApiKey(
  ctx: ExtensionContext,
): Promise<string> {
  // Reuse the key the user configured via `/login zai` (seamless — no extra setup).
  try {
    const key = await ctx.modelRegistry.getApiKeyForProvider("zai");
    if (key && key.trim()) return key.trim();
  } catch {
    // modelRegistry resolution may fail in some contexts; fall through to env.
  }

  const envKey = process.env.ZAI_API_KEY;
  if (envKey && envKey.trim()) return envKey.trim();

  throw new Error(
    "No Z.AI API key found. Either:\n" +
      "  1. Run `/login` in pi and configure the `zai` (ZAI Coding Plan) provider, or\n" +
      "  2. Export ZAI_API_KEY in your shell.\n" +
      "Get a key at https://z.ai/manage-apikey/apikey-list",
  );
}

function baseUrl(): string {
  return process.env.ZAI_BASE_URL?.trim() || DEFAULT_BASE_URL;
}

export interface ZaiError extends Error {
  status: number;
  body: string;
}

async function zaiRequest<T>(
  path: string,
  body: Record<string, unknown>,
  apiKey: string,
  signal?: AbortSignal,
): Promise<T> {
  const res = await fetch(`${baseUrl()}${path}`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${apiKey}`,
      "Accept-Language": "en-US,en",
    },
    body: JSON.stringify(body),
    signal,
  });

  const text = await res.text();
  if (!res.ok) {
    const err = new Error(
      `Z.AI request to ${path} failed (HTTP ${res.status}): ${truncate(text, 500)}`,
    ) as ZaiError;
    err.status = res.status;
    err.body = text;
    throw err;
  }

  try {
    return JSON.parse(text) as T;
  } catch {
    throw new Error(`Z.AI returned non-JSON from ${path}: ${truncate(text, 200)}`);
  }
}

// ─── Image generation ──────────────────────────────────────────────────────

export interface GenerateImageOptions {
  prompt: string;
  model?: string; // default: glm-image
  size?: string; // default: 1280x1280
  apiKey: string;
  signal?: AbortSignal;
}

export interface GenerateImageResult {
  url: string;
  model: string;
}

interface ImageApiResponse {
  model?: string;
  data?: Array<{ url?: string; b64_json?: string }>;
}

export async function generateImage(
  opts: GenerateImageOptions,
): Promise<GenerateImageResult> {
  const model = opts.model?.trim() || "glm-image";
  const json = await zaiRequest<ImageApiResponse>(
    "/images/generations",
    {
      model,
      prompt: opts.prompt,
      size: opts.size?.trim() || "1280x1280",
    },
    opts.apiKey,
    opts.signal,
  );

  const entry = json.data?.[0];
  const url = entry?.url;
  if (!url) {
    throw new Error(
      `Z.AI image generation returned no image URL. Response: ${truncate(JSON.stringify(json), 300)}`,
    );
  }
  return { url, model: json.model ?? model };
}

// ─── OCR / layout parsing ──────────────────────────────────────────────────

export interface ParseLayoutOptions {
  /** Image or PDF — URL or base64 data URI. */
  file: string;
  model?: string; // default: glm-ocr
  returnCropImages?: boolean;
  apiKey: string;
  signal?: AbortSignal;
}

export interface ParseLayoutResult {
  model: string;
  /** Recognized content in Markdown. */
  content: string;
  /** URLs of recognized/segment images, if returned. */
  images: string[];
  segments: unknown[];
}

interface LayoutApiResponse {
  model?: string;
  content?: string;
  images?: string[];
  segments?: unknown[];
  choices?: unknown;
}

export async function parseLayout(
  opts: ParseLayoutOptions,
): Promise<ParseLayoutResult> {
  const model = opts.model?.trim() || "glm-ocr";
  const body: Record<string, unknown> = {
    model,
    file: opts.file,
  };
  if (opts.returnCropImages !== undefined) {
    body.return_crop_images = opts.returnCropImages;
  }

  const json = await zaiRequest<LayoutApiResponse>(
    "/layout_parsing",
    body,
    opts.apiKey,
    opts.signal,
  );

  const content =
    typeof json.content === "string"
      ? json.content
      : // Some responses nest content under choices[].message.content; fall back gracefully.
        extractNestedContent(json) ?? "";

  return {
    model: json.model ?? model,
    content,
    images: Array.isArray(json.images) ? json.images : [],
    segments: Array.isArray(json.segments) ? json.segments : [],
  };
}

function extractNestedContent(json: LayoutApiResponse): string | undefined {
  const choices = json.choices as
    | Array<{ message?: { content?: string } }>
    | undefined;
  return choices?.[0]?.message?.content;
}

// ─── helpers ───────────────────────────────────────────────────────────────

function truncate(s: string, max: number): string {
  return s.length <= max ? s : `${s.slice(0, max)}…[truncated]`;
}