pi-ext

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

const CODEX_API_ENDPOINT = "https://chatgpt.com/backend-api/codex/responses";
const DEFAULT_MODEL = "gpt-5.5";
const MAX_SOURCES = 10;
const JWT_CLAIM_PATH = "https://api.openai.com/auth";

interface CodexSearchSource {
	title: string;
	url: string;
	snippet: string;
}

interface CodexSearchOutput {
	summary?: string;
	sources?: CodexSearchSource[];
}

function isRecord(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

function decodeJwtPayload(token: string): Record<string, unknown> | undefined {
	try {
		const payload = token.split(".")[1];
		if (!payload) return undefined;
		return JSON.parse(Buffer.from(payload, "base64url").toString("utf8"));
	} catch {
		return undefined;
	}
}

function accountIdFromToken(token: string): string | undefined {
	const payload = decodeJwtPayload(token);
	const auth = payload?.[JWT_CLAIM_PATH];
	if (!isRecord(auth)) return undefined;
	return typeof auth.chatgpt_account_id === "string"
		? auth.chatgpt_account_id
		: undefined;
}

async function codexAuth(
	ctx: ExtensionContext,
): Promise<{ accessToken: string; accountId: string } | undefined> {
	const accessToken = await ctx.modelRegistry
		.getApiKeyForProvider("openai-codex")
		.catch(() => undefined);
	if (!accessToken) return undefined;
	const accountId = accountIdFromToken(accessToken);
	return accountId ? { accessToken, accountId } : undefined;
}

export async function isCodexAvailable(
	ctx?: ExtensionContext,
): Promise<boolean> {
	if (!ctx) return false;
	return (await codexAuth(ctx)) !== undefined;
}

function schema(): string {
	return JSON.stringify({
		type: "object",
		additionalProperties: false,
		properties: {
			summary: { type: "string" },
			sources: {
				type: "array",
				items: {
					type: "object",
					additionalProperties: false,
					properties: {
						title: { type: "string" },
						url: { type: "string" },
						snippet: { type: "string" },
					},
					required: ["title", "url", "snippet"],
				},
			},
		},
		required: ["summary", "sources"],
	});
}

function buildPrompt(query: string, count: number): string {
	return [
		"Search the public web for the user's query.",
		"Return ONLY JSON matching this schema:",
		schema(),
		"Do not wrap JSON in markdown fences. Do not add commentary.",
		`Limit sources to at most ${count}. Prefer official or primary sources.`,
		"Keep snippets short and directly relevant.",
		"",
		`Query: ${query}`,
	].join("\n");
}

function parseSseBlock(
	block: string,
): { type: string; data: Record<string, unknown> } | undefined {
	let eventType = "";
	const dataLines: string[] = [];

	for (const raw of block.split("\n")) {
		const line = raw.endsWith("\r") ? raw.slice(0, -1) : raw;
		if (line.startsWith("event: ")) eventType = line.slice(7).trim();
		else if (line.startsWith("data: ")) dataLines.push(line.slice(6));
		else if (line.startsWith("data:")) dataLines.push(line.slice(5).trim());
	}

	const text = dataLines.join("\n").trim();
	if (!text || text === "[DONE]") return undefined;

	try {
		const data = JSON.parse(text) as Record<string, unknown>;
		const type = eventType || (typeof data.type === "string" ? data.type : "");
		return type ? { type, data } : undefined;
	} catch {
		return undefined;
	}
}

async function readOutputText(
	response: Response,
	signal?: AbortSignal,
): Promise<string> {
	if (!response.body) throw new Error("Codex search response has no body.");
	const reader = response.body.getReader();
	const decoder = new TextDecoder();
	let buffer = "";
	let output = "";

	try {
		while (true) {
			if (signal?.aborted) throw new Error("Codex search aborted.");
			const { value, done } = await reader.read();
			if (done) break;
			buffer += decoder.decode(value, { stream: true });

			let idx = buffer.search(/\r?\n\r?\n/);
			while (idx >= 0) {
				const block = buffer.slice(0, idx);
				const sep = buffer.slice(idx).match(/^\r?\n\r?\n/)?.[0] ?? "\n\n";
				buffer = buffer.slice(idx + sep.length);
				const event = parseSseBlock(block);
				if (
					event?.type === "response.output_text.delta" &&
					typeof event.data.delta === "string"
				) {
					output += event.data.delta;
				}
				idx = buffer.search(/\r?\n\r?\n/);
			}
		}

		const trailing = parseSseBlock(buffer.trim());
		if (
			trailing?.type === "response.output_text.delta" &&
			typeof trailing.data.delta === "string"
		) {
			output += trailing.data.delta;
		}
	} finally {
		reader.releaseLock();
	}

	return output;
}

function asResults(parsed: CodexSearchOutput, count: number): SearchResult[] {
	const results = (parsed.sources ?? []).slice(0, count).map((source) => ({
		title: source.title,
		url: source.url,
		snippet: source.snippet,
	}));

	if (results.length > 0 || !parsed.summary?.trim()) return results;
	return [
		{ title: "Codex search summary", url: "", snippet: parsed.summary.trim() },
	];
}

export async function codexSearch(
	query: string,
	ctx: ExtensionContext,
	signal?: AbortSignal,
	count = 5,
): Promise<SearchResult[]> {
	const auth = await codexAuth(ctx);
	if (!auth) throw new Error("OpenAI Codex auth unavailable.");

	const maxSources = Math.min(Math.max(Math.trunc(count), 1), MAX_SOURCES);
	const response = await fetch(CODEX_API_ENDPOINT, {
		method: "POST",
		headers: {
			"Content-Type": "application/json",
			Authorization: `Bearer ${auth.accessToken}`,
			"ChatGPT-Account-ID": auth.accountId,
			accept: "text/event-stream",
		},
		body: JSON.stringify({
			model: DEFAULT_MODEL,
			instructions: buildPrompt(query, maxSources),
			input: [{ role: "user", content: `Search the web for: ${query}` }],
			tools: [{ type: "web_search" }],
			store: false,
			stream: true,
		}),
		signal,
	});

	if (!response.ok) {
		const text = await response.text().catch(() => "");
		throw new Error(`Codex search ${response.status}: ${text.slice(0, 200)}`);
	}

	const raw = (await readOutputText(response, signal)).trim();
	if (!raw) throw new Error("Codex search returned no output.");

	const parsed = JSON.parse(raw) as CodexSearchOutput;
	if (!Array.isArray(parsed.sources))
		throw new Error("Codex search returned invalid JSON.");
	return asResults(parsed, maxSources);
}

export const _test = { accountIdFromToken, parseSseBlock };