import { NextRequest, NextResponse } from 'next/server';
import { db } from '@/lib/db';
import { verifyAdminPermission } from '@/lib/auth';
import { checkRateLimit, getRateLimitIdentifier, rateLimitError, RATE_LIMITS } from '@/lib/rate-limit';
import ZAI from 'z-ai-web-dev-sdk';
import type { VisionMessage } from 'z-ai-web-dev-sdk';

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const MAX_FILE_SIZE = 20 * 1024 * 1024; // 20 MB
const ALLOWED_MIME_TYPES = ['application/pdf'];
const MAX_PAGES = 30; // Process at most 30 pages per request

// ---------------------------------------------------------------------------
// VLM prompt
// ---------------------------------------------------------------------------

const VLM_SYSTEM_PROMPT = `You are an expert MCQ extraction assistant. Your job is to read a scanned or digital page image and extract every multiple-choice question (MCQ) from it.

Rules:
1. Extract ALL MCQs from the image – do not skip any.
2. Support Bangla, English, Math, and Science questions.
3. If a question contains an image, diagram, graph, or table that is essential to answering it, set hasImage to true.
4. For math questions, preserve the mathematical notation as closely as possible in plain text.
5. If the correct answer is visible (e.g. answer key on the same page), include it. Otherwise, make your best inference and set correctAnswer accordingly.
6. Difficulty should be one of: "easy", "medium", "hard". Infer from question complexity.
7. If an explanation is available, include it. Otherwise, leave explanation as an empty string.

You MUST respond with ONLY a valid JSON object – no markdown fences, no extra text. The format is:

{
  "questions": [
    {
      "question": "The full question text",
      "optionA": "First option text",
      "optionB": "Second option text",
      "optionC": "Third option text",
      "optionD": "Fourth option text",
      "correctAnswer": "A",
      "explanation": "Explanation or empty string",
      "difficulty": "easy",
      "hasImage": false
    }
  ]
}

If no MCQs are found on the page, return:
{
  "questions": []
}`;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

interface ExtractedQuestion {
  question: string;
  optionA: string;
  optionB: string;
  optionC: string;
  optionD: string;
  correctAnswer: string;
  explanation: string;
  difficulty: string;
  hasImage: boolean;
}

interface VLMExtractionResult {
  questions: ExtractedQuestion[];
}

/**
 * Parse the VLM response text into structured questions.
 * Tries to be resilient to minor formatting issues.
 */
function parseVLMResponse(text: string): ExtractedQuestion[] {
  try {
    // Strip markdown code fences if present
    let cleaned = text.trim();
    if (cleaned.startsWith('```')) {
      cleaned = cleaned.replace(/^```(?:json)?\s*\n?/, '').replace(/\n?```\s*$/, '');
    }

    const parsed: VLMExtractionResult = JSON.parse(cleaned);

    if (!parsed.questions || !Array.isArray(parsed.questions)) {
      return [];
    }

    // Validate & normalise each question
    return parsed.questions
      .map((q: Record<string, unknown>) => {
        const question = typeof q.question === 'string' ? q.question.trim() : '';
        const optionA = typeof q.optionA === 'string' ? q.optionA.trim() : '';
        const optionB = typeof q.optionB === 'string' ? q.optionB.trim() : '';
        const optionC = typeof q.optionC === 'string' ? q.optionC.trim() : '';
        const optionD = typeof q.optionD === 'string' ? q.optionD.trim() : '';
        const rawAnswer = typeof q.correctAnswer === 'string' ? q.correctAnswer.trim().toUpperCase() : '';
        const explanation = typeof q.explanation === 'string' ? q.explanation.trim() : '';
        const rawDifficulty = typeof q.difficulty === 'string' ? q.difficulty.trim().toLowerCase() : 'medium';
        const hasImage = q.hasImage === true;

        // Skip entries without a question or at least 2 options
        if (!question) return null;

        const validAnswer = ['A', 'B', 'C', 'D'].includes(rawAnswer) ? rawAnswer : 'A';
        const validDifficulty = ['easy', 'medium', 'hard'].includes(rawDifficulty) ? rawDifficulty : 'medium';

        return {
          question,
          optionA: optionA || 'N/A',
          optionB: optionB || 'N/A',
          optionC: optionC || 'N/A',
          optionD: optionD || 'N/A',
          correctAnswer: validAnswer,
          explanation,
          difficulty: validDifficulty,
          hasImage,
        };
      })
      .filter(Boolean) as ExtractedQuestion[];
  } catch {
    console.error('[extract-pdf] Failed to parse VLM JSON response');
    return [];
  }
}

