@kody/ai

import {
	completePreview,
	generateText,
	resolveComplete,
	type CompleteInput,
	type CompletePreview,
} from './complete.ts'
import { isPreviewDryRun } from './validation.ts'

export type JsonSchema = Record<string, unknown>

export type GenerateObjectInput = CompleteInput & {
	schema: JsonSchema
	schemaName?: string
	schemaDescription?: string
}

export type GenerateObjectResult = {
	ok: boolean
	error?: string
	object: Record<string, unknown> | Array<unknown> | null
	text: string
	provider: string
	model: string
	host: string
}

function parseJson(text: string): unknown {
	const trimmed = text.trim().replace(/^```(?:json)?\s*/i, '').replace(/\s*```$/i, '')
	return JSON.parse(trimmed)
}

/**
 * Structured generation from a JSON Schema. Pass `dryRun: true` to preview.
 *
 * @example
 * import generateObject from 'kody:@kody/ai/generate-object'
 * const preview = await generateObject({
 *   provider: 'openai',
 *   messages: [{ role: 'user', content: 'Label: great job' }],
 *   schema: {
 *     type: 'object',
 *     properties: { label: { type: 'string' } },
 *     required: ['label'],
 *   },
 *   dryRun: true,
 * })
 */
export async function generateObject(
	input: GenerateObjectInput,
): Promise<GenerateObjectResult | CompletePreview> {
	if (!input?.schema || typeof input.schema !== 'object') {
		throw new Error('schema must be a JSON Schema object.')
	}
	const schemaName = input.schemaName?.trim() || 'result'
	const resolved = await resolveComplete(input)
	const request =
		resolved.provider === 'anthropic'
			? {
					...input,
					tools: [
						{
							name: schemaName,
							description:
								input.schemaDescription ||
								'Return a single object that matches the required JSON schema.',
							input_schema: input.schema,
						},
					],
					toolChoice: { type: 'tool', name: schemaName },
				}
			: {
					...input,
					responseFormat: {
						type: 'json_schema',
						json_schema: {
							name: schemaName,
							description: input.schemaDescription,
							schema: input.schema,
							strict: true,
						},
					},
				}

	if (isPreviewDryRun(input)) {
		const previewResolved = await resolveComplete(request)
		return completePreview(previewResolved)
	}

	const result = await generateText(request)
	if ('dryRun' in result) return result
	if (!result.ok) {
		return {
			ok: false,
			error: result.error,
			object: null,
			text: result.text,
			provider: result.provider,
			model: result.model,
			host: result.host,
		}
	}

	if (result.toolCalls.length === 1) {
		return {
			ok: true,
			object: result.toolCalls[0].input,
			text: result.text,
			provider: result.provider,
			model: result.model,
			host: result.host,
		}
	}

	try {
		const object = parseJson(result.text)
		if (object === null || typeof object !== 'object') {
			throw new Error('Model did not return a JSON object or array.')
		}
		return {
			ok: true,
			object,
			text: result.text,
			provider: result.provider,
			model: result.model,
			host: result.host,
		}
	} catch (error) {
		return {
			ok: false,
			error:
				(error instanceof Error ? error.message : String(error)) +
				' Raw text: ' +
				result.text.slice(0, 400),
			object: null,
			text: result.text,
			provider: result.provider,
			model: result.model,
			host: result.host,
		}
	}
}

export default generateObject