import { openai } from '../agent.js';
import type { InterpreterOutput, RetrievalOutput, MediaType } from '../types/agents.js';
import { z } from 'zod';
import { zodTextFormat } from 'openai/helpers/zod';

const RetrievalSchema = z.object({
  candidates: z.array(z.object({
    title: z.string(),
    reason: z.string()
  }))
});

export async function runRetrieval(
  input: InterpreterOutput,
  brainstormCount = 100,
  mediaType: MediaType = 'tv_show',
  useWebSearch = false,
): Promise<RetrievalOutput> {
  const mediaLabel = mediaType === 'movie' ? 'movie' : 'TV show';
  const mediaLabelPlural = mediaType === 'movie' ? 'movies' : 'TV shows';

  const response = await openai.responses.parse({
    model: 'gpt-5.4',
    temperature: 0.9,
    service_tier: 'flex',
    ...(useWebSearch ? { tools: [{ type: 'web_search' as const }] } : {}),
    text: { format: zodTextFormat(RetrievalSchema, "candidates") },
    instructions: `You are a ${mediaLabel} candidate generator. Your goal is to brainstorm a LARGE, DIVERSE pool of ${brainstormCount} ${mediaLabel} candidates that match the user's structured preferences.${useWebSearch ? '\n\nUse web search to find recent and accurate titles, including newer releases.' : ''}

Rules:
- Include both well-known and obscure ${mediaLabelPlural}
- Prioritize RECALL over precision — it's better to include too many than too few
- Each "reason" should briefly explain why the ${mediaLabel} matches the preferences
- Avoid duplicates
- Include ${mediaLabelPlural} from different decades, countries${mediaType === 'tv_show' ? ', and networks' : ', and directors'}
- Aim for ${brainstormCount} candidates minimum`,
    input: `Structured preferences:
Liked ${mediaLabelPlural}: ${JSON.stringify(input.liked)}
Disliked ${mediaLabelPlural}: ${JSON.stringify(input.disliked)}
Themes: ${JSON.stringify(input.themes)}
Character preferences: ${JSON.stringify(input.character_preferences)}
Tone: ${JSON.stringify(input.tone)}
Avoid: ${JSON.stringify(input.avoid)}

Generate a large, diverse pool of ${mediaLabel} candidates.`,
  });

  return (response.output_parsed as RetrievalOutput) ?? { candidates: [] };
}
