import z from 'zod';
import { search } from '../search.js';
import { videoSearchPrompt, videoSearchFewShots } from '../prompts/video.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 VideoResult = {
  img_src: string;
  url: string;
  title: string;
  iframe_src: string;
};

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

  const res = await llm.generateObject<{ query: string }>({
    schema,
    messages: [
      { role: 'system', content: videoSearchPrompt },
      ...videoSearchFewShots,
      {
        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: ['youtube'] });

  const videos: VideoResult[] = [];
  for (const r of results) {
    const thumb = r.thumbnail ?? r.img_src;
    if (thumb && r.url && r.title && r.iframe_src) {
      videos.push({
        img_src: thumb,
        url: r.url,
        title: r.title,
        iframe_src: r.iframe_src,
      });
    }
  }
  return videos.slice(0, 10);
}
