db.py

from __future__ import annotations

import os
import traceback
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any

from bson.binary import Binary
from dotenv import load_dotenv
from pymongo import MongoClient
import qdrant_client
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from sentence_transformers import SentenceTransformer

from loguru import logger

load_dotenv(Path(".env"))

# logger.info(f"Qdrant version: {qdrant_client.__version__}")

MONGO_URL = os.getenv("MONGO_URL", "mongodb://localhost:27017/")
MONGO_DB_NAME = os.getenv("MONGO_DB_NAME", "grandma_voice")
QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")

_embedding_model = None

def get_embedding_model():
    global _embedding_model
    if _embedding_model is None:
        model_name = "sentence-transformers/all-MiniLM-L6-v2"
        logger.info(f"Loading embedding model: {model_name}...")

        _embedding_model = SentenceTransformer(
            model_name,
            device="cpu"   # <-- force CPU
        )

        logger.info("SentenceTransformer loaded successfully")
    return _embedding_model

def get_embedding(text: str):
    model = get_embedding_model()
    embedding = model.encode(
        text,
        convert_to_numpy=True,
        normalize_embeddings=True
    )
    return embedding.tolist()


def get_db() -> Any:
    try:
        client = MongoClient(MONGO_URL, serverSelectionTimeoutMS=5000)
        client.admin.command('ping')
        logger.info(f"Successfully connected to MongoDB at {MONGO_URL[:15]}...")
        return client[MONGO_DB_NAME]
    except Exception as e:
        logger.error(f"Failed to connect to MongoDB: {traceback.format_exc()}")
        raise


def get_qdrant_client() -> QdrantClient:
    try:
        if QDRANT_URL:
            client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
            logger.info(f"Successfully connected to Qdrant at {QDRANT_URL}")
        else:
            client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT, api_key=QDRANT_API_KEY)
            logger.info(f"Successfully connected to Qdrant at {QDRANT_HOST}:{QDRANT_PORT}")
        
        # Simple health check
        client.get_collections()
        return client
    except Exception as e:
        logger.error(f"Failed to connect to Qdrant: {traceback.format_exc()}")
        raise


def save_story_to_qdrant(hf_username: str, relationship: str, story: str) -> None:
    collection_name = f"{hf_username}-{relationship}".lower().replace("_", "-")
    logger.info(f"=== SAVE STORY TO QDRANT START ===")
    logger.info(f"Collection: {collection_name}")
    logger.info(f"Story length: {len(story)}")
    
    try:
        client = get_qdrant_client()
        
        # Ensure collection exists
        if not client.collection_exists(collection_name):
            logger.info(f"Creating new Qdrant collection: {collection_name} (384 dims, COSINE)")
            client.create_collection(
                collection_name=collection_name,
                vectors_config=VectorParams(size=384, distance=Distance.COSINE),
            )
        
        logger.info("Generating embedding...")
        embedding = get_embedding(story)
        logger.info(f"Embedding length: {len(embedding)}")
        
        point_id = str(uuid.uuid4())
        client.upsert(
            collection_name=collection_name,
            points=[
                PointStruct(
                    id=point_id,
                    vector=embedding,
                    payload={
                        "text": story,
                        "created_at": datetime.utcnow().isoformat(),
                        "relationship": relationship,
                        "user": hf_username,
                    }
                )
            ]
        )
        logger.info(f"Story successfully indexed in Qdrant. Inserted point ID: {point_id}")
    except Exception as e:
        logger.error(f"Qdrant storage failed: {traceback.format_exc()}")
        raise


