kaboot/server/src/routes/generate.ts

import { Router, Response } from 'express';
import { GoogleGenAI, Type, createUserContent, createPartFromUri } from '@google/genai';
import { requireAuth, AuthenticatedRequest, requireAIAccess } from '../middleware/auth.js';
import { incrementGenerationCount, GENERATION_LIMIT, FREE_TIER_LIMIT } from '../services/stripe.js';
import { v4 as uuidv4 } from 'uuid';
import { buildQuizPrompt } from '../shared/quizPrompt.js';

const router = Router();

const GEMINI_API_KEY = process.env.GEMINI_API_KEY;
const DEFAULT_MODEL = 'gemini-3-flash-preview';
const MAX_CONCURRENT_GENERATIONS = 2;

type GenerationJob<T> = {
  priority: number;
  run: () => Promise<T>;
  resolve: (value: T) => void;
  reject: (error: unknown) => void;
};

const generationQueue: GenerationJob<any>[] = [];
let activeGenerations = 0;

const processGenerationQueue = () => {
  while (activeGenerations < MAX_CONCURRENT_GENERATIONS && generationQueue.length > 0) {
    const next = generationQueue.shift();
    if (!next) return;
    activeGenerations += 1;
    next
      .run()
      .then((result) => next.resolve(result))
      .catch((err) => next.reject(err))
      .finally(() => {
        activeGenerations -= 1;
        processGenerationQueue();
      });
  }
};

const enqueueGeneration = <T,>(priority: number, run: () => Promise<T>): Promise<T> => {
  return new Promise((resolve, reject) => {
    generationQueue.push({ priority, run, resolve, reject });
    generationQueue.sort((a, b) => b.priority - a.priority);
    processGenerationQueue();
  });
};

interface GenerateRequest {
  topic: string;
  questionCount?: number;
  documents?: Array<{
    type: 'text' | 'native';
    content: string;
    mimeType?: string;
  }>;
}

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"]
};

const buildPrompt = (topic: string, questionCount: number, hasDocuments: boolean): string =>
  buildQuizPrompt({ topic, questionCount, hasDocuments });

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) {
  const shapes = ['triangle', 'diamond', 'circle', 'square'] as const;
  const colors = ['red', 'blue', 'yellow', 'green'] as const;
  
  const questions = data.questions.map((q: any) => {
    const shuffledOpts = shuffleArray(q.options);
    
    const options = 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,
      timeLimit: 20
    };
  });

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

router.get('/status', (_req, res: Response) => {
  res.json({ 
    available: !!GEMINI_API_KEY,
    model: DEFAULT_MODEL
  });
});

router.post('/', requireAuth, requireAIAccess, async (req: AuthenticatedRequest, res: Response) => {
  if (!GEMINI_API_KEY) {
    res.status(503).json({ error: 'System AI is not configured' });
    return;
  }

  const { topic, questionCount = 10, documents = [] } = req.body as GenerateRequest;

  if (!topic && documents.length === 0) {
    res.status(400).json({ error: 'Topic or documents required' });
    return;
  }

  try {
    const accessInfo = (req as any).aiAccessInfo as { accessType?: 'group' | 'subscription' | 'none'; remaining?: number } | undefined;
    if (accessInfo?.accessType === 'none' && documents.length > 1) {
      res.status(403).json({ error: 'Free plan allows a single document per generation.' });
      return;
    }

    const priority = accessInfo?.accessType === 'subscription' || accessInfo?.accessType === 'group' ? 1 : 0;
    const queuedAt = Date.now();

    const quiz = await enqueueGeneration(priority, async () => {
      const ai = new GoogleGenAI({ apiKey: GEMINI_API_KEY });
      const hasDocuments = documents.length > 0;
      const prompt = buildPrompt(topic, questionCount, hasDocuments);

      let contents: any;

      if (hasDocuments) {
        const parts: any[] = [];

        for (const doc of documents) {
          if (doc.type === 'native' && doc.mimeType) {
            const buffer = Buffer.from(doc.content, 'base64');
            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) {
              parts.push(createPartFromUri(uploadedFile.uri, uploadedFile.mimeType));
            }
          } else if (doc.type === 'text') {
            parts.push({ text: doc.content });
          }
        }

        parts.push({ text: prompt });
        contents = createUserContent(parts);
      } else {
        contents = prompt;
      }

      const response = await ai.models.generateContent({
        model: DEFAULT_MODEL,
        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);
    });

    const waitMs = Date.now() - queuedAt;
    if (waitMs > 0) {
      console.log('AI generation queued', { waitMs, priority });
    }

    const groups = req.user!.groups || [];
    if (!groups.includes('kaboot-ai-access')) {
      const newCount = incrementGenerationCount(req.user!.sub);
      const limit = accessInfo?.accessType === 'subscription' ? GENERATION_LIMIT : FREE_TIER_LIMIT;
      const remaining = Math.max(0, limit - newCount);
      res.setHeader('X-Generations-Remaining', remaining.toString());
    }

    res.json(quiz);
  } catch (err: any) {
    console.error('AI generation error:', err);
    res.status(500).json({ error: err.message || 'Failed to generate quiz' });
  }
});

export default router;