import z from 'zod';
import { search } from '../search.js';
import { imageSearchPrompt, imageSearchFewShots } from '../prompts/image.js';
import formatChatHistoryAsString from '../utils/formatHistory.js';
import type { BaseLLM } from '../models/base/llm.js';
import type { ChatTurnMessage } from '../types.js';

type Input = { chatHistory: ChatTurnMessage[]; query: string };

export type ImageResult = { img_src: string; url: string; title: string };

export async function searchImages(input: Input, llm: BaseLLM): Promise<ImageResult[]> {
  const schema = z.object({ query: z.string().describe('The image search query.') });

  const res = await llm.generateObject<{ query: string }>({
    schema,
    messages: [
      { role: 'system', content: imageSearchPrompt },
      ...imageSearchFewShots,
      {
        role: 'user',
        content: `<conversation>\n${formatChatHistoryAsString(input.chatHistory)}\n</conversation>\n<follow_up>\n${input.query}\n</follow_up>`,
      },
    ],
  });

  const { results } = await search(res.query, {
    engines: ['bing images', 'google images'],
  });

  const images: ImageResult[] = [];
  for (const r of results) {
    if (r.img_src && r.url && r.title) {
      images.push({ img_src: r.img_src, url: r.url, title: r.title });
    }
  }
  return images.slice(0, 10);
}
