"""
LangGraph Japanese Lesson Generator - Direct Topic to Detailed Lesson
Takes a topic name and generates comprehensive 3-level interactive lessons
"""

import os
import sys
from dotenv import load_dotenv
from typing import Optional, List, TypedDict

load_dotenv()

if not os.getenv("GROQ_API_KEY"):
    print("⚠️  WARNING: GROQ_API_KEY not set in environment")

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

try:
    from langchain_groq import ChatGroq
    from langgraph.graph import StateGraph, END
except ImportError as e:
    print(f"❌ Import Error: {e}")
    sys.exit(1)


# ============ PYDANTIC MODELS ============
class VocabularyItem(BaseModel):
    japanese: str = Field(..., description="Japanese word")
    romaji: str = Field(..., description="Romanized version")
    meaning: str = Field(..., description="English meaning")
    example: str = Field(..., description="Example sentence")
    hiragana: Optional[str] = Field(None, description="Hiragana reading")


class GrammarPoint(BaseModel):
    pattern: str = Field(..., description="Grammar pattern")
    level: str = Field(..., description="JLPT Level: N5, N4, N3")
    explanation: str = Field(..., description="Detailed explanation")
    example: str = Field(..., description="Example sentence")
    translation: str = Field(..., description="English translation")


class InteractiveActivity(BaseModel):
    type: str = Field(..., description="Activity type: roleplay, translation, formation")
    instruction: str = Field(..., description="Activity instruction")
    example: str = Field(..., description="Example with solution")


class Question(BaseModel):
    question: str = Field(..., description="Question text")
    options: List[str] = Field(..., min_items=4, max_items=4)
    correctIndex: int = Field(..., ge=0, le=3)
    explanation: str = Field(..., description="Why this answer is correct")


class SubSection(BaseModel):
    title: str = Field(..., description="Subsection title")
    level: str = Field(..., description="BEGINNER, INTERMEDIATE, ADVANCED")
    conversation: str = Field(..., description="Realistic dialogue")
    conversationTranslation: str = Field(..., description="English translation of dialogue")
    vocabularies: List[VocabularyItem] = Field(..., min_items=10)
    grammarPoints: List[GrammarPoint] = Field(..., min_items=4)
    interactiveActivities: List[InteractiveActivity] = Field(..., min_items=2)
    questions: List[Question] = Field(..., min_items=5, max_items=5)
    culturalNote: Optional[str] = Field(None, description="Cultural context or nuance")


class Lesson(BaseModel):
    title: str = Field(..., description="Lesson title")
    description: str = Field(..., description="Lesson overview")
    objectives: List[str] = Field(..., min_items=3, description="Learning objectives")
    color: str = Field(..., pattern="^#[0-9a-fA-F]{6}$", description="Hex color")
    difficulty: str = Field(..., description="Overall difficulty")
    estimatedTime: str = Field(..., description="Time to complete")
    subSections: List[SubSection] = Field(..., min_items=3, max_items=3)


class LessonRequest(BaseModel):
    lesson_topic: str = Field(..., description="Japanese lesson topic")
    additional_context: Optional[str] = Field(None, description="Optional additional context")


# ============ STATE ============
class State(TypedDict):
    topic: str
    final_lesson: Optional[Lesson]
    error: Optional[str]


# ============ INITIALIZE LLM ============
print("🚀 Initializing Lesson Generator...")

try:
    llm = ChatGroq(
        model="llama-3.3-70b-versatile",
        temperature=0.8,
        api_key=os.getenv("GROQ_API_KEY")
    )
    print("✅ Groq LLM initialized")
except Exception as e:
    print(f"❌ Error initializing Groq: {e}")
    llm = None

try:
    generator_llm = llm.with_structured_output(Lesson)
    print("✅ Structured output configured")
except Exception as e:
    print(f"❌ Error configuring structured output: {e}")


