import { GoogleGenAI, Type } from '@google/genai';
import { TranslationSentenceInput, TranslationSentenceResult } from './subtitlePipeline';

export interface TranslateSentencesWithGeminiOptions {
  targetLanguage: string;
  sentences: TranslationSentenceInput[];
  generateContent: (request: any) => Promise<{ text?: string | null }>;
}

const stripJsonFences = (text: string) => text.replace(/```json\n?|\n?```/g, '').trim();

export const translateSentencesWithGemini = async ({
  targetLanguage,
  sentences,
  generateContent,
}: TranslateSentencesWithGeminiOptions): Promise<TranslationSentenceResult[]> => {
  const prompt = `Translate the following subtitle segments into ${targetLanguage}.
CRITICAL INSTRUCTIONS:
1. You MUST return a valid JSON array of objects.
2. Keep the EXACT same "id" and "originalText" for each segment.
3. Add a "translatedText" field with the translation.
4. Add a "speaker" field if you can infer a useful display label.
5. DO NOT wrap in markdown blocks. Return ONLY JSON.

Input JSON:
${JSON.stringify(sentences, null, 2)}`;

  const response = await generateContent({
    model: 'gemini-2.5-flash',
    contents: [{ role: 'user', parts: [{ text: prompt }] }],
    config: {
      responseMimeType: 'application/json',
      responseSchema: {
        type: Type.ARRAY,
        items: {
          type: Type.OBJECT,
          properties: {
            id: { type: Type.STRING },
            translatedText: { type: Type.STRING },
            speaker: { type: Type.STRING },
          },
          required: ['id', 'translatedText'],
        },
      },
    },
  });

  const text = stripJsonFences(response.text || '[]');
  return JSON.parse(text) as TranslationSentenceResult[];
};

export const createGeminiSentenceTranslator = ({
  apiKey,
}: {
  apiKey: string;
}) => {
  const ai = new GoogleGenAI({ apiKey });

  return async (
    sentences: TranslationSentenceInput[],
    targetLanguage: string,
  ): Promise<TranslationSentenceResult[]> =>
    translateSentencesWithGemini({
      targetLanguage,
      sentences,
      generateContent: async (request) => {
        const response = await ai.models.generateContent(request);
        return { text: response.text };
      },
    });
};