src/database.py

"""MongoDB database client and utilities with connection pooling and Redis caching."""

import os
from pymongo import MongoClient, ASCENDING, TEXT
from pymongo.errors import ServerSelectionTimeoutError, BulkWriteError
import logging
from typing import Optional, Dict, Any, List
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

# MongoDB connection with connection pooling
_client: Optional[MongoClient] = None
_db = None

# Optional Redis client for session-based caching
_redis_client = None
try:
    import redis
    HAS_REDIS = True
except ImportError:
    HAS_REDIS = False
    logger.debug("Redis not installed - session caching disabled")


def get_mongo_client() -> MongoClient:
    """Get or create MongoDB client with connection pooling."""
    global _client
    if _client is None:
        mongo_uri = os.getenv("MONGO_URI", "mongodb://localhost:27017")
        max_pool_size = int(os.getenv("MONGO_MAX_POOL_SIZE", "10"))

        try:
            _client = MongoClient(
                mongo_uri,
                serverSelectionTimeoutMS=5000,
                maxPoolSize=max_pool_size,  # Connection pooling
                minPoolSize=2,  # Keep 2 connections ready
                maxIdleTimeMS=45000,  # Close idle connections after 45s
                socketTimeoutMS=20000,  # Socket timeout
                connectTimeoutMS=10000,  # Connection timeout
            )
            # Test connection
            _client.admin.command('ping')
            logger.info(f"Successfully connected to MongoDB (pool size: {max_pool_size})")
        except ServerSelectionTimeoutError:
            logger.error(f"Failed to connect to MongoDB at {mongo_uri}")
            raise
    return _client


def get_redis_client():
    """Get or create Redis client for session caching (optional)."""
    global _redis_client

    if not HAS_REDIS:
        return None

    if _redis_client is None:
        redis_host = os.getenv("REDIS_HOST", "localhost")
        redis_port = int(os.getenv("REDIS_PORT", "6379"))
        redis_db = int(os.getenv("REDIS_DB", "0"))
        redis_password = os.getenv("REDIS_PASSWORD")

        try:
            _redis_client = redis.Redis(
                host=redis_host,
                port=redis_port,
                db=redis_db,
                password=redis_password,
                decode_responses=True,
                socket_timeout=5,
                socket_connect_timeout=5,
            )
            # Test connection
            _redis_client.ping()
            logger.info(f"Successfully connected to Redis at {redis_host}:{redis_port}")
        except Exception as e:
            logger.warning(f"Failed to connect to Redis: {e} - session caching disabled")
            _redis_client = None

    return _redis_client


def get_database():
    """Get database instance."""
    global _db
    if _db is None:
        client = get_mongo_client()
        db_name = os.getenv("MONGO_DB_NAME", "grant_analyst")
        _db = client[db_name]
    return _db


def close_mongo_connection():
    """Close MongoDB connection."""
    global _client, _redis_client

    if _client is not None:
        _client.close()
        _client = None
        logger.info("MongoDB connection closed")

    if _redis_client is not None:
        _redis_client.close()
        _redis_client = None
        logger.info("Redis connection closed")


