from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from typing import List, Optional
from llama_index.llms.ollama import Ollama
from llama_index.core import PromptTemplate
from llama_index.llms.groq import Groq
import os
import dotenv
from enum import Enum

dotenv.load_dotenv()

app = FastAPI()

# Enable CORS for all origins (for development purposes)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load LLM model
# llm = Ollama(model="llama3", request_timeout=120.0)
# llama3-70b-8192
llm = Groq(model="llama-3.3-70b-versatile", 
           api_key="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",
    )



class Options(BaseModel):
    """A quiz with a title and a list of questions."""
    questions: List[str] = Field(..., description="List of questions in the quiz")

class Quiz(BaseModel):
    """An album inspired by a movie, containing an artist and a list of songs."""
    question: str = Field(..., description="The question to be answered")
    correct_answer: str = Field(..., description="The correct answer to the question")
    options: List[Options] = Field(..., description="List of options for the question")

class Difficulty(str, Enum):
    easy = "Easy"
    medium = "Medium"
    hard = "Hard"

class QuizFormat(str, Enum):
    multiple_choice = "Multiple Choice"
    open_ended = "Open ended"
    true_false = "True/False"

class QuizRequest(BaseModel):
    topic: str = Field(..., description="Topic of the quiz")
    subject: str = Field(..., description="Subject of the quiz")
    grade_level: str = Field(..., description="Grade level (e.g., Primary 3, JSS 1)")
    description: Optional[str] = Field(None, description="Brief description of the quiz (optional)")
    difficulty: Difficulty = Field(..., description="Difficulty level")
    format: QuizFormat = Field(..., description="Format of the quiz")
    num_questions: int = Field(..., gt=0, le=50, description="Number of questions (1-50)")
    language: str = Field(..., description="Language to translate the quiz into")

# Define the prompt template for quiz generation
PROMPT_TEMPLATE = """
Generate a set of quizzes that align with the given requirements:
Topic: {topic}
Subject: {subject}
Grade Level: {grade_level}
{description_part}
Difficulty: {difficulty}
Format: {format}
Number of Questions: {num_questions}
Ensure the format follows and ensure its in json format. No extra information just your response please and make it well formatted: 
[
    {{
        "question": "Question text here?",
        "correct_answer": "Correct answer",
        "options": ["Option 1", "Option 2", "Option 3", "Option 4"]
    }},
]
"""

FORMAT_PROMPT_TEMPLATE = """
You are a formatting assistant. Your ONLY job is to format quizzes into a structured, teacher-friendly format. 

STRICT INSTRUCTIONS:
1. Do NOT include any explanations, introductions, or thought processes. Output ONLY the formatted quiz.
2. Do NOT add or remove content from questions, answers, or options.
3. Format each quiz as follows:
   - Start with the question number (e.g., "1. Question text here?").
   - On the next line, write "Answer: " followed by the correct answer.
   - On the next line, write "Options: " followed by the list of options, each on a new line with a bullet point.
4. Separate each quiz with two blank lines for clarity.
5. Use STRICTLY these HTML tags for proper formatting:
   - `<br>` for line breaks.
   - `<p>` for paragraphs.
   - `<ul>` and `<li>` for lists.

Output ONLY the formatted quiz. No extra text.

Here is the input JSON:
{quiz_json}
"""

# Define the prompt template for translation
TRANSLATION_PROMPT_TEMPLATE = """
Translate the following text into {target_language}. Ensure the translation is accurate and maintains the original meaning:
{text}
"""

def format_quiz_stream(prompt: str):
    """Stream responses from the LLM for formatting"""
    response = llm.complete(prompt)
    if hasattr(response, "text"):
        for chunk in response.text.split("\n"):
            yield chunk
    elif hasattr(response, "content"):
        for chunk in response.content.split("\n"):
            yield chunk
    else:
        raise AttributeError("Response does not contain text content.")

def translate_text(text: str, target_language: str) -> str:
    """Translate text to the desired language using the LLM"""
    translation_prompt = TRANSLATION_PROMPT_TEMPLATE.format(
        target_language=target_language,
        text=text
    )
    translation_response = llm.complete(translation_prompt)
    return translation_response.text


@app.get("/")
def read_root():
    return {"Hello": "World"}

@app.post("/generate-quizzes/")
def generate_quiz(quiz_request: QuizRequest):
    print(quiz_request)
    """Generates quizzes dynamically based on user input"""
    try:
        # Handle optional description
        description_part = f"Description: {quiz_request.description}" if quiz_request.description else ""

        # Step 1: Generate the raw quiz JSON using the first LLM
        prompt = PROMPT_TEMPLATE.format(
            topic=quiz_request.topic,
            subject=quiz_request.subject,
            grade_level=quiz_request.grade_level,
            format=quiz_request.format,
            description_part=description_part,  # This avoids adding "Description: None" if it's not provided
            difficulty=quiz_request.difficulty,
            num_questions=quiz_request.num_questions
        )

        # Generate the raw quiz JSON
        raw_quiz_response = llm.complete(prompt).text

        # Step 2: Translate the raw quiz JSON into the desired language
        translated_quiz_response = translate_text(raw_quiz_response, quiz_request.language)

        # Step 3: Format the translated quizzes into a teacher-friendly format
        format_prompt = FORMAT_PROMPT_TEMPLATE.format(
            quiz_json=translated_quiz_response,
            
            )

        # Return the formatted quizzes as a streaming response
        return StreamingResponse(format_quiz_stream(format_prompt), media_type="text/plain")
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")