# agents/teacher.py
import os
import json
import time
from dotenv import load_dotenv
import openai
from .base_agent import BaseAgent

load_dotenv()


class TeacherAgent(BaseAgent):
    def __init__(self, level):
        super().__init__(f"{level.capitalize()}TeacherAgent")
        api_key = os.getenv("OPENAI_API_KEY")
        self.client = openai.OpenAI(api_key=api_key)
        self.level = level.lower()

    def teach_module(self, module):
        """
        Returns a rich, k-shot based explanation string for the given module.
        Will attempt a small re-prompt if the output is too short.
        """
        system_prompt = (
            "You are an engaging and adaptive AI Tutor. Produce a thorough, structured explanation "
            "of the requested module tuned to the user's learning level. Use the structure below."
            "\n\nTeaching style differences by level:\n"
            "- Novice: Simple language, relatable analogies, many concrete examples.\n"
            "- Intermediate: Explain 'how' and 'why', connect to related concepts, include practical examples.\n"
            "- Advanced: Discuss nuances, edge-cases, performance, comparative methods and deeper insights.\n\n"
            "Explanation structure (MANDATORY):\n"
            "1) Core Concept — clear, precise description (several sentences)\n"
            "2) Worked Example or Analogy — at least one detailed example that illustrates application\n"
            "4) Key Takeaway — one concise sentence\n\n"
            "### Few-shot examples (follow style):\n\n"
            "Novice Example:\n"
            "Module: Variables\n"
            "Core Concept: Variables are containers that store values like numbers or text. Example: x = 5.\n"
            "Analogy/Example: Think of a labeled jar.\n"
            "Intermediate Example:\n"
            "Module: For loops\n"
            "Core Concept: For loops iterate over a collection. Example: for item in collection: process(item)\n"
            "Application: Useful for batch-processing and iteration in algorithms.\n\n"
            "Advanced Example:\n"
            "Module: Tail recursion\n"
            "Core Concept: Tail recursion preserves state for compiler optimizations; consider stack usage.\n\n"
            "Produce detailed output (minimum ~130 words). Do NOT output JSON — just plain text explanation."
        )

        user_prompt = (
            f"Module: {module.name}\n"
            f"Learning objective: {getattr(module, 'learning_objective', '')}\n\n"
            "Please produce the explanation now."
        )

        resp = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            # you can tune temperature if needed
        )
        explanation = resp.choices[0].message.content.strip()

        # If too short, ask for expansion once
        if len(explanation.split()) < 120:
            followup = (
                "The previous explanation was too short. Please expand the explanation, add another worked example "
                "and a short code or pseudo-code snippet where appropriate. Keep same style & level."
            )
            resp2 = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                    {"role": "user", "content": followup}
                ],
            )
            extra = resp2.choices[0].message.content.strip()
            # concatenate but keep readable
            explanation = explanation + "\n\n" + extra

        # return explanation string (do not print in agent; app will show it)
        return explanation

    def evaluate_example(self, module, student_example):
        """
        Evaluate the student's example for the given module.
        Returns a dict: { "is_correct": bool, "feedback": str, "confidence": float (0-1) }
        The model MUST output JSON only; we robustly parse it and fallback if needed.
        """
        system_prompt = (
            "You are an expert educational assessor. Evaluate the student's example strictly with respect "
            "to the module's learning objective. Respond with JSON ONLY (no extra text). The JSON object MUST contain:\n"
            " - is_correct: true/false\n"
            " - feedback: short one-sentence constructive feedback\n"
            " - confidence: numeric between 0 and 1\n\n"
            f"Evaluation sensitivity is based on learner level: {self.level}."
        )

        user_prompt = (
            f"Module: {module.name}\n"
            f"Learning objective: {getattr(module, 'learning_objective', '')}\n\n"
            f"Student example: {student_example}\n\n"
            "Evaluate and return JSON only."
        )

        resp = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            # lower temperature can reduce hallucinations
        )
        raw = resp.choices[0].message.content.strip()

        # Try to parse JSON; be forgiving (extract first {...})
        try:
            parsed = json.loads(raw)
        except Exception:
            import re
            m = re.search(r'(\{.*\})', raw, re.DOTALL)
            if m:
                try:
                    parsed = json.loads(m.group(1))
                except Exception:
                    parsed = None
            else:
                parsed = None

        if not parsed:
            # fallback: simple heuristics (very conservative)
            is_correct = len(student_example.strip()) > 20
            feedback = "Could not parse evaluator output; using conservative heuristic. Provide a more concrete example." \
                if not is_correct else "Example seems plausible but automatic evaluation failed to parse."
            confidence = 0.45 if not is_correct else 0.6
            return {"is_correct": bool(is_correct), "feedback": feedback, "confidence": confidence, "raw": raw}

        # normalize fields
        is_correct = bool(parsed.get("is_correct", parsed.get("correct", False)))
        feedback = str(parsed.get("feedback", parsed.get("explanation", "")))
        confidence = float(parsed.get("confidence", parsed.get("score", 0))) if parsed.get("confidence") is not None else 0.9

        return {"is_correct": is_correct, "feedback": feedback, "confidence": confidence, "raw": raw}

    def check_example(self, module):
        """
        Interactive method to check if student understands the module by asking for an example.
        Returns True if the student provides a satisfactory example, False otherwise.
        """
        print(f"\n🎯 Let's test your understanding of '{module.name}'!")
        print("Please provide a specific example or application that demonstrates this concept.")
        
        learning_objective = getattr(module, 'learning_objective', '')
        if learning_objective:
            print(f"Learning Objective: {learning_objective}")
        
        print("\nYour example should be:")
        print("• Specific and concrete (not just a definition)")
        print("• Relevant to the module content")
        print("• Show your understanding of how to apply the concept")
        print("• At least 2-3 sentences with clear reasoning")
        
        while True:  # Allow multiple attempts
            student_example = input("\nYour example: ").strip()
            
            if not student_example:
                print("❌ Please provide an example to demonstrate your understanding.")
                continue
            
            # Check for minimum length and specificity
            if len(student_example.split()) < 10:
                print("❌ Your example is too brief. Please provide a more detailed example with specific details.")
                continue
            
            # Evaluate the example
            evaluation = self.evaluate_example(module, student_example)
            
            if evaluation['is_correct']:
                print("✅ Correct! Your example demonstrates good understanding.")
                return True
            else:
                print("❌ Incorrect. Your example doesn't demonstrate sufficient understanding.")
                
                # Provide a simple hint
                hint = self.get_simple_hint(module)
                if hint:
                    print(f"💡 Hint: {hint}")
                
                retry = input("\nWould you like to try again with a different example? (y/n): ").lower().strip()
                if retry == 'y':
                    continue
                else:
                    print("📚 Let's review the concept again and then try the example.")
                    return False

    def get_simple_hint(self, module):
        """Generate a simple hint to guide the student toward a better example."""
        system_prompt = (
            "You are an expert tutor providing a simple hint. Give a brief, helpful suggestion (1-2 sentences) "
            "to guide the student toward providing a better example. Don't give away the answer, just nudge them in the right direction."
        )
        
        user_prompt = (
            f"Module: {module.name}\n"
            f"Learning Objective: {getattr(module, 'learning_objective', '')}\n\n"
            "Provide a simple hint to help the student think of a better example."
        )
        
        try:
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ]
            )
            return response.choices[0].message.content.strip()
        except Exception:
            return "Think about how this concept applies in real-world situations or practical scenarios."


# convenience subclasses (optional)
class NoviceTeacherAgent(TeacherAgent):
    def __init__(self):
        super().__init__("novice")


class IntermediateTeacherAgent(TeacherAgent):
    def __init__(self):
        super().__init__("intermediate")


class AdvancedTeacherAgent(TeacherAgent):
    def __init__(self):
        super().__init__("advanced")