"""
LLM Chain implementation using Langchain for educational concept analysis
"""

from typing import Dict, Any, List
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.output_parsers import PydanticOutputParser
from langchain.chains import LLMChain
from pydantic import BaseModel, Field
from config import OPENAI_API_KEY, OPENAI_MODEL

# Define Pydantic models for structured output
class Concept(BaseModel):
    """Model for a single concept"""
    id: str = Field(description="Unique identifier for the concept")
    name: str = Field(description="Name of the concept")
    description: str = Field(description="Brief description of the concept")
    difficulty: str = Field(description="Difficulty level: basic, intermediate, or advanced")

class Relationship(BaseModel):
    """Model for relationship between concepts"""
    source: str = Field(description="Source concept ID")
    target: str = Field(description="Target concept ID")
    type: str = Field(description="Type of relationship: prerequisite or related")
    explanation: str = Field(description="Explanation of why this relationship exists")

class ConceptMap(BaseModel):
    """Model for complete concept map"""
    main_concept: str = Field(description="Main concept being analyzed")
    sub_concepts: List[Concept] = Field(description="List of sub-concepts")
    relationships: List[Relationship] = Field(description="List of relationships between concepts")

class Example(BaseModel):
    """Model for concept examples"""
    problem: str = Field(description="Example problem")
    solution: str = Field(description="Step-by-step solution")
    difficulty: str = Field(description="Difficulty level: Easy, Medium, or Hard")

class Resource(BaseModel):
    """Model for learning resources"""
    type: str = Field(description="Type of resource (Video/Article/Interactive/Book)")
    title: str = Field(description="Resource title")
    description: str = Field(description="Resource description")
    link: str = Field(description="Optional resource link")

class ConceptExplanation(BaseModel):
    """Model for detailed concept explanation"""
    explanation: str = Field(description="Detailed concept explanation")
    examples: List[Example] = Field(description="List of example problems and solutions")
    resources: List[Resource] = Field(description="List of learning resources")
    practice_questions: List[Example] = Field(description="List of practice questions")

class EducationalLLMChain:
    """
    Chain for processing educational concepts using LLM
    """
    def __init__(self):
        """Initialize the LLM and parsers"""
        self.llm = ChatOpenAI(
            model=OPENAI_MODEL,
            temperature=0.1,
            openai_api_key=OPENAI_API_KEY
        )
        
        # Initialize output parsers
        self.concept_parser = PydanticOutputParser(pydantic_object=ConceptMap)
        self.explanation_parser = PydanticOutputParser(pydantic_object=ConceptExplanation)
        
        # Create decomposition chain
        self.decomposition_chain = self._create_decomposition_chain()
        
        # Create explanation chain
        self.explanation_chain = self._create_explanation_chain()

    def _create_decomposition_chain(self) -> LLMChain:
        """
        Create chain for concept decomposition
        
        Returns:
            LLMChain for decomposing concepts
        """
        template = """You are an expert educational AI tutor.
        
        Analyze this question for a {grade} level student studying {subject}.
        
        Question: {question}
        
        Student Background:
        - Grade Level: {grade}
        - Subject: {subject}
        - Learning Needs: {learning_needs}
        
        Break down the concepts needed to understand this question into a knowledge graph.
        Consider the student's grade level and background knowledge.
        
        {format_instructions}
        """
        
        prompt = ChatPromptTemplate.from_template(
            template=template,
            partial_variables={
                "format_instructions": self.concept_parser.get_format_instructions()
            }
        )
        
        return LLMChain(llm=self.llm, prompt=prompt)

    def _create_explanation_chain(self) -> LLMChain:
        """
        Create chain for concept explanation
        
        Returns:
            LLMChain for explaining concepts
        """
        template = """You are an expert educational tutor.
        
        Explain this concept for a {grade} level student studying {subject}:
        
        Concept: {concept_name}
        Description: {concept_description}
        
        Student Background:
        - Grade Level: {grade}
        - Subject: {subject}
        - Learning Needs: {learning_needs}
        
        Provide a detailed explanation, examples, resources, and practice questions.
        
        {format_instructions}
        """
        
        prompt = ChatPromptTemplate.from_template(
            template=template,
            partial_variables={
                "format_instructions": self.explanation_parser.get_format_instructions()
            }
        )
        
        return LLMChain(llm=self.llm, prompt=prompt)

    async def decompose_concepts(
        self,
        question: str,
        grade: str,
        subject: str,
        learning_needs: str
    ) -> ConceptMap:
        """
        Decompose a question into concepts
        
        Args:
            question: User's question
            grade: Educational grade level
            subject: Subject area
            learning_needs: Learning needs/goals
            
        Returns:
            Structured concept map
        """
        response = await self.decomposition_chain.arun({
            "question": question,
            "grade": grade,
            "subject": subject,
            "learning_needs": learning_needs
        })
        
        return self.concept_parser.parse(response)

    async def explain_concept(
        self,
        concept_name: str,
        concept_description: str,
        grade: str,
        subject: str,
        learning_needs: str
    ) -> ConceptExplanation:
        """
        Generate detailed concept explanation
        
        Args:
            concept_name: Name of concept to explain
            concept_description: Brief concept description
            grade: Educational grade level
            subject: Subject area
            learning_needs: Learning needs/goals
            
        Returns:
            Structured concept explanation
        """
        response = await self.explanation_chain.arun({
            "concept_name": concept_name,
            "concept_description": concept_description,
            "grade": grade,
            "subject": subject,
            "learning_needs": learning_needs
        })
        
        return self.explanation_parser.parse(response)