main.py

import os
from dotenv import load_dotenv
from agno.agent import Agent
from agno.models.google import Gemini
from agno.team.team import Team
from textwrap import dedent
import re

# Load API Key
load_dotenv()
model = Gemini(api_key=os.getenv("GEMINI_API_KEY"))

# Agent 1: Document Topic Extractor
document_parser = Agent(
    name="DocumentParser",
    model=model,
    instructions=dedent("""
        You are a document analyzer. Extract the key topics and subtopics from the given study material.
        These topics should guide question paper generation.
        Output should be clean, in bullet points or numbered format.
    """)
)

# Agent 2: Question Paper Generator
question_generator = Agent(
    name="QuestionGenerator",
    model=model,
    instructions=dedent("""
        You are a smart AI that generates customized question papers based on study material and structure requirements.

        Input will specify:
        - Total marks
        - Number of questions per type (MCQ, 1m, 2m, 3m, ..., 20m)
        - Difficulty level (Easy, Medium, Hard)
        - Structure: First section with questions, second with answer key

        Ensure:
        - Each question includes marks, Bloom's Taxonomy level, and difficulty
        - MCQs have 4 options
        - Answer Key section starts after all questions
        - All questions and answers are aligned with the input material
    """)
)

# Combine into a Team
team = Team(
    name="QuestionPaperTeam",
    model=model,
    members=[document_parser, question_generator]
)


def build_question_prompt(total_marks: int, counts: dict, difficulty: str) -> str:
    """
    Build the custom prompt based on user selection of marks, question count, and difficulty level.
    """
    parts = [
        f"Generate a question paper of total {total_marks} marks with the following format:",
        f"The overall difficulty level should be **{difficulty}**."
    ]

    mapping = {
        "mcq": "MCQs (1 mark each, 4 options)",
        "q1": "1-mark questions",
        "q2": "2-mark questions",
        "q3": "3-mark questions",
        "q4": "4-mark questions",
        "q5": "5-mark questions",
        "q10": "10-mark questions",
        "q15": "15-mark questions",
        "q20": "20-mark questions"
    }

    for key, label in mapping.items():
        count = counts.get(key, 0)
        if count > 0:
            parts.append(f"- {count} {label}")

    parts.append("First show only the questions in a clean, numbered format with difficulty and Bloom’s level.")
    parts.append("Then, on a new section titled 'Answer Key', provide all answers in the same order.")
    parts.append("In the Answer Key, DO NOT repeat the question — only write numbered answers corresponding to each question.")

    return "\n".join(parts)


def run_question_pipeline(raw_text: str, total_marks: int, counts: dict, difficulty: str) -> dict:
    try:
        # Step 1: Extract Topics
        topics_response = document_parser.run(message=f"Analyze the following study material:\n{raw_text}")
        topics = topics_response.final_output if hasattr(topics_response, "final_output") else topics_response.content

        # Step 2: Build Prompt with difficulty
        prompt = build_question_prompt(total_marks, counts, difficulty)
        full_input = f"{prompt}\n\nUse this content:\n{raw_text}"

        # Step 3: Generate Questions + Answer Key
        questions_response = question_generator.run(message=full_input)
        questions = questions_response.final_output if hasattr(questions_response, "final_output") else questions_response.content

        return {
            "topics": topics.strip(),
            "questions": questions.strip()
        }

    except Exception as e:
        return {"error": str(e)}