# buddy_math_server/exercise_generator.py
import json
import logging
import time
import google.generativeai as genai
from pydantic import BaseModel, Field
from typing import Optional
from config import GEMINI_MODEL
from firebase_manager import firebase_manager

logger = logging.getLogger(__name__)

class GeneratedExerciseSchema(BaseModel):
    new_problem_latex: str = Field(description="The new exercise in LaTeX format")
    hint_for_student: str = Field(description="A helpful hint for the student")
    correct_answer_internal: str = Field(description="The correct answer for system verification")

class ExerciseGenerator:
    def __init__(self):
        # We reuse the same model as orchestrator
        self.model = genai.GenerativeModel(GEMINI_MODEL)

    async def generate_challenge(self, original_problem: str, category: str, uid: str) -> Optional[dict]:
        """
        V318.0: Generates a similar challenge exercise based on a solved one.
        """
        prompt = f"""אתה מחולל תרגילים למתמטיקה. 
המטרה: לייצר תרגיל 'אתגר' דומה לתרגיל שהתלמיד פתר בהצלחה.

התרגיל המקורי:
[ORIGINAL_PROBLEM]
{original_problem}

הקטגוריה:
{category}

הנחיות:
1. ייצר תרגיל חדש באותה רמת קושי בדיוק, המבוסס על אותו עיקרון מתמטי.
2. השתמש במספרים שונים.
3. וודא שהפתרון יוצא מספר 'נח' (שלם או שבר פשוט) אם אפשר.
4. החזר JSON בלבד לפי הסכימה המוגדרת.
"""
        
        try:
            generation_config = genai.GenerationConfig(
                response_mime_type="application/json",
                response_schema=GeneratedExerciseSchema,
                temperature=0.8
            )
            
            res = await self.model.generate_content_async(
                prompt,
                generation_config=generation_config
            )
            
            data = json.loads(res.text)
            
            # Save to Firestore
            self._save_to_suggestions(uid, data, category)
            
            return data
            
        except Exception as e:
            logger.error(f"❌ [EXERCISE-GENERATOR] Error: {e}")
            return None

    def _save_to_suggestions(self, uid: str, data: dict, category: str):
        """
        Saves the generated exercise to users/{uid}/suggestions.
        """
        try:
            db = firebase_manager.get_db()
            if not db: return
            
            suggestion_ref = db.collection('users').document(uid).collection('suggestions').document()
            suggestion_ref.set({
                "problem_latex": data.get("new_problem_latex"),
                "hint": data.get("hint_for_student"),
                "answer": data.get("correct_answer_internal"),
                "category": category,
                "status": "pending",
                "reminder_sent": False,
                "created_at": time.time()
            })
            logger.info(f"✅ [EXERCISE-GENERATOR] Suggestion saved for user {uid}")
        except Exception as e:
            logger.error(f"❌ [EXERCISE-GENERATOR] Firestore Save Error: {e}")

exercise_generator = ExerciseGenerator()