import dspy
from typing import List, Dict, Any
import os

SMART_MODEL = "google/gemini-2.5-flash"

# Initialize LM with error handling for missing API key
api_key = os.getenv("OPENROUTER_API_KEY")
if not api_key:
    raise ValueError(
        "OPENROUTER_API_KEY environment variable is required but not set. "
        "Please ensure the API key is configured in the environment."
    )

lm = dspy.LM(
    model=f"openrouter/{SMART_MODEL}",
    api_base="https://openrouter.ai/api/v1",
    api_key=api_key,
    cache=True,
    temperature=0.1,
)
dspy.settings.configure(lm=lm)

class ProblemAnalysis(dspy.Signature):
    """Analyze a user's problem to identify key themes and emotions for finding relevant Bhagavad Gita teachings."""
    
    user_problem: str = dspy.InputField(desc="The user's problem or situation they need guidance on")
    
    key_themes: str = dspy.OutputField(desc="Main themes present in the problem (comma-separated keywords like 'duty, purpose, fear, decision-making')")
    emotional_state: str = dspy.OutputField(desc="The emotional state of the user (e.g., 'anxious, confused, overwhelmed, seeking purpose')")
    core_question: str = dspy.OutputField(desc="The fundamental question the user is asking, restated clearly")

class ContextualizeTeaching(dspy.Signature):
    """Contextualize Bhagavad Gita verses to address a specific user problem with compassionate guidance."""
    
    user_problem: str = dspy.InputField(desc="The user's specific problem or situation")
    verse_chapter: int = dspy.InputField(desc="Chapter number of the Bhagavad Gita verse")
    verse_number: int = dspy.InputField(desc="Verse number within the chapter")
    sanskrit_text: str = dspy.InputField(desc="Original Sanskrit verse")
    english_translation: str = dspy.InputField(desc="English translation of the verse")
    user_emotional_state: str = dspy.InputField(desc="The user's emotional state and key concerns")
    
    contextual_guidance: str = dspy.OutputField(desc="Compassionate guidance that connects the verse's wisdom to the user's specific situation, written in a warm, understanding tone as if Lord Krishna is speaking directly to the user")
    practical_application: str = dspy.OutputField(desc="Practical steps or mindset shifts the user can apply based on this teaching")
    reflection_question: str = dspy.OutputField(desc="A thoughtful question to help the user reflect on how this wisdom applies to their life")

class SynthesizeWisdom(dspy.Signature):
    """Synthesize multiple Bhagavad Gita teachings into a cohesive response addressing the user's problem."""
    
    user_problem: str = dspy.InputField(desc="The original problem the user brought")
    contextual_teachings: str = dspy.InputField(desc="Multiple contextualized teachings from different verses")
    core_message: str = dspy.InputField(desc="The main spiritual message to convey")
    
    final_response: str = dspy.OutputField(desc="A cohesive, compassionate response that weaves together the teachings into practical wisdom for the user's situation")
    closing_blessing: str = dspy.OutputField(desc="A brief, warm closing blessing or encouragement in the spirit of the Bhagavad Gita")

class ProblemAnalyzer(dspy.Module):
    def __init__(self):
        super().__init__()
        self.analyze = dspy.Predict(ProblemAnalysis)
    
    def forward(self, user_problem: str):
        return self.analyze(user_problem=user_problem)

class TeachingContextualizer(dspy.Module):
    def __init__(self):
        super().__init__()
        self.contextualize = dspy.Predict(ContextualizeTeaching)
    
    def forward(self, user_problem: str, verse_data: Dict[str, Any], user_emotional_state: str):
        return self.contextualize(
            user_problem=user_problem,
            verse_chapter=verse_data['chapter_num'],
            verse_number=verse_data['verse_num'],
            sanskrit_text=verse_data['sanskrit_text'],
            english_translation=verse_data['english_text'],
            user_emotional_state=user_emotional_state
        )

class WisdomSynthesizer(dspy.Module):
    def __init__(self):
        super().__init__()
        self.synthesize = dspy.Predict(SynthesizeWisdom)
    
    def forward(self, user_problem: str, contextual_teachings: str, core_message: str):
        return self.synthesize(
            user_problem=user_problem,
            contextual_teachings=contextual_teachings,
            core_message=core_message
        )