kaboot/services/geminiService.ts

import { GoogleGenAI, Type, createUserContent, createPartFromUri } from "@google/genai";
import { Quiz, Question, AnswerOption, GenerateQuizOptions, ProcessedDocument } from "../types";
import { v4 as uuidv4 } from 'uuid';

const getClient = () => {
  const apiKey = process.env.API_KEY;
  if (!apiKey) {
    throw new Error("API_KEY environment variable is missing");
  }
  return new GoogleGenAI({ apiKey });
};

const QUIZ_SCHEMA = {
  type: Type.OBJECT,
  properties: {
    title: { type: Type.STRING, description: "A catchy title for the quiz" },
    questions: {
      type: Type.ARRAY,
      items: {
        type: Type.OBJECT,
        properties: {
          text: { type: Type.STRING, description: "The question text" },
          options: {
            type: Type.ARRAY,
            items: {
              type: Type.OBJECT,
              properties: {
                text: { type: Type.STRING },
                isCorrect: { type: Type.BOOLEAN },
                reason: { type: Type.STRING, description: "Brief explanation of why this answer is correct or incorrect" }
              },
              required: ["text", "isCorrect", "reason"]
            },
          }
        },
        required: ["text", "options"]
      }
    }
  },
  required: ["title", "questions"]
};

function buildPrompt(options: GenerateQuizOptions, hasDocuments: boolean): string {
  const questionCount = options.questionCount || 10;
  
  const baseInstructions = `Create ${questionCount} engaging multiple-choice questions. Each question must have exactly 4 options, and exactly one correct answer. Vary the difficulty.

IMPORTANT: For each option's reason, write as if you are directly explaining facts - never reference "the document", "the text", "the material", or "the source". Write explanations as standalone factual statements.`;

  if (hasDocuments) {
    const topicContext = options.topic 
      ? ` Focus on aspects related to "${options.topic}".` 
      : '';
    return `Generate a quiz based on the provided content.${topicContext}

${baseInstructions}`;
  }
  
  return `Generate a trivia quiz about "${options.topic}".

${baseInstructions}`;
}

function shuffleArray<T>(array: T[]): T[] {
  const shuffled = [...array];
  for (let i = shuffled.length - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * (i + 1));
    [shuffled[i], shuffled[j]] = [shuffled[j], shuffled[i]];
  }
  return shuffled;
}

function transformToQuiz(data: any): Quiz {
  const shapes = ['triangle', 'diamond', 'circle', 'square'] as const;
  const colors = ['red', 'blue', 'yellow', 'green'] as const;
  
  const questions: Question[] = data.questions.map((q: any) => {
    const shuffledOpts = shuffleArray(q.options);
    
    const options: AnswerOption[] = shuffledOpts.map((opt: any, index: number) => ({
      text: opt.text,
      isCorrect: opt.isCorrect,
      shape: shapes[index % 4],
      color: colors[index % 4],
      reason: opt.reason
    }));

    return {
      id: uuidv4(),
      text: q.text,
      options: options,
      timeLimit: 20
    };
  });

  return {
    title: data.title,
    questions
  };
}

async function uploadNativeDocument(ai: GoogleGenAI, doc: ProcessedDocument): Promise<{ uri: string; mimeType: string }> {
  const buffer = typeof doc.content === 'string'
    ? Buffer.from(doc.content, 'base64')
    : doc.content;
  
  const blob = new Blob([buffer], { type: doc.mimeType });
  
  const uploadedFile = await ai.files.upload({
    file: blob,
    config: { mimeType: doc.mimeType! }
  });
  
  if (!uploadedFile.uri || !uploadedFile.mimeType) {
    throw new Error("Failed to upload document to Gemini");
  }
  
  return { uri: uploadedFile.uri, mimeType: uploadedFile.mimeType };
}

export const generateQuiz = async (options: GenerateQuizOptions): Promise<Quiz> => {
  const ai = getClient();
  
  const docs = options.documents || [];
  const hasDocuments = docs.length > 0;
  const prompt = buildPrompt(options, hasDocuments);
  
  let contents: any;
  
  if (hasDocuments) {
    const parts: any[] = [];
    
    for (const doc of docs) {
      if (doc.type === 'native' && doc.mimeType) {
        const uploaded = await uploadNativeDocument(ai, doc);
        parts.push(createPartFromUri(uploaded.uri, uploaded.mimeType));
      } else if (doc.type === 'text') {
        parts.push({ text: doc.content as string });
      }
    }
    
    parts.push({ text: prompt });
    contents = createUserContent(parts);
  } else {
    contents = prompt;
  }

  const response = await ai.models.generateContent({
    model: "gemini-3-flash-preview",
    contents,
    config: {
      responseMimeType: "application/json",
      responseSchema: QUIZ_SCHEMA
    }
  });

  if (!response.text) {
    throw new Error("Failed to generate quiz content");
  }

  const data = JSON.parse(response.text);
  return transformToQuiz(data);
};