import { openai, defaultModel, serviceOptions, supportsWebSearch, parseWithRetry } 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,
  hardRequirements = false,
  previousFullMatches: string[] = [],
): Promise<RetrievalOutput> {
  const mediaLabel = mediaType === 'movie' ? 'movie' : 'TV show';
  const mediaLabelPlural = mediaType === 'movie' ? 'movies' : 'TV shows';

  const canSearch = useWebSearch && supportsWebSearch;
  const response = await parseWithRetry(() => openai.responses.parse({
    model: defaultModel,
    temperature: 0.9,
    max_completion_tokens: 16384,
    ...serviceOptions,
    ...(canSearch ? { 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${previousFullMatches.length > 0 ? '\n- Do NOT suggest titles already in the Previous Full Matches list — generate NEW candidates inspired by what made those successful' : ''}${hardRequirements ? '\n\nIMPORTANT: Strictly follow ALL requirements. Exclude any candidate that does not meet every stated requirement.' : ''}`,
    input: `Structured preferences:
Liked ${mediaLabelPlural}: ${input.liked.join(', ') || '(none)'}
Disliked ${mediaLabelPlural}: ${input.disliked.join(', ') || '(none)'}
Themes: ${input.themes.join(', ') || '(none)'}
Character preferences: ${input.character_preferences.join(', ') || '(none)'}
Tone: ${input.tone.join(', ') || '(none)'}
Avoid: ${input.avoid.join(', ') || '(none)'}
Requirements: ${input.requirements.join(', ') || '(none)'}${previousFullMatches.length > 0 ? `\n\nPrevious Full Match titles (DO NOT repeat these; use them as inspiration for NEW candidates with similar qualities): ${previousFullMatches.join(', ')}` : ''}

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

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