import type * as GoogleGenAi from "@google/genai";
import type { GenerateContentParameters } from "@google/genai";
import { Type } from "typebox";
import { describe, expect, it, vi } from "vitest";

vi.mock("@google/genai", async (importOriginal) => {
	const actual = await importOriginal<typeof GoogleGenAi>();
	class GoogleGenAI {
		models = {
			generateContentStream: async function* () {
				yield {
					candidates: [{ content: { parts: [{ text: '{"decision":"commit"}' }] }, finishReason: "STOP" }],
					usageMetadata: { promptTokenCount: 1, candidatesTokenCount: 1, totalTokenCount: 2 },
				};
			},
		};
	}

	return {
		...actual,
		GoogleGenAI,
		ResourceScope: { COLLECTION: "COLLECTION" },
		ThinkingLevel: {
			THINKING_LEVEL_UNSPECIFIED: "THINKING_LEVEL_UNSPECIFIED",
			MINIMAL: "MINIMAL",
			LOW: "LOW",
			MEDIUM: "MEDIUM",
			HIGH: "HIGH",
		},
	};
});

vi.mock("openai", () => {
	class FakeOpenAI {
		chat = {
			completions: {
				create: () => {
					const stream = {
						async *[Symbol.asyncIterator]() {
							yield {
								choices: [{ delta: { content: '{"decision":"commit"}' }, finish_reason: "stop" }],
								usage: {
									prompt_tokens: 1,
									completion_tokens: 1,
									prompt_tokens_details: { cached_tokens: 0 },
									completion_tokens_details: { reasoning_tokens: 0 },
								},
							};
						},
					};
					const promise = Promise.resolve(stream) as Promise<typeof stream> & {
						withResponse: () => Promise<{ data: typeof stream; response: { status: number; headers: Headers } }>;
					};
					promise.withResponse = async () => ({
						data: stream,
						response: { status: 200, headers: new Headers() },
					});
					return promise;
				},
			},
		};
	}

	return { default: FakeOpenAI };
});

import { getModel } from "../src/models.js";
import { streamGoogle } from "../src/providers/google.js";
import { streamSimple } from "../src/stream.js";
import type { Context, ResponseFormat } from "../src/types.js";

const context: Context = {
	messages: [{ role: "user", content: "commit or hold?", timestamp: Date.now() }],
};

const responseFormat: ResponseFormat = {
	type: "json_schema",
	name: "commitment",
	schema: Type.Object(
		{
			decision: Type.Union([Type.Literal("commit"), Type.Literal("hold")]),
		},
		{ additionalProperties: false },
	),
};

interface SchemaNode {
	type?: string;
	enum?: unknown[];
	const?: unknown;
	anyOf?: SchemaNode[];
	oneOf?: SchemaNode[];
	items?: SchemaNode;
	properties?: Record<string, SchemaNode>;
	$defs?: Record<string, SchemaNode>;
	$ref?: string;
	additionalProperties?: boolean;
}

async function capturePayload<T>(run: (onPayload: (payload: unknown) => void) => Promise<void>): Promise<T> {
	let captured: unknown;
	await run((payload) => {
		captured = payload;
	});
	if (captured === undefined) {
		throw new Error("Expected a provider payload to be captured");
	}
	return captured as T;
}

async function captureGoogleSchema(format: ResponseFormat): Promise<SchemaNode> {
	const model = getModel("google", "gemini-3-flash-preview")!;
	const payload = await capturePayload<GenerateContentParameters>(async (onPayload) => {
		await streamGoogle(model, context, { apiKey: "fake-key", responseFormat: format, onPayload }).result();
	});
	return payload.config?.responseJsonSchema as SchemaNode;
}

