src/apis/routes/lesson_route.py

from fastapi import (
    APIRouter,
    status,
    Depends,
    BackgroundTasks,
    HTTPException,
    File,
    UploadFile,
    Form,
)
from fastapi.responses import JSONResponse
from src.utils.logger import logger
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional
from src.agents.lesson_practice.flow import lesson_practice_agent
from src.agents.lesson_practice_2.flow import lesson_practice_2_agent
from src.apis.models.lesson_models import Lesson, LessonResponse, LessonDetailResponse
import json
import os
import uuid
from datetime import datetime
import base64

router = APIRouter(prefix="/lesson", tags=["AI"])


class LessonPracticeRequest(BaseModel):
    unit: str = Field(..., description="Unit of the lesson")
    vocabulary: list = Field(..., description="Vocabulary for the lesson")
    key_structures: list = Field(..., description="Key structures for the lesson")
    practice_questions: list = Field(
        ..., description="Practice questions for the lesson"
    )
    student_level: str = Field("beginner", description="Student's level of English")
    query: str = Field(..., description="User query for the lesson")
    session_id: str = Field(..., description="Session ID for the lesson")


# Helper function to load lessons from JSON file
def load_lessons_from_file() -> List[Lesson]:
    """Load lessons from the JSON file"""
    try:
        lessons_file_path = os.path.join(os.path.dirname(__file__), "..", "..", "data", "lessons.json")
        
        if not os.path.exists(lessons_file_path):
            logger.warning(f"Lessons file not found at {lessons_file_path}")
            return []
            
        with open(lessons_file_path, 'r', encoding='utf-8') as file:
            lessons_data = json.load(file)
            
        # Convert to Lesson objects
        lessons = []
        for lesson_data in lessons_data:
            try:
                lesson = Lesson(**lesson_data)
                lessons.append(lesson)
            except Exception as e:
                logger.error(f"Error parsing lesson {lesson_data.get('id', 'unknown')}: {str(e)}")
                continue
                
        return lessons
    except Exception as e:
        logger.error(f"Error loading lessons: {str(e)}")
        return []


@router.get("/all", response_model=LessonResponse)
async def get_all_lessons():
    """
    Get all available lessons
    
    Returns:
        LessonResponse: Contains list of all lessons and total count
    """
    try:
        lessons = load_lessons_from_file()
        
        return LessonResponse(
            lessons=lessons,
            total=len(lessons)
        )
    except Exception as e:
        logger.error(f"Error retrieving lessons: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to retrieve lessons"
        )


@router.get("/{lesson_id}", response_model=LessonDetailResponse)
async def get_lesson_by_id(lesson_id: str):
    """
    Get a specific lesson by ID
    
    Args:
        lesson_id (str): The unique identifier of the lesson
        
    Returns:
        LessonDetailResponse: Contains the lesson details
    """
    try:
        lessons = load_lessons_from_file()
        
        # Find the lesson with the specified ID
        lesson = next((l for l in lessons if l.id == lesson_id), None)
        
        if not lesson:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"Lesson with ID '{lesson_id}' not found"
            )
            
        return LessonDetailResponse(lesson=lesson)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error retrieving lesson {lesson_id}: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to retrieve lesson"
        )


@router.get("/search/unit/{unit_name}")
async def search_lessons_by_unit(unit_name: str):
    """
    Search lessons by unit name (case-insensitive partial match)
    
    Args:
        unit_name (str): Part of the unit name to search for
        
    Returns:
        LessonResponse: Contains list of matching lessons
    """
    try:
        lessons = load_lessons_from_file()
        
        # Filter lessons by unit name (case-insensitive partial match)
        matching_lessons = [
            lesson for lesson in lessons 
            if unit_name.lower() in lesson.unit.lower()
        ]
        
        return LessonResponse(
            lessons=matching_lessons,
            total=len(matching_lessons)
        )
    except Exception as e:
        logger.error(f"Error searching lessons by unit '{unit_name}': {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to search lessons"
        )


