database.py

import os
import json
import time
import asyncio
from logger_config import get_logger
from datetime import datetime
from dotenv import load_dotenv
from pymongo import AsyncMongoClient, UpdateOne
from pymongo.errors import DuplicateKeyError, ServerSelectionTimeoutError, ConnectionFailure

# Load environment variables
load_dotenv()

logger = get_logger(__name__)


class Database:
    def __init__(self, uri=None, db_name="aspira_db"):
        if uri is None:
            uri = os.getenv("MONGODB_URI", "mongodb://localhost:27017/")

        # Configure client to handle idle drops and timeouts properly
        self.client = AsyncMongoClient(
            uri,
            # Proactively close connections before load balancers (usually ~4 mins) drop them
            maxIdleTimeMS=210000,
            serverSelectionTimeoutMS=10000  # Wait 10s per attempt to find a server
        )
        self.db = self.client[db_name]

    async def initialize(self, retries=5, delay=5):
        """Tries to connect and initialize indexes, waiting if the Atlas cluster is paused."""
        for attempt in range(1, retries + 1):
            try:
                # Ping the database to force connection and check availability
                await self.client.admin.command('ping')
                await self._init_indexes()
                logger.info(
                    "Successfully connected to MongoDB and initialized indexes.")
                return
            except (ServerSelectionTimeoutError, ConnectionFailure) as e:
                logger.warning(f"MongoDB connection failed on attempt {attempt}/{retries}. Cluster might be waking up. Retrying in {delay}s...")
                if attempt == retries:
                    logger.error(
                        "Failed to connect to MongoDB after multiple retries. The backend might fail.")
                    raise e
                await asyncio.sleep(delay)

    async def _init_indexes(self):
        """Create indexes for performance and uniqueness"""
        # User uniqueness
        await self.db.users.create_index("username", unique=True)

        # Meta: user_id + key unique
        await self.db.meta.create_index([("user_id", 1), ("key", 1)], unique=True)

        # Keywords: user_id + conversation_id + keyword unique
        await self.db.keywords.create_index(
            [("user_id", 1), ("conversation_id", 1), ("keyword", 1)], unique=True)

        # QA Cache: user_id + question unique
        await self.db.qa_cache.create_index(
            [("user_id", 1), ("question", 1)], unique=True)

        # Conversation: index on user_id and conversation_id for fast retrieval
        await self.db.conversation.create_index(
            [("user_id", 1), ("conversation_id", 1), ("id", 1)])

        # Scores: user_id unique
        await self.db.scores.create_index("user_id", unique=True)

    # --- USER MANAGEMENT ---
    async def create_user(self, username, password_hash, groq_api_key_encrypted=""):
        try:
            result = await self.db.users.insert_one({
                "username": username,
                "password_hash": password_hash,
                "groq_api_key": groq_api_key_encrypted,
                "created_at": datetime.utcnow()
            })
            return str(result.inserted_id)
        except DuplicateKeyError:
            return None

    async def get_user(self, username):
        return await self.db.users.find_one({"username": username})

    async def get_user_by_id(self, user_id: str):
        from bson.objectid import ObjectId
        try:
            return await self.db.users.find_one({"_id": ObjectId(user_id)})
        except Exception:
            return None

    # --- SESSION_COUNTER ---
    async def get_session_counter(self, user_id: str) -> int:
        doc = await self.db.meta.find_one(
            {"user_id": user_id, "key": "session_counter"})
        return int(doc["value"]) if doc else 0

    async def increment_session_counter(self, user_id: str) -> int:
        # Atomic increment
        result = await self.db.meta.find_one_and_update(
            {"user_id": user_id, "key": "session_counter"},
            {"$inc": {"value": 1}},
            upsert=True,
            return_document=True
        )
        return result["value"]

    # --- RESUME ---
    async def save_resume(self, user_id: str, text: str):
        await self.db.meta.update_one(
            {"user_id": user_id, "key": "resume"},
            {"$set": {"value": text, "updated_at": datetime.utcnow()}},
            upsert=True
        )

    async def get_resume(self, user_id: str) -> str:
        doc = await self.db.meta.find_one({"user_id": user_id, "key": "resume"})
        return doc["value"] if doc else None

    # --- KEYWORDS (KW) ---
    async def get_keywords(self, user_id: str, conversation_id: str = "default") -> dict:
        cursor = self.db.keywords.find({"user_id": user_id})
        kw = {}
        async for doc in cursor:
            # Default to 1 if it's missing
            version = doc.get("schema_version", 1)

            # If it's the old state, pretend it belongs to the "default" conversation
            if version == 1:
                if conversation_id == "default":
                    kw[doc["keyword"]] = doc["scores"]

            # If it's the new state, strictly match the conversation_id
            elif version == 2:
                if doc.get("conversation_id") == conversation_id:
                    kw[doc["keyword"]] = doc["scores"]
        return kw

    async def update_keywords(self, user_id: str, kw_dict: dict, conversation_id: str = "default"):
        """Replaces keywords for the user."""
        try:
            await self.db.keywords.delete_many(
                {"user_id": user_id, "conversation_id": conversation_id})
            if kw_dict:
                await self.db.keywords.insert_many([
                    {"user_id": user_id, "conversation_id": conversation_id, "keyword": k, "scores": v,
                        "updated_at": datetime.utcnow(), "schema_version": 2}
                    for k, v in kw_dict.items()
                ])
        except Exception as e:
            logger.error(f"Error updating keywords: {e}")

    # --- QA ---
    async def update_qa(self, user_id: str, qa_dict: dict):
        ops = []
        for q, a in qa_dict.items():
            ops.append(UpdateOne(
                {"user_id": user_id, "question": q},
                {"$set": {"answer_chunk": a, "created_at": datetime.utcnow()}},
                upsert=True
            ))
        if ops:
            try:
                await self.db.qa_cache.bulk_write(ops)
            except Exception as e:
                logger.error(f"Error updating QA: {e}")

    async def get_qa(self, user_id: str) -> dict:
        cursor = self.db.qa_cache.find({"user_id": user_id})
        return {doc["question"]: doc["answer_chunk"] async for doc in cursor}

    # --- CONVERSATION HISTORY ---
    async def add_conversation_message(self, user_id: str, message: str, conversation_id: str = "default"):
        role = "user" if message.startswith("User:") else "interviewer"
        await self.db.conversation.insert_one({
            "user_id": user_id,
            "conversation_id": conversation_id,
            "message": message,
            "role": role,
            "timestamp": datetime.utcnow(),
            "schema_version": 2
        })

    async def get_conversation_history(self, user_id: str, conversation_id: str = "default") -> list:
        cursor = self.db.conversation.find(
            {"user_id": user_id}).sort("timestamp", 1)
        history = []
        async for doc in cursor:
            version = doc.get("schema_version", 1)  # Default to 1 if missing

            # Legacy records (version 1) default to "default" conversation
            if version == 1 and conversation_id == "default":
                history.append(doc["message"])

            # New records (version 2) check the explicit conversation_id
            elif version == 2 and doc.get("conversation_id") == conversation_id:
                history.append(doc["message"])

        return history

    async def get_conversations(self, user_id: str) -> list:
        """Returns a list of unique conversation IDs for the user."""
        return await self.db.conversation.distinct("conversation_id", {"user_id": user_id})

    # --- SCORES ---
    async def get_score(self, user_id: str) -> dict:
        doc = await self.db.scores.find_one({"user_id": user_id})
        return doc.get("data", {}) if doc else {}

    async def update_score(self, user_id: str, score_data: dict):
        """Update user's score data (accumulated interview performance)."""
        try:
            await self.db.scores.update_one(
                {"user_id": user_id},
                {"$set": {"data": score_data, "updated_at": datetime.utcnow()}},
                upsert=True
            )
        except Exception as e:
            logger.error(f"Error updating score: {e}")

    # --- METADATA ---
    async def get_interview_metadata(self, user_id: str, conversation_id: str) -> dict:
        doc = await self.db.metadata.find_one({"user_id": user_id, "conversation_id": conversation_id})
        return doc.get("metadata", {}) if doc else {}

    async def save_interview_metadata(self, user_id: str, conversation_id: str, metadata: dict):
        try:
            await self.db.metadata.update_one(
                {"user_id": user_id, "conversation_id": conversation_id},
                {"$set": {"metadata": metadata, "updated_at": datetime.utcnow()}},
                upsert=True
            )
        except Exception as e:
            logger.error(f"Error updating interview metadata: {e}")

    # --- EVALUATION ---
    async def get_evaluation(self, user_id: str, conversation_id: str) -> dict:
        doc = await self.db.evaluations.find_one({"user_id": user_id, "conversation_id": conversation_id})
        return doc.get("evaluation", {}) if doc else {}

    async def save_evaluation(self, user_id: str, conversation_id: str, evaluation: dict):
        try:
            await self.db.evaluations.update_one(
                {"user_id": user_id, "conversation_id": conversation_id},
                {"$set": {"evaluation": evaluation, "updated_at": datetime.utcnow()}},
                upsert=True
            )
        except Exception as e:
            logger.error(f"Error updating evaluation: {e}")

    # --- KNOWLEDGE GRAPH ---
    async def get_knowledge_graph(self, user_id: str, conversation_id: str) -> dict:
        """Retrieve the knowledge graph for a specific conversation."""
        doc = await self.db.knowledge_graphs.find_one(
            {"user_id": user_id, "conversation_id": conversation_id}
        )
        return doc.get("graph_data", {}) if doc else {}

    async def save_knowledge_graph(self, user_id: str, conversation_id: str, graph_data: dict):
        """Upsert the knowledge graph for a specific conversation."""
        try:
            await self.db.knowledge_graphs.update_one(
                {"user_id": user_id, "conversation_id": conversation_id},
                {"$set": {"graph_data": graph_data, "updated_at": datetime.utcnow()}},
                upsert=True
            )
        except Exception as e:
            logger.error(f"Error saving knowledge graph: {e}")