class SummaryStore:
    """Handle pre-computed grant summaries with optimized indexing and bulk operations."""

    def __init__(self):
        """Initialize summary store with compound indexes and text search."""
        self.db = get_database()
        self.collection = self.db["summaries"]
        self.redis = get_redis_client()

        # Create compound index for cache lookups (most common query pattern)
        self.collection.create_index(
            [("grant_id", ASCENDING), ("summary_type", ASCENDING)],
            unique=True,
            name="grant_summary_lookup"
        )

        # Create text index for full-text search on summaries
        try:
            self.collection.create_index(
                [("summary_text", TEXT)],
                name="summary_text_search",
                default_language="english"
            )
        except Exception as e:
            logger.debug(f"Text index may already exist: {e}")

        # Create index on created_at for time-based queries
        self.collection.create_index("created_at", name="created_at_idx")

        # Create index on metadata fields for analytics
        self.collection.create_index(
            [("metadata.model", ASCENDING)],
            name="model_idx",
            sparse=True
        )

    def save_summary(
        self,
        grant_id: str,
        summary_type: str,
        summary_text: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """
        Save a pre-computed summary to database.

        Args:
            grant_id: Grant ID (e.g., "competition-2315")
            summary_type: Type of summary ("layman", "technical", "exec")
            summary_text: The summary content
            metadata: Optional metadata (model used, tokens, etc.)

        Returns:
            True if saved successfully
        """
        try:
            from datetime import datetime

            doc_key = f"{grant_id}_{summary_type}"
            document = {
                "grant_id": grant_id,
                "summary_type": summary_type,
                "summary_text": summary_text,
                "created_at": datetime.utcnow(),
                "metadata": metadata or {}
            }

            # Upsert: update if exists, insert if not
            self.collection.update_one(
                {"grant_id": grant_id, "summary_type": summary_type},
                {"$set": document},
                upsert=True
            )
            logger.info(f"Summary saved: {doc_key}")
            return True

        except Exception as e:
            logger.error(f"Error saving summary for {grant_id}: {e}")
            return False

    def get_summary(self, grant_id: str, summary_type: str = "layman") -> Optional[str]:
        """
        Retrieve a pre-computed summary from cache (Redis → MongoDB).

        Args:
            grant_id: Grant ID
            summary_type: Type of summary to retrieve

        Returns:
            Summary text if found, None otherwise
        """
        cache_key = f"summary:{grant_id}:{summary_type}"

        try:
            # Try Redis first (if available)
            if self.redis is not None:
                try:
                    cached = self.redis.get(cache_key)
                    if cached:
                        logger.info(f"⚡ Redis HIT: {grant_id}_{summary_type}")
                        # Record cache hit
                        try:
                            from src.monitoring import record_cache_hit
                            record_cache_hit()
                        except Exception:
                            pass
                        return cached
                except Exception as redis_err:
                    logger.debug(f"Redis read error: {redis_err}")

            # Fall back to MongoDB
            doc = self.collection.find_one(
                {"grant_id": grant_id, "summary_type": summary_type},
                {"summary_text": 1, "_id": 0}  # Project only needed field
            )

            if doc:
                summary_text = doc.get("summary_text")
                logger.info(f"📦 MongoDB HIT: {grant_id}_{summary_type}")

                # Record cache hit
                try:
                    from src.monitoring import record_cache_hit
                    record_cache_hit()
                except Exception:
                    pass

                # Cache in Redis for next time (TTL: 1 hour)
                if self.redis is not None and summary_text:
                    try:
                        self.redis.setex(cache_key, 3600, summary_text)
                    except Exception as redis_err:
                        logger.debug(f"Redis write error: {redis_err}")

                return summary_text

            # Record cache miss
            try:
                from src.monitoring import record_cache_miss
                record_cache_miss()
            except Exception:
                pass

            return None

        except Exception as e:
            logger.error(f"Error retrieving summary for {grant_id}: {e}")
            return None

    def get_all_summaries(self, grant_id: str) -> Dict[str, str]:
        """
        Get all summary types for a grant.

        Returns:
            Dict with keys: layman, technical, exec (if available)
        """
        try:
            docs = self.collection.find({"grant_id": grant_id})
            return {doc["summary_type"]: doc["summary_text"] for doc in docs}

        except Exception as e:
            logger.error(f"Error retrieving summaries for {grant_id}: {e}")
            return {}

    def bulk_save_summaries(self, summaries: List[Dict[str, Any]]) -> int:
        """
        Bulk save multiple summaries using bulk write operations.

        Args:
            summaries: List of dicts with keys: grant_id, summary_type, summary_text, metadata

        Returns:
            Number of summaries saved

        Example:
            summaries = [
                {"grant_id": "comp-123", "summary_type": "layman", "summary_text": "...", "metadata": {}},
                {"grant_id": "comp-124", "summary_type": "layman", "summary_text": "...", "metadata": {}},
            ]
            store.bulk_save_summaries(summaries)
        """
        if not summaries:
            return 0

        try:
            from pymongo import UpdateOne

            operations = []
            for summary in summaries:
                grant_id = summary.get("grant_id")
                summary_type = summary.get("summary_type", "layman")
                summary_text = summary.get("summary_text", "")
                metadata = summary.get("metadata", {})

                if not grant_id or not summary_text:
                    logger.warning(f"Skipping invalid summary: {summary}")
                    continue

                document = {
                    "grant_id": grant_id,
                    "summary_type": summary_type,
                    "summary_text": summary_text,
                    "created_at": datetime.utcnow(),
                    "metadata": metadata
                }

                # Upsert operation
                operations.append(
                    UpdateOne(
                        {"grant_id": grant_id, "summary_type": summary_type},
                        {"$set": document},
                        upsert=True
                    )
                )

            if not operations:
                return 0

            # Execute bulk write
            result = self.collection.bulk_write(operations, ordered=False)

            saved_count = result.upserted_count + result.modified_count
            logger.info(f"💾 Bulk saved {saved_count} summaries ({result.upserted_count} new, {result.modified_count} updated)")

            return saved_count

        except BulkWriteError as bwe:
            # Log errors but don't fail completely
            logger.error(f"Bulk write errors: {bwe.details}")
            # Return count of successful writes
            return bwe.details.get("nInserted", 0) + bwe.details.get("nModified", 0)
        except Exception as e:
            logger.error(f"Error in bulk save: {e}")
            return 0

    def search_summaries(self, query: str, summary_type: Optional[str] = None, limit: int = 10) -> List[Dict[str, Any]]:
        """
        Full-text search across summaries using text index.

        Args:
            query: Search query string
            summary_type: Optional filter by summary type
            limit: Maximum results to return

        Returns:
            List of matching summaries with grant_id and summary_text
        """
        try:
            filter_dict = {"$text": {"$search": query}}

            if summary_type:
                filter_dict["summary_type"] = summary_type

            # Text search with relevance score
            results = self.collection.find(
                filter_dict,
                {"grant_id": 1, "summary_type": 1, "summary_text": 1, "score": {"$meta": "textScore"}}
            ).sort([("score", {"$meta": "textScore"})]).limit(limit)

            return list(results)

        except Exception as e:
            logger.error(f"Error searching summaries: {e}")
            return []


class GrantStore:
    """Handle grant data operations with bulk write support."""

    def __init__(self):
        """Initialize grant store with text indexes."""
        self.db = get_database()
        self.collection = self.db["grants"]

        # Compound index for common queries
        self.collection.create_index(
            [("grant_id", ASCENDING), ("status", ASCENDING)],
            name="grant_status_lookup"
        )

        # Text index for full-text search on grant titles and descriptions
        try:
            self.collection.create_index(
                [("title", TEXT), ("summary", TEXT)],
                name="grant_text_search",
                default_language="english"
            )
        except Exception as e:
            logger.debug(f"Text index may already exist: {e}")

        # Index on deadline for sorting
        self.collection.create_index("deadline", name="deadline_idx")

    def bulk_update_grants(self, grants: List[Dict[str, Any]]) -> int:
        """
        Bulk update grants from crawler.

        Args:
            grants: List of grant documents to upsert

        Returns:
            Number of grants updated
        """
        if not grants:
            return 0

        try:
            from pymongo import UpdateOne

            operations = []
            for grant in grants:
                grant_id = grant.get("id") or grant.get("grant_id")
                if not grant_id:
                    logger.warning(f"Skipping grant without ID: {grant.get('title', 'unknown')}")
                    continue

                # Upsert operation
                operations.append(
                    UpdateOne(
                        {"grant_id": grant_id},
                        {"$set": grant},
                        upsert=True
                    )
                )

            if not operations:
                return 0

            # Execute bulk write
            result = self.collection.bulk_write(operations, ordered=False)

            updated_count = result.upserted_count + result.modified_count
            logger.info(f"💾 Bulk updated {updated_count} grants ({result.upserted_count} new, {result.modified_count} updated)")

            return updated_count

        except BulkWriteError as bwe:
            logger.error(f"Bulk write errors: {bwe.details}")
            return bwe.details.get("nInserted", 0) + bwe.details.get("nModified", 0)
        except Exception as e:
            logger.error(f"Error in bulk grant update: {e}")
            return 0


class FeedbackStore:
    """Handle feedback data operations with optimized indexing."""

    def __init__(self):
        """Initialize feedback store with compound indexes."""
        self.db = get_database()
        self.collection = self.db["feedback"]

        # Compound index for user feedback history
        self.collection.create_index(
            [("user_id", ASCENDING), ("created_at", ASCENDING)],
            name="user_feedback_history"
        )

        # Index on rating for statistics
        self.collection.create_index("rating", name="rating_idx")

        # Index on created_at for time-based queries
        self.collection.create_index("created_at", name="feedback_created_at_idx")

    def save_feedback(self, feedback_data: Dict[str, Any]) -> str:
        """
        Save feedback to database.

        Args:
            feedback_data: Dictionary with feedback information

        Returns:
            String ID of the inserted feedback
        """
        try:
            result = self.collection.insert_one(feedback_data)
            logger.info(f"Feedback saved with ID: {result.inserted_id}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"Error saving feedback: {e}")
            raise

    def get_feedback_stats(self) -> Dict[str, Any]:
        """
        Get feedback statistics.

        Returns:
            Dictionary with feedback statistics
        """
        try:
            total_feedback = self.collection.count_documents({})

            # Calculate average rating (only count feedback with ratings)
            pipeline = [
                {"$match": {"rating": {"$exists": True, "$ne": None}}},
                {"$group": {"_id": None, "avg_rating": {"$avg": "$rating"}}}
            ]
            result = list(self.collection.aggregate(pipeline))
            avg_rating = result[0]["avg_rating"] if result else 0.0

            return {
                "total_feedback": total_feedback,
                "average_rating": round(avg_rating, 2)
            }
        except Exception as e:
            logger.error(f"Error getting feedback stats: {e}")
            return {
                "total_feedback": 0,
                "average_rating": 0.0
            }