# ============ AGENTS ============
def generator_agent(state: State) -> State:
    """Generator agent - creates comprehensive lesson directly from topic"""
    
    topic = state["topic"]
    
    generate_prompt = f"""You are an expert Japanese language instructor. Create a detailed, comprehensive lesson about: {topic}

IMPORTANT: You MUST return ONLY valid JSON. NO markdown, NO backticks, NO explanations.

For each vocabulary item, MUST include:
- Actual Japanese characters (e.g., こんにちは, 銀行, 友達)
- Romaji transcription
- English meaning
- Hiragana reading
- Example sentence with Japanese characters

Generate exactly this JSON structure (MUST be valid JSON):

{{
  "title": "Create an engaging title related to {topic}",
  "description": "Write 2-3 sentences explaining what students will learn about {topic}",
  "difficulty": "Choose: BEGINNER, INTERMEDIATE, or ADVANCED based on topic complexity",
  "estimatedTime": "Choose: 30-45 minutes, 45-60 minutes, or 60-90 minutes",
  "objectives": [
    "First learning objective related to {topic}",
    "Second learning objective related to {topic}",
    "Third learning objective related to {topic}"
  ],
  "color": "#3498db",
  "subSections": [
    {{
      "title": "Beginner Foundation for {topic}",
      "level": "BEGINNER",
      "conversation": "Create a realistic beginner-level dialogue related to {topic} with A and B speakers using simple Japanese. Include line breaks with \\n",
      "conversationTranslation": "Provide English translation of the dialogue",
      "vocabularies": [
        {{"japanese": "こんにちは", "romaji": "Konnichiwa", "meaning": "Hello", "hiragana": "こんにちは", "example": "こんにちは。お元気ですか?"}},
        {{"japanese": "友達", "romaji": "Tomodachi", "meaning": "Friend", "hiragana": "ともだち", "example": "私の友達は学生です。"}},
        {{"japanese": "話す", "romaji": "Hanasu", "meaning": "To talk/speak", "hiragana": "はなす", "example": "友達と話します。"}},
        {{"japanese": "一緒に", "romaji": "Issho ni", "meaning": "Together", "hiragana": "いっしょに", "example": "友達と一緒に遊びます。"}},
        {{"japanese": "楽しい", "romaji": "Tanoshii", "meaning": "Fun/enjoyable", "hiragana": "たのしい", "example": "友達との会話は楽しいです。"}},
        {{"japanese": "笑う", "romaji": "Warau", "meaning": "To laugh", "hiragana": "わらう", "example": "一緒に笑います。"}},
        {{"japanese": "学校", "romaji": "Gakkou", "meaning": "School", "hiragana": "がっこう", "example": "学校の友達です。"}},
        {{"japanese": "遊ぶ", "romaji": "Asobu", "meaning": "To play/hang out", "hiragana": "あそぶ", "example": "放課後に遊びます。"}},
        {{"japanese": "会う", "romaji": "Au", "meaning": "To meet", "hiragana": "あう", "example": "友達に会いました。"}},
        {{"japanese": "大切", "romaji": "Taisetsu", "meaning": "Important/precious", "hiragana": "たいせつ", "example": "友達は大切です。"}}
      ],
      "grammarPoints": [
        {{"pattern": "N は N です", "level": "N5", "explanation": "Basic statement pattern. Used to state what something or someone is.", "example": "私は学生です。", "translation": "I am a student."}},
        {{"pattern": "V ます (Polite present)", "level": "N5", "explanation": "Polite form of action verbs in present/future tense.", "example": "友達と話します。", "translation": "I talk with my friend."}},
        {{"pattern": "と一緒に", "level": "N5", "explanation": "Means 'together with' or 'at the same time as'. Shows doing something with someone.", "example": "友達と一緒に遊びます。", "translation": "I play together with my friend."}},
        {{"pattern": "N は 形容詞 です", "level": "N5", "explanation": "Using adjectives to describe nouns in polite form.", "example": "友達は楽しいです。", "translation": "My friend is fun."}}
      ],
      "interactiveActivities": [
        {{"type": "roleplay", "instruction": "Beginner-level roleplay instruction related to {topic}", "example": "Example dialogue or sentences"}},
        {{"type": "translation", "instruction": "Beginner-level translation instruction with 2-3 items to translate", "example": "Example translations"}}
      ],
      "questions": [
        {{"question": "Beginner question about {topic} or the dialogue", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 1, "explanation": "Explanation of correct answer"}},
        {{"question": "Beginner vocabulary question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 0, "explanation": "Explanation of correct answer"}},
        {{"question": "Beginner comprehension question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 2, "explanation": "Explanation of correct answer"}},
        {{"question": "Beginner grammar question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 3, "explanation": "Explanation of correct answer"}},
        {{"question": "Beginner practical application question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 1, "explanation": "Explanation of correct answer"}}
      ],
      "culturalNote": "Add a relevant cultural note about Japanese etiquette or customs related to {topic}"
    }},
    {{
      "title": "Intermediate Conversations for {topic}",
      "level": "INTERMEDIATE",
      "conversation": "Create a realistic intermediate-level dialogue related to {topic} with natural flow and varied grammar. Use \\n for line breaks",
      "conversationTranslation": "Provide English translation of the dialogue",
      "vocabularies": [
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}}
      ],
      "grammarPoints": [
        {{"pattern": "Grammar pattern 1", "level": "N4", "explanation": "Intermediate grammar explanation", "example": "Example sentence", "translation": "Translation"}},
        {{"pattern": "Grammar pattern 2", "level": "N4", "explanation": "Intermediate grammar explanation", "example": "Example sentence", "translation": "Translation"}},
        {{"pattern": "Grammar pattern 3", "level": "N4", "explanation": "Intermediate grammar explanation", "example": "Example sentence", "translation": "Translation"}},
        {{"pattern": "Grammar pattern 4", "level": "N4", "explanation": "Intermediate grammar explanation", "example": "Example sentence", "translation": "Translation"}}
      ],
      "interactiveActivities": [
        {{"type": "roleplay", "instruction": "Intermediate roleplay scenario related to {topic}", "example": "Example dialogue"}},
        {{"type": "formation", "instruction": "Intermediate sentence formation with specific vocabulary/grammar", "example": "Example sentences"}}
      ],
      "questions": [
        {{"question": "Intermediate question about {topic}", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 2, "explanation": "Explanation"}},
        {{"question": "Intermediate grammar question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 1, "explanation": "Explanation"}},
        {{"question": "Intermediate comprehension question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 0, "explanation": "Explanation"}},
        {{"question": "Intermediate practical question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 3, "explanation": "Explanation"}},
        {{"question": "Intermediate nuance question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 1, "explanation": "Explanation"}}
      ],
      "culturalNote": "Add intermediate-level cultural insight about {topic}"
    }},
    {{
      "title": "Advanced Nuanced Usage for {topic}",
      "level": "ADVANCED",
      "conversation": "Create an advanced-level dialogue with sophisticated expressions, keigo (敬語), and nuanced usage related to {topic}. Use \\n for line breaks",
      "conversationTranslation": "Provide English translation",
      "vocabularies": [
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}},
        {{"japanese": "Japanese word", "romaji": "romanization", "meaning": "English meaning", "hiragana": "hiragana", "example": "Example sentence"}}
      ],
      "grammarPoints": [
        {{"pattern": "Advanced grammar pattern 1", "level": "N3", "explanation": "Advanced grammar explanation", "example": "Example sentence", "translation": "Translation"}},
        {{"pattern": "Advanced grammar pattern 2", "level": "N3", "explanation": "Advanced grammar explanation", "example": "Example sentence", "translation": "Translation"}},
        {{"pattern": "Advanced grammar pattern 3", "level": "N3", "explanation": "Advanced grammar explanation", "example": "Example sentence", "translation": "Translation"}},
        {{"pattern": "Advanced grammar pattern 4", "level": "N3", "explanation": "Advanced grammar explanation", "example": "Example sentence", "translation": "Translation"}}
      ],
      "interactiveActivities": [
        {{"type": "roleplay", "instruction": "Advanced roleplay scenario with formal/business context related to {topic}", "example": "Example dialogue"}},
        {{"type": "translation", "instruction": "Advanced translation with complex sentences", "example": "Example translations"}}
      ],
      "questions": [
        {{"question": "Advanced question about {topic}", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 0, "explanation": "Explanation"}},
        {{"question": "Advanced nuance question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 3, "explanation": "Explanation"}},
        {{"question": "Advanced keigo question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 1, "explanation": "Explanation"}},
        {{"question": "Advanced comprehension question", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 2, "explanation": "Explanation"}},
        {{"question": "Advanced practical application", "options": ["Option 1", "Option 2", "Option 3", "Option 4"], "correctIndex": 1, "explanation": "Explanation"}}
      ],
      "culturalNote": "Add advanced cultural and social context about {topic}"
    }}
  ]
}}

CRITICAL REQUIREMENTS:
- Return ONLY valid JSON, no markdown or backticks
- Each subsection must have EXACTLY 10 vocabularies, 4 grammar points, 2 activities, 5 questions
- All Japanese text must use actual Japanese characters (hiragana, kanji, katakana)
- All romaji and meanings must be accurate
- Difficulty levels must progress: BEGINNER → INTERMEDIATE → ADVANCED
- Grammatical levels must be appropriate: N5 for beginner, N4 for intermediate, N3 for advanced
- All conversations must include \\n for line breaks between speakers
- All questions must have exactly 4 options with one correct answer
- Color must be valid hex code
- Make content realistic, practical, and culturally accurate for {topic}"""

    try:
        print(f"🔄 Generating lesson for: {topic}...")
        lesson_result = generator_llm.invoke(generate_prompt)
        state["final_lesson"] = lesson_result
        print(f"✅ Lesson generated: {lesson_result.title}")
        return state
    except Exception as e:
        state["error"] = f"Generator Error: {str(e)}"
        print(f"❌ Error: {e}")
        return state


