import chainlit as cl
from pilates_evaluator import PilatesVideoEvaluator
from rag_processor import RAGProcessor
import os
from dotenv import load_dotenv

load_dotenv()

# Initialize RAG processor
rag_processor = RAGProcessor()

# Sample Pilates knowledge base - you can replace this with your own documents
PILATES_KNOWLEDGE = [
    "The Hundred is a fundamental Pilates exercise that helps build core strength and breathing control.",
    "Proper alignment is crucial in Pilates - maintain neutral spine position during exercises.",
    "Pilates focuses on six principles: concentration, control, centering, flow, precision, and breathing.",
    "The Pilates reformer is a versatile piece of equipment that provides resistance through springs.",
    "Pilates exercises should be performed with controlled movements and proper breathing patterns.",
    "The Pilates mat work series includes exercises like the Roll Up, Single Leg Stretch, and Double Leg Stretch.",
    "Pilates helps improve posture, flexibility, and overall body awareness.",
    "Joseph Pilates developed the method as a system of exercises to strengthen the mind and body.",
    "Pilates exercises can be modified for different fitness levels and physical conditions.",
    "Regular Pilates practice can help prevent injuries and improve athletic performance."
]

# Add documents to RAG system
rag_processor.add_documents(PILATES_KNOWLEDGE)

@cl.on_chat_start
async def start():
    await cl.Message(content="Welcome! You can chat with me about Pilates or upload a video to analyze your Pilates exercise.").send()

@cl.on_message
async def main(message: cl.Message):
    # Handle video uploads
    if message.elements:
        video_file = message.elements[0]
        if not video_file.name.endswith(('.mp4', '.avi', '.mov')):
            await cl.Message(content="Please upload a valid video file (mp4, avi, mov).").send()
            return

        await cl.Message(content=f"Analyzing video: {video_file.name}...").send()

        evaluator = PilatesVideoEvaluator()
        try:
            evaluator.process_video(video_file.path)
            report_path = "pilates_evaluation_report.json"
            evaluator.generate_report(report_path)
            await cl.Message(content=f"Analysis complete! Report saved to {report_path}.").send()
        except Exception as e:
            await cl.Message(content=f"Error analyzing video: {e}").send()
        return

    # Handle chat messages
    if message.content:
        try:
            # Create a message placeholder
            msg = cl.Message(content="")
            await msg.send()

            # Stream the response
            async for chunk in rag_processor.generate_response(message.content):
                await msg.stream_token(chunk)
        except Exception as e:
            await cl.Message(content=f"Error processing your message: {e}").send()
        return

    await cl.Message(content="Please either send a message or upload a video file.").send()