/**
 * Convert a single PDF page image (Buffer) to base64 string.
 */
function bufferToBase64(buf: Buffer): string {
  return buf.toString('base64');
}

// ---------------------------------------------------------------------------
// POST handler
// ---------------------------------------------------------------------------

export async function POST(request: NextRequest) {
  try {
    // 1. Auth check
    const authResult = await verifyAdminPermission(request, 'canManageMCQ');
    if ('error' in authResult) return authResult.error;

    // 2. Rate limit – stricter than default because VLM calls are expensive
    const identifier = `admin-extract-pdf:${getRateLimitIdentifier(request)}`;
    const rateLimit = checkRateLimit(identifier, { maxRequests: 5, windowMs: 60 * 1000 });
    if (!rateLimit.allowed) return rateLimitError(rateLimit.resetTime);

    // 3. Parse multipart form data
    const formData = await request.formData();
    const file = formData.get('file');

    if (!file || !(file instanceof File)) {
      return NextResponse.json(
        { success: false, error: 'PDF file is required' },
        { status: 400 }
      );
    }

    // Validate MIME type
    if (!ALLOWED_MIME_TYPES.includes(file.type) && !file.name.toLowerCase().endsWith('.pdf')) {
      return NextResponse.json(
        { success: false, error: 'Only PDF files are accepted' },
        { status: 400 }
      );
    }

    // Validate file size
    if (file.size > MAX_FILE_SIZE) {
      return NextResponse.json(
        { success: false, error: `File too large. Maximum size is ${MAX_FILE_SIZE / 1024 / 1024}MB` },
        { status: 400 }
      );
    }

    // Optional assignment fields
    const classId = (formData.get('classId') as string) || null;
    const subjectId = (formData.get('subjectId') as string) || null;
    const topicId = (formData.get('topicId') as string) || null;

    // Validate optional IDs against DB if provided
    if (classId) {
      const mcqClass = await db.mcqClass.findUnique({ where: { id: classId } });
      if (!mcqClass) {
        return NextResponse.json(
          { success: false, error: 'Invalid classId' },
          { status: 400 }
        );
      }
    }
    if (subjectId) {
      const subject = await db.mcqSubject.findUnique({ where: { id: subjectId } });
      if (!subject) {
        return NextResponse.json(
          { success: false, error: 'Invalid subjectId' },
          { status: 400 }
        );
      }
    }
    if (topicId) {
      const topic = await db.mcqTopic.findUnique({ where: { id: topicId } });
      if (!topic) {
        return NextResponse.json(
          { success: false, error: 'Invalid topicId' },
          { status: 400 }
        );
      }
    }

    // 4. Read the PDF into a Buffer
    const arrayBuffer = await file.arrayBuffer();
    const pdfBuffer = Buffer.from(arrayBuffer);

    // 5. Convert PDF pages to images
    console.log(`[extract-pdf] Converting PDF "${file.name}" (${(pdfBuffer.length / 1024).toFixed(1)}KB) to images...`);

    let pageImages: Buffer[];
    try {
      // pdf-img-convert returns an array of Buffer (one per page)
      const convert = (await import('pdf-img-convert')).convert;
      const imageArrays = await convert(pdfBuffer, {
        width: 1200,           // Reasonable resolution for VLM
        height: 1600,
        page_numbers: Array.from({ length: MAX_PAGES }, (_, i) => i + 1), // limit pages
      });

      pageImages = imageArrays.map((arr: Uint8Array | number[]) =>
        Buffer.from(arr)
      );
    } catch (conversionError) {
      console.error('[extract-pdf] PDF to image conversion failed:', conversionError);
      return NextResponse.json(
        {
          success: false,
          error: 'Failed to convert PDF to images. The PDF may be corrupted or password-protected.',
        },
        { status: 422 }
      );
    }

    if (pageImages.length === 0) {
      return NextResponse.json(
        { success: false, error: 'No pages could be extracted from the PDF' },
        { status: 422 }
      );
    }

    console.log(`[extract-pdf] Converted ${pageImages.length} page(s). Starting VLM extraction...`);

    // 6. Initialise VLM client
    const vlm = await ZAI.create();

    // 7. Process each page through VLM
    const allExtractedQuestions: ExtractedQuestion[] = [];
    const pagesProcessed = pageImages.length;
    const errors: string[] = [];

    for (let i = 0; i < pageImages.length; i++) {
      const pageImage = pageImages[i];
      const base64Image = bufferToBase64(pageImage);

      console.log(`[extract-pdf] Processing page ${i + 1}/${pageImages.length}...`);

      try {
        const messages: VisionMessage[] = [
          {
            role: 'user',
            content: [
              {
                type: 'image_url',
                image_url: { url: `data:image/png;base64,${base64Image}` },
              },
              {
                type: 'text',
                text: VLM_SYSTEM_PROMPT,
              },
            ],
          },
        ];

        const response = await vlm.chat.completions.createVision({
          model: 'glm-4v-flash',
          messages,
          stream: false,
        });

        // Extract the text content from the response
        let responseText = '';
        if (response && response.choices && response.choices.length > 0) {
          const choice = response.choices[0];
          responseText = choice.message?.content || choice.message?.text || '';
          // If content is an array (some models return array)
          if (Array.isArray(responseText)) {
            responseText = responseText
              .filter((c: Record<string, unknown>) => c.type === 'text')
              .map((c: Record<string, unknown>) => c.text)
              .join('\n');
          }
        }

        if (!responseText) {
          errors.push(`Page ${i + 1}: VLM returned empty response`);
          continue;
        }

        const pageQuestions = parseVLMResponse(responseText);
        console.log(`[extract-pdf] Page ${i + 1}: extracted ${pageQuestions.length} question(s)`);

        // Attach optional assignment metadata
        for (const q of pageQuestions) {
          allExtractedQuestions.push({
            ...q,
            // These are not part of ExtractedQuestion type but we'll include them
            // in the response so the admin frontend can auto-fill the form
          } as ExtractedQuestion & { classId?: string | null; subjectId?: string | null; topicId?: string | null });
        }
      } catch (vlmError) {
        console.error(`[extract-pdf] VLM error on page ${i + 1}:`, vlmError);
        errors.push(`Page ${i + 1}: VLM processing failed`);
        // Continue with other pages instead of failing completely
      }
    }

    // 8. Build the response with optional assignment metadata
    const extractedQuestions = allExtractedQuestions.map((q) => ({
      ...q,
      classId: classId || undefined,
      subjectId: subjectId || undefined,
      topicId: topicId || undefined,
    }));

    const totalExtracted = extractedQuestions.length;

    console.log(`[extract-pdf] Done. Extracted ${totalExtracted} MCQ(s) from ${pagesProcessed} page(s).`);

    return NextResponse.json({
      success: true,
      extractedQuestions,
      totalExtracted,
      pagesProcessed,
      errors: errors.length > 0 ? errors : undefined,
      message: `Extracted ${totalExtracted} MCQs from ${pagesProcessed} page(s)`,
    });
  } catch (error) {
    console.error('[extract-pdf] Unhandled error:', error);
    return NextResponse.json(
      { success: false, error: 'An unexpected error occurred during PDF processing' },
      { status: 500 }
    );
  }
}
