# src/db/mongodb_store.py
from motor.motor_asyncio import AsyncIOMotorClient
from datetime import datetime
import json
from typing import List, Dict, Optional
from bson import ObjectId

class MongoDBStore:
    def __init__(self, mongo_uri: str = "mongodb://localhost:27017"):
        """Initialize MongoDB connection"""
        self.client = AsyncIOMotorClient(mongo_uri)
        self.db = self.client.rag_chatbot
        self.chat_history = self.db.chat_history

    async def store_message(
        self,
        conversation_id: str,
        query: str,
        response: str,
        context: List[str],
        sources: List[Dict],
        llm_provider: str
    ) -> str:
        """Store chat message in MongoDB"""
        document = {
            "conversation_id": conversation_id,
            "timestamp": datetime.now(),
            "query": query,
            "response": response,
            "context": context,
            "sources": sources,
            "llm_provider": llm_provider,
            "feedback": None,
            "rating": None
        }
        
        result = await self.chat_history.insert_one(document)
        return str(result.inserted_id)

    async def get_conversation_history(self, conversation_id: str) -> List[Dict]:
        """Retrieve conversation history"""
        cursor = self.chat_history.find(
            {"conversation_id": conversation_id}
        ).sort("timestamp", 1)
        
        history = []
        async for document in cursor:
            document["_id"] = str(document["_id"])
            history.append(document)
            
        return history

    async def update_feedback(
        self,
        conversation_id: str,
        feedback: Optional[str],
        rating: Optional[int]
    ) -> bool:
        """Update feedback for a conversation"""
        result = await self.chat_history.update_many(
            {"conversation_id": conversation_id},
            {
                "$set": {
                    "feedback": feedback,
                    "rating": rating
                }
            }
        )
        return result.modified_count > 0

    async def get_messages_for_summary(self, conversation_id: str) -> List[Dict]:
        """Get messages in format suitable for summarization"""
        cursor = self.chat_history.find(
            {"conversation_id": conversation_id}
        ).sort("timestamp", 1)
        
        messages = []
        async for doc in cursor:
            messages.append({
                'role': 'user' if doc['query'] else 'assistant',
                'content': doc['query'] or doc['response'],
                'timestamp': doc['timestamp'],
                'sources': doc['sources']
            })
            
        return messages