# ============ GRAPH ============
def build_graph():
    graph = StateGraph(State)
    graph.add_node("generator", generator_agent)
    graph.add_edge("generator", END)
    graph.set_entry_point("generator")
    return graph.compile()


try:
    workflow = build_graph()
    print("✅ LangGraph workflow built")
except Exception as e:
    print(f"❌ Error building graph: {e}")
    workflow = None


# ============ FASTAPI ============
app = FastAPI(title="Japanese Lesson Generator")


@app.on_event("startup")
async def startup_event():
    print("\n" + "="*60)
    print("🎌 Japanese Lesson Generator - Direct Topic to Lesson")
    print("   Simply provide a topic, get a complete interactive lesson")
    print("="*60 + "\n")


@app.post("/api/generate-lesson")
async def generate_lesson(request: LessonRequest):
    """Generate a comprehensive Japanese lesson from a topic"""
    try:
        if not request.lesson_topic or request.lesson_topic.strip() == "":
            raise HTTPException(status_code=400, detail="lesson_topic cannot be empty")

        if not workflow:
            raise HTTPException(status_code=500, detail="Workflow not initialized")

        print(f"\n📚 Generating lesson: {request.lesson_topic}")

        initial_state: State = {
            "topic": request.lesson_topic,
            "final_lesson": None,
            "error": None
        }

        result = workflow.invoke(initial_state)

        if result.get("error"):
            print(f"⚠️  Error: {result['error']}")
            raise HTTPException(status_code=500, detail=result["error"])

        if not result.get("final_lesson"):
            raise HTTPException(status_code=500, detail="Failed to generate lesson")

        print("✅ Lesson generated successfully!")
        return result["final_lesson"].model_dump()

    except HTTPException:
        raise
    except Exception as e:
        print(f"❌ API Error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health_check():
    return {
        "status": "healthy",
        "message": "Japanese Lesson Generator is running",
        "version": "2.1"
    }


@app.get("/")
async def root():
    return {
        "name": "Japanese Lesson Generator",
        "version": "2.1",
        "description": "Generate complete Japanese lessons from just a topic name",
        "usage": {
            "endpoint": "POST /api/generate-lesson",
            "example": {"lesson_topic": "Asking for directions", "additional_context": "For tourists"},
            "response": "Complete lesson with 3 difficulty levels, vocabularies, grammar, activities, and questions"
        },
        "features": [
            "✅ 3-level progressive difficulty (Beginner → Intermediate → Advanced)",
            "✅ Realistic dialogues with translations",
            "✅ 10 vocabularies per section with examples",
            "✅ 4 grammar points per section with JLPT levels",
            "✅ 2 interactive activities per section",
            "✅ 5 comprehension questions with explanations",
            "✅ Cultural notes for context"
        ]
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")