describe("structured output response format", () => {
	it("sends responseJsonSchema for google-generative-ai", async () => {
		const model = getModel("google", "gemini-3-flash-preview")!;
		const payload = await capturePayload<GenerateContentParameters>(async (onPayload) => {
			await streamSimple(model, context, { apiKey: "fake-key", responseFormat, onPayload }).result();
		});

		expect(payload.config?.responseMimeType).toBe("application/json");
		expect(payload.config?.responseSchema).toBeUndefined();
		expect(payload.config?.responseJsonSchema).toMatchObject({
			type: "object",
			additionalProperties: false,
			properties: { decision: { type: "string", enum: ["commit", "hold"] } },
		});
	});

	it("sends responseJsonSchema for google-vertex", async () => {
		const model = getModel("google-vertex", "gemini-2.5-flash")!;
		const payload = await capturePayload<GenerateContentParameters>(async (onPayload) => {
			await streamSimple(model, context, { apiKey: "fake-key", responseFormat, onPayload }).result();
		});

		expect(payload.config?.responseMimeType).toBe("application/json");
		expect(payload.config?.responseJsonSchema).toMatchObject({
			properties: { decision: { type: "string", enum: ["commit", "hold"] } },
		});
	});

	it("omits the schema fields for google when responseFormat is absent", async () => {
		const model = getModel("google", "gemini-3-flash-preview")!;
		const payload = await capturePayload<GenerateContentParameters>(async (onPayload) => {
			await streamSimple(model, context, { apiKey: "fake-key", onPayload }).result();
		});

		expect(payload.config?.responseMimeType).toBeUndefined();
		expect(payload.config?.responseJsonSchema).toBeUndefined();
	});

	it("normalizes a TypeBox literal union to enum and keeps declaration order", async () => {
		const schema = await captureGoogleSchema({
			type: "json_schema",
			schema: Type.Object({
				evidence: Type.String(),
				status: Type.Union([Type.Literal("done"), Type.Literal("pending"), Type.Literal("blocked")]),
			}),
		});

		expect(Object.keys(schema.properties ?? {})).toEqual(["evidence", "status"]);
		const status = schema.properties?.status;
		expect(status?.enum).toEqual(["done", "pending", "blocked"]);
		expect(status?.type).toBe("string");
		expect(status?.anyOf).toBeUndefined();
		expect(status?.const).toBeUndefined();
	});

	it("normalizes a single Type.Literal to a one-value enum", async () => {
		const schema = await captureGoogleSchema({
			type: "json_schema",
			schema: Type.Object({ kind: Type.Literal("commitment") }),
		});

		const kind = schema.properties?.kind;
		expect(kind?.enum).toEqual(["commitment"]);
		expect(kind?.type).toBe("string");
		expect(kind?.const).toBeUndefined();
	});

	it("leaves a genuine non-const union unchanged", async () => {
		const schema = await captureGoogleSchema({
			type: "json_schema",
			schema: Type.Object({
				value: Type.Union([Type.String(), Type.Number()]),
				mixed: Type.Union([Type.Literal("auto"), Type.Number()]),
			}),
		});

		expect(schema.properties?.value?.anyOf).toEqual([{ type: "string" }, { type: "number" }]);
		expect(schema.properties?.value?.enum).toBeUndefined();
		const mixed = schema.properties?.mixed;
		expect(mixed?.enum).toBeUndefined();
		expect(mixed?.anyOf?.[0]).toEqual({ type: "string", enum: ["auto"] });
		expect(mixed?.anyOf?.[1]).toEqual({ type: "number" });
	});

	it("normalizes unions nested in array items and in $defs", async () => {
		const schema = await captureGoogleSchema({
			type: "json_schema",
			schema: Type.Object({
				labels: Type.Array(Type.Union([Type.Literal("bug"), Type.Literal("feat")])),
				status: Type.Union([Type.Literal("open"), Type.Literal("closed")]),
			}),
		});

		const withDefs = await captureGoogleSchema({
			type: "json_schema",
			schema: {
				type: "object",
				properties: { status: { $ref: "#/$defs/Status" } },
				$defs: { Status: Type.Union([Type.Literal("open"), Type.Literal("closed")]) },
			},
		});

		expect(schema.properties?.labels?.items?.enum).toEqual(["bug", "feat"]);
		expect(schema.properties?.labels?.items?.anyOf).toBeUndefined();
		expect(withDefs.$defs?.Status?.enum).toEqual(["open", "closed"]);
		expect(withDefs.$defs?.Status?.anyOf).toBeUndefined();
		expect(withDefs.properties?.status?.$ref).toBe("#/$defs/Status");
	});

	it("sends response_format json_schema for openai-completions", async () => {
		const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!;
		const model = { ...baseModel, api: "openai-completions" } as const;
		const payload = await capturePayload<{
			response_format?: { type: string; json_schema: { name: string; schema: unknown; strict?: boolean } };
		}>(async (onPayload) => {
			await streamSimple(model, context, { apiKey: "fake-key", responseFormat, onPayload }).result();
		});

		expect(payload.response_format?.type).toBe("json_schema");
		expect(payload.response_format?.json_schema.name).toBe("commitment");
		expect(payload.response_format?.json_schema.schema).toMatchObject({
			type: "object",
			additionalProperties: false,
			properties: { decision: { type: "string", enum: ["commit", "hold"] } },
		});
		expect(payload.response_format?.json_schema.strict).toBe(true);
	});

	it("normalizes literal unions and keeps genuine unions for openai-completions", async () => {
		const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!;
		const model = { ...baseModel, api: "openai-completions" } as const;
		const payload = await capturePayload<{
			response_format?: { json_schema: { schema: SchemaNode } };
		}>(async (onPayload) => {
			await streamSimple(model, context, {
				apiKey: "fake-key",
				responseFormat: {
					type: "json_schema",
					schema: Type.Object({
						status: Type.Union([Type.Literal("open"), Type.Literal("closed")]),
						value: Type.Union([Type.String(), Type.Number()]),
					}),
				},
				onPayload,
			}).result();
		});

		const schema = payload.response_format?.json_schema.schema;
		expect(schema?.properties?.status).toEqual({ type: "string", enum: ["open", "closed"] });
		expect(schema?.properties?.value?.anyOf).toEqual([{ type: "string" }, { type: "number" }]);
		expect(schema?.properties?.value?.enum).toBeUndefined();
	});

	it("omits response_format for non-standard openai-compatible endpoints by default", async () => {
		const model = getModel("deepseek", "deepseek-flash")!;
		const payload = await capturePayload<{ response_format?: unknown }>(async (onPayload) => {
			await streamSimple(model, context, { apiKey: "fake-key", responseFormat, onPayload }).result();
		});

		expect(payload.response_format).toBeUndefined();
	});

	it("keeps response_format for opencode-go, which is measured to support it", async () => {
		const model = getModel("opencode-go", "deepseek-v4-flash")!;
		const payload = await capturePayload<{ response_format?: { type: string } }>(async (onPayload) => {
			await streamSimple(model, context, { apiKey: "fake-key", responseFormat, onPayload }).result();
		});

		expect(payload.response_format?.type).toBe("json_schema");
	});

	it("sends response_format for a non-standard endpoint that opts in", async () => {
		const baseModel = getModel("deepseek", "deepseek-flash")!;
		const model = {
			...baseModel,
			compat: { ...baseModel.compat, supportsResponseFormatJsonSchema: true },
		} as const;
		const payload = await capturePayload<{ response_format?: { type: string } }>(async (onPayload) => {
			await streamSimple(model, context, { apiKey: "fake-key", responseFormat, onPayload }).result();
		});

		expect(payload.response_format?.type).toBe("json_schema");
	});

	it("omits response_format when compat disables json_schema support", async () => {
		const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!;
		const model = {
			...baseModel,
			api: "openai-completions",
			compat: { supportsResponseFormatJsonSchema: false },
		} as const;
		const payload = await capturePayload<{ response_format?: unknown }>(async (onPayload) => {
			await streamSimple(model, context, { apiKey: "fake-key", responseFormat, onPayload }).result();
		});

		expect(payload.response_format).toBeUndefined();
	});
});