import fitz  # PyMuPDF
import json
import logging
import asyncio
from typing import List, Dict, Any
import os
from openai import AsyncOpenAI

logger = logging.getLogger("pdf_parser")

client = AsyncOpenAI(
    base_url=os.getenv("LLM_BASE_URL", "http://host.docker.internal:11434/v1"),
    api_key=os.getenv("OPENAI_API_KEY", "ollama"),
)

async def parse_pdf_to_questions(pdf_bytes: bytes, exam_pattern: str) -> List[Dict[str, Any]]:
    """
    Extracts text from a PDF and uses the LLM to structure it into a JSON array of questions.
    """
    text = ""
    try:
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        for page in doc:
            text += page.get_text("text") + "\n\n"
    except Exception as e:
        logger.error(f"Failed to parse PDF with PyMuPDF: {e}")
        return []

    system_prompt = f"""
    You are an AI expert at extracting {exam_pattern} test questions from messy PDF text.
    The text contains complex Chemistry and Math formulas. 
    CRITICAL: You MUST format all formulas perfectly in LaTeX wrapping inline math in $ and block math in $$.
    
    Extract the questions and return them STRICTLY as a JSON array of objects. 
    If your API requires a JSON object root, return {{"questions": [...]}}.
    
    Schema per question:
    {{
      "subject": "Physics" | "Chemistry" | "Maths" | "Biology",
      "topic": "String",
      "subtopic": "String",
      "difficulty": "Easy" | "Medium" | "Hard",
      "type": "MCQ_Single" | "MCQ_Multiple" | "Numerical",
      "question_text": "Full text of the question with LaTeX",
      "options": {{"A": "Option A text", "B": "Option B text", "C": "Option C text", "D": "Option D text"}},
      "correct_answer": {{"value": "A"}},
      "solution_text": "Detailed step-by-step solution if available, else 'Solution not provided'",
      "language": "English"
    }}
    """
    
    try:
        truncated_text = text[:15000] 
        
        response = await client.chat.completions.create(
            model=os.getenv("LLM_MODEL_NAME", "gemma"),
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": truncated_text}
            ],
            response_format={ "type": "json_object" }
        )
        
        output = response.choices[0].message.content
        if not output:
            return []
            
        if "```json" in output:
            output = output.split("```json")[1].split("```")[0].strip()
        elif "```" in output:
            output = output.split("```")[1].split("```")[0].strip()
            
        data = json.loads(output)
        
        if isinstance(data, list):
            return data
        elif isinstance(data, dict) and "questions" in data:
            return data["questions"]
        elif isinstance(data, dict):
            return [data]
        else:
            return []
    except Exception as e:
        logger.error(f"LLM PDF parsing failed: {e}")
        return []