"""Teaching Assistant Agent - Specialized agent for teaching programming."""

from __future__ import annotations

from typing import Any

from agenai.agent.specialized import SpecializedAgent
from agenai.agent.config import AgentConfig
from agenai.agent.context import AgentContext
from agenai.agent.core import AgentResult
from agenai.agent.tools.base import Tool
from agenai.agent.tools.read_file import ReadFileTool
from agenai.agent.tools.write_file import WriteFileTool

from agenai.agents.teaching_assistant.prompts import get_system_prompt
from agenai.agents.teaching_assistant.tools import (
    CodeRunnerTool,
    CodeReviewerTool,
    ExerciseGeneratorTool,
    QuizGeneratorTool,
    RAGSearchTool,
)


class TeachingAssistantAgent(SpecializedAgent):
    """Specialized agent for teaching Python programming to students.

    Features:
    - Explain programming concepts with examples and analogies
    - Review student code with constructive feedback
    - Generate practice exercises with test cases
    - Create quizzes to test understanding
    - Execute code to demonstrate concepts

    Example:
        agent = TeachingAssistantAgent()
        result = agent.run("Explain what a for loop is and give me an example")
        print(result.output)
    """

    name = "teaching_assistant"
    description = "An AI teaching assistant that teaches Python programming, reviews code, and creates exercises"

    def __init__(
        self,
        config: AgentConfig | None = None,
        difficulty_level: str = "adaptive",
        rag_retriever: Any | None = None,
        lesson_context: dict[str, Any] | None = None,
        **kwargs: Any,
    ):
        """Initialize the Teaching Assistant agent.

        Args:
            config: Agent configuration (model, max_steps, etc.)
            difficulty_level: Target difficulty ("beginner", "intermediate", "advanced", "adaptive")
            rag_retriever: Optional RAGRetriever for knowledge base search.
            lesson_context: Optional context about the current lesson being studied.
            **kwargs: Additional arguments passed to parent
        """
        self.difficulty_level = difficulty_level
        self._rag_retriever = rag_retriever
        self._lesson_context = lesson_context
        self._session_topics: list[str] = []
        self._questions_answered: int = 0
        self._exercises_created: int = 0
        self._code_reviewed: int = 0

        super().__init__(config=config, **kwargs)

    def get_system_prompt(self) -> str:
        """Return the teaching-focused system prompt."""
        base_prompt = get_system_prompt()

        # Add current configuration context
        config_note = f"""

## Current Session Configuration
- Difficulty Level: {self.difficulty_level}
- Focus: Python programming
"""
        # Add lesson context if available
        lesson_context_prompt = self._build_lesson_context_prompt()
        if lesson_context_prompt:
            config_note += lesson_context_prompt

        return base_prompt + config_note

    def _build_lesson_context_prompt(self) -> str:
        """Build lesson-specific context for system prompt."""
        if not self._lesson_context:
            return ""

        parts = ["\n## Current Lesson Context"]

        # Lesson info
        lesson = self._lesson_context.get("lesson", {})
        if lesson:
            parts.append(f"- **Lesson**: {lesson.get('title', 'Unknown')}")
            parts.append(f"- **Difficulty**: {lesson.get('difficulty', 'Unknown')}")
            if lesson.get("description"):
                parts.append(f"- **Description**: {lesson.get('description')}")

        # Current section
        section = self._lesson_context.get("current_section", {})
        if section:
            parts.append(f"\n### Current Section")
            parts.append(f"- **Section**: {section.get('title', 'Unknown')} (Type: {section.get('type', 'Unknown')})")
            if section.get("content"):
                # Truncate long content
                content = section.get("content", "")
                if len(content) > 500:
                    content = content[:500] + "..."
                parts.append(f"- **Content**: {content}")
            if section.get("instruction"):
                parts.append(f"- **Instruction**: {section.get('instruction')}")

        # Student profile
        student = self._lesson_context.get("student", {})
        if student:
            parts.append(f"\n### Student Profile")
            age_group = student.get("age_group", "adult")
            parts.append(f"- **Age Group**: {age_group}")
            interests = student.get("interests", [])
            if interests:
                parts.append(f"- **Interests**: {', '.join(interests)}")

        # Progress
        progress = self._lesson_context.get("progress", {})
        if progress:
            parts.append(f"\n### Progress")
            parts.append(f"- **Status**: {progress.get('status', 'not_started')}")
            if progress.get("score") is not None:
                parts.append(f"- **Quiz Score**: {progress.get('score')}%")

        # Teaching guidelines based on context
        parts.append(f"\n### Teaching Guidelines for This Lesson")
        parts.append("- Reference the current lesson and section content in your responses")
        parts.append("- Tailor explanations to the student's age group and interests")
        parts.append("- If the student asks about something in the lesson, provide specific help")
        parts.append("- Encourage the student based on their progress status")
        parts.append("- If the student asks about Python topics not covered in this lesson, briefly answer then gently guide them back (e.g., 'Great question! That's covered in a later lesson. For now, let's focus on...')")
        parts.append("- If the student asks completely off-topic questions (not Python-related), politely redirect them to the lesson (e.g., 'I'm here to help with Python! Let's get back to learning about...')")

        return "\n".join(parts)

    def get_tools(self) -> list[Tool]:
        """Return teaching tools."""
        tools: list[Tool] = [
            # Custom teaching tools
            CodeRunnerTool(),
            CodeReviewerTool(),
            ExerciseGeneratorTool(),
            QuizGeneratorTool(),
            # Built-in file tools
            ReadFileTool(),
            WriteFileTool(),
        ]

        # Add RAG search tool if configured
        if self._rag_retriever is not None:
            tools.insert(
                0,
                RAGSearchTool(
                    retriever=self._rag_retriever,
                    default_language="python",
                ),
            )

        return tools

    def setup(self) -> None:
        """Initialize session tracking."""
        self._session_topics = []
        self._questions_answered = 0
        self._exercises_created = 0
        self._code_reviewed = 0

    def on_start(self, context: AgentContext) -> None:
        """Log session start and set up context."""
        context.set("assistant_difficulty", self.difficulty_level)
        context.set("session_active", True)

    def on_complete(self, result: AgentResult) -> None:
        """Track session completion."""
        self._questions_answered += 1

        if result.success:
            # Track what kind of interaction this was
            for step in result.steps:
                if step.tool_name == "create_exercise":
                    self._exercises_created += 1
                elif step.tool_name == "review_code":
                    self._code_reviewed += 1

    def get_session_stats(self) -> dict[str, Any]:
        """Get statistics about the current teaching session."""
        return {
            "questions_answered": self._questions_answered,
            "exercises_created": self._exercises_created,
            "code_reviewed": self._code_reviewed,
            "topics_covered": self._session_topics,
        }