def search_qdrant(hf_username: str, relationship: str, query: str, min_score: float = 0.65) -> str | None:
    collection_name = f"{hf_username}-{relationship}".lower().replace("_", "-")
    logger.info(f"Searching Qdrant for context in {collection_name} | Query: {query}")
    
    try:
        client = get_qdrant_client()
        if not client.collection_exists(collection_name):
            logger.warning(f"Collection {collection_name} does not exist. No context found.")
            return None
        
        query_embedding = get_embedding(query)
        
        # Use query_points which is the modern API
        results_response = client.query_points(
            collection_name=collection_name,
            query=query_embedding,
            limit=5,
            with_payload=True,
        )
        results = results_response.points
        
        logger.info(f"Query: {query}")
        logger.info(f"Results returned: {len(results)}")
        
        if not results:
            logger.info(f"No results found in Qdrant for query: {query}")
            return None
        
        # Log all scores for debugging
        for i, point in enumerate(results):
            score_pct = point.score * 100
            text_preview = point.payload.get("text", "")[:60] if point.payload else ""
            logger.info(f"  Result {i+1}: score={point.score:.4f} ({score_pct:.1f}%) | {text_preview}...")
        
        # Filter by minimum score threshold (75% confidence ≈ cosine ≥ 0.65)
        filtered = [p for p in results if p.score >= min_score]
        
        if not filtered:
            best_score = results[0].score
            logger.info(
                f"Best match score {best_score:.4f} ({best_score*100:.1f}%) "
                f"below threshold {min_score} ({min_score*100:.0f}%). Returning no context."
            )
            return None
        
        logger.info(
            f"Found {len(filtered)} relevant memories above threshold "
            f"{min_score} ({min_score*100:.0f}%) for {hf_username}/{relationship}"
        )
        memory_texts = [p.payload.get("text", "") for p in filtered if p.payload]
        return "\n".join(memory_texts)
        
    except Exception as e:
        logger.error(f"Qdrant search failed: {traceback.format_exc()}")
        return None

def profiles_collection() -> Any:
    return get_db()["voice_profiles"]


def save_profile(hf_username: str, relationship: str, audio_bytes: bytes) -> str:
    logger.info(f"Saving voice profile for user={hf_username}, relationship={relationship}")
    document = {
        "hf_username": hf_username,
        "relationship": relationship,
        "audio_bytes": Binary(audio_bytes),
        "created_at": datetime.utcnow(),
    }
    
    coll = profiles_collection()
    try:
        result = coll.insert_one(document)
        logger.info(f"Inserted MongoDB document ID: {result.inserted_id} for {hf_username}/{relationship}")
        return str(result.inserted_id)
    except Exception as exc:
        logger.error(f"MongoDB save failed: {traceback.format_exc()}")
        raise


def list_profiles_for_user(hf_username: str) -> list[dict[str, str]]:
    if not hf_username:
        return []

    logger.info(f"Fetching profiles for user: {hf_username}")
    coll = profiles_collection()
    
    query = {"hf_username": hf_username}
    results = list(coll.find(query).sort("created_at", -1))
    logger.info(f"Retrieved {len(results)} profiles")
    
    return [
        {
            "id": str(doc["_id"]),
            "relationship": doc.get("relationship", ""),
            "created_at": doc.get("created_at", datetime.utcnow()).strftime("%Y-%m-%d %H:%M:%S UTC"),
        }
        for doc in results
    ]


def list_relationships_for_user(hf_username: str) -> list[str]:
    if not hf_username:
        return []

    logger.info(f"Listing relationships for user: {hf_username}")
    query = {"hf_username": hf_username}
    results = profiles_collection().find(query, {"relationship": 1})
    relationships: set[str] = set()
    for doc in results:
        relationship = doc.get("relationship")
        if relationship:
            relationships.add(relationship)
    
    logger.info(f"Found {len(relationships)} unique relationships for user {hf_username}")
    return sorted(relationships)


def relationship_exists_for_user(hf_username: str, relationship: str) -> bool:
    if not hf_username or not relationship:
        return False
    query = {"hf_username": hf_username, "relationship": relationship}
    exists = profiles_collection().count_documents(query, limit=1) > 0
    logger.info(f"Checking if relationship '{relationship}' exists for {hf_username}: {exists}")
    return exists


def get_latest_profile_audio(hf_username: str, relationship: str) -> bytes | None:
    if not hf_username or not relationship:
        return None

    logger.info(f"Fetching latest audio for {hf_username}/{relationship}")
    query = {"hf_username": hf_username, "relationship": relationship}
    document = profiles_collection().find_one(query, sort=[("created_at", -1)])
    if document is None:
        logger.warning(f"No audio document found for {hf_username}/{relationship}")
        return None

    logger.info(f"Audio found for {hf_username}/{relationship}, size: {len(document['audio_bytes'])} bytes")
    return bytes(document["audio_bytes"])

def debug_database():
    try:
        db = get_db()
        logger.info("=== DATABASE DIAGNOSTICS ===")
        logger.info(f"Database Name: {db.name}")

        collections = db.list_collection_names()
        logger.info(f"Collections: {collections}")

        for collection in collections:
            count = db[collection].count_documents({})
            logger.info(f"{collection}: {count} documents")
        logger.info("============================")
    except Exception as e:
        logger.error(f"Database diagnostics failed: {traceback.format_exc()}")