@router.post("/chat")
async def chat(
    session_id: str = Form(
        ..., description="Session ID for tracking user interactions"
    ),
    lesson_data: str = Form(
        ..., description="The lesson data as JSON string"
    ),
    text_message: Optional[str] = Form(None, description="Text message from user"),
    audio_file: Optional[UploadFile] = File(None, description="Audio file from user"),
):
    """Send a message (text or audio) to the lesson practice agent"""

    # Validate that at least one input is provided
    if not text_message and not audio_file:
        raise HTTPException(
            status_code=400, detail="Either text_message or audio_file must be provided"
        )

    # Parse lesson data from JSON string
    try:
        lesson_dict = json.loads(lesson_data)
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid lesson_data JSON format")

    if not lesson_dict:
        raise HTTPException(status_code=400, detail="Lesson data not provided")

    # Prepare message content
    message_content = []

    # Handle text input
    if text_message:
        message_content.append({"type": "text", "text": text_message})

    # Handle audio input
    if audio_file:
        try:
            # Read audio file content
            audio_data = await audio_file.read()

            # Convert to base64
            audio_base64 = base64.b64encode(audio_data).decode("utf-8")

            # Determine mime type based on file extension
            file_extension = (
                audio_file.filename.split(".")[-1].lower()
                if audio_file.filename
                else "wav"
            )
            mime_type_map = {
                "wav": "audio/wav",
                "mp3": "audio/mpeg",
                "ogg": "audio/ogg",
                "webm": "audio/webm",
                "m4a": "audio/mp4",
            }
            mime_type = mime_type_map.get(file_extension, "audio/wav")

            message_content.append(
                {
                    "type": "audio",
                    "source_type": "base64",
                    "data": audio_base64,
                    "mime_type": mime_type,
                }
            )

        except Exception as e:
            logger.error(f"Error processing audio file: {str(e)}")
            raise HTTPException(
                status_code=400, detail=f"Error processing audio file: {str(e)}"
            )

    # Create message in the required format
    message = {"role": "user", "content": message_content}

    try:
        response = await lesson_practice_agent().ainvoke(
            {
                "messages": [message],
                "unit": lesson_dict.get("unit", ""),
                "vocabulary": lesson_dict.get("vocabulary", []),
                "key_structures": lesson_dict.get("key_structures", []),
                "practice_questions": lesson_dict.get("practice_questions", []),
                "student_level": lesson_dict.get("student_level", "beginner"),
            },
            {"configurable": {"thread_id": session_id}},
        )
        
        # Extract AI response content
        ai_response = response["messages"][-1].content
        logger.info(f"AI response: {ai_response}")

        return JSONResponse(content={"response": ai_response})

    except Exception as e:
        logger.error(f"Error in lesson practice: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")


@router.post("/chat_v2")
async def chat_v2(
    session_id: str = Form(
        ..., description="Session ID for tracking user interactions"
    ),
    lesson_data: str = Form(
        ..., description="The lesson data as JSON string"
    ),
    text_message: Optional[str] = Form(None, description="Text message from user"),
    audio_file: Optional[UploadFile] = File(None, description="Audio file from user"),
):
    """Send a message (text or audio) to the lesson practice v2 agent with Practice and Teaching agents"""

    # Validate that at least one input is provided
    if not text_message and not audio_file:
        raise HTTPException(
            status_code=400, detail="Either text_message or audio_file must be provided"
        )

    # Parse lesson data from JSON string
    try:
        lesson_dict = json.loads(lesson_data)
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid lesson_data JSON format")

    if not lesson_dict:
        raise HTTPException(status_code=400, detail="Lesson data not provided")

    # Prepare message content
    message_content = []

    # Handle text input
    if text_message:
        message_content.append({"type": "text", "text": text_message})

    # Handle audio input
    if audio_file:
        try:
            # Read audio file content
            audio_data = await audio_file.read()

            # Convert to base64
            audio_base64 = base64.b64encode(audio_data).decode("utf-8")

            # Determine mime type based on file extension
            file_extension = (
                audio_file.filename.split(".")[-1].lower()
                if audio_file.filename
                else "wav"
            )
            mime_type_map = {
                "wav": "audio/wav",
                "mp3": "audio/mpeg",
                "ogg": "audio/ogg",
                "webm": "audio/webm",
                "m4a": "audio/mp4",
            }
            mime_type = mime_type_map.get(file_extension, "audio/wav")

            message_content.append(
                {
                    "type": "audio",
                    "source_type": "base64",
                    "data": audio_base64,
                    "mime_type": mime_type,
                }
            )

        except Exception as e:
            logger.error(f"Error processing audio file: {str(e)}")
            raise HTTPException(
                status_code=400, detail=f"Error processing audio file: {str(e)}"
            )

    # Create message in the required format
    message = {"role": "user", "content": message_content}

    try:
        response = await lesson_practice_2_agent().ainvoke(
            {
                "messages": [message],
                "unit": lesson_dict.get("unit", ""),
                "vocabulary": lesson_dict.get("vocabulary", []),
                "key_structures": lesson_dict.get("key_structures", []),
                "practice_questions": lesson_dict.get("practice_questions", []),
                "student_level": lesson_dict.get("student_level", "beginner"),
            },
            {"configurable": {"thread_id": session_id}},
        )
        
        # Extract AI response content
        ai_response = response["messages"][-1].content
        logger.info(f"AI response (v2): {ai_response}")

        return JSONResponse(content={"response": ai_response})

    except Exception as e:
        logger.error(f"Error in lesson practice v2: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")