# src/db/mongodb_store.py
from motor.motor_asyncio import AsyncIOMotorClient
from datetime import datetime
import json
from typing import List, Dict, Optional, Any
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
        self.documents = self.db.documents  # Collection for original documents

    async def store_document(
        self,
        document_id: str,
        filename: str,
        content: str,
        content_type: str,
        file_size: int
    ) -> str:
        """Store original document in MongoDB"""
        document = {
            "document_id": document_id,
            "filename": filename,
            "content": content,
            "content_type": content_type,
            "file_size": file_size,
            "upload_timestamp": datetime.now()
        }
        
        await self.documents.insert_one(document)
        return document_id

    async def get_document(self, document_id: str) -> Optional[Dict]:
        """Retrieve document by ID"""
        return await self.documents.find_one(
            {"document_id": document_id},
            {"_id": 0}  # Exclude MongoDB's _id
        )

    async def get_all_documents(self) -> List[Dict]:
        """Retrieve all documents"""
        cursor = self.documents.find({}, {"_id": 0})
        return await cursor.to_list(length=None)

    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