"""
NeuraPrompt Agent — Memory Module v7.6 (MongoDB Integrated)
Properly connects to your existing MongoDB system from main.py
"""

from typing import List, Dict, Optional
import logging
from datetime import datetime, timezone

log = logging.getLogger("agent.memory.v7.6")

# These will be injected from main.py
long_term_memory_col = None
chat_history_col = None


def set_collections(long_term_col, chat_history_col_ref):
    """Call this from main.py to inject MongoDB collections"""
    global long_term_memory_col, chat_history_col
    long_term_memory_col = long_term_col
    chat_history_col = chat_history_col_ref


class MemoryManager:
    def __init__(self, user_id: str):
        self.user_id = user_id
        self.short_term: List[Dict] = []

    def load_long_term_memory(self) -> Dict:
        """Load user's long-term memory from MongoDB"""
        if not long_term_memory_col:
            return {}
        try:
            doc = long_term_memory_col.find_one({"user_id": self.user_id})
            return doc or {}
        except Exception as e:
            log.error(f"Failed to load long-term memory: {e}")
            return {}

    def load_recent_chat(self, limit: int = 8) -> List[Dict]:
        """Load recent chat history from MongoDB"""
        if not chat_history_col:
            return []
        try:
            cursor = chat_history_col.find(
                {"user_id": self.user_id}
            ).sort("timestamp", -1).limit(limit)
            messages = list(cursor)
            messages.reverse()
            return messages
        except Exception as e:
            log.error(f"Failed to load chat history: {e}")
            return []

    def get_context_for_agent(self) -> str:
        """Build rich context string for the agent"""
        long_term = self.load_long_term_memory()
        recent = self.load_recent_chat(6)

        context_parts = []

        # Long-term facts
        if long_term:
            facts = []
            for key, value in long_term.items():
                if key not in ["_id", "user_id", "last_updated"]:
                    facts.append(f"{key}: {value}")
            if facts:
                context_parts.append("Known about user: " + ", ".join(facts[:7]))

        # Recent conversation
        if recent:
            chat_lines = []
            for msg in recent[-5:]:
                role = msg.get("role", "unknown").capitalize()
                content = msg.get("content", "")[:180]
                chat_lines.append(f"{role}: {content}")
            context_parts.append("Recent conversation:\n" + "\n".join(chat_lines))

        return "\n\n".join(context_parts) if context_parts else "No previous memory available."

    def save_important_fact(self, key: str, value: str):
        """Save important fact to long-term memory"""
        if not long_term_memory_col:
            return
        try:
            long_term_memory_col.update_one(
                {"user_id": self.user_id},
                {"$set": {
                    key: value,
                    "last_updated": datetime.now(timezone.utc)
                }},
                upsert=True
            )
            log.info(f"Saved fact '{key}' for user {self.user_id}")
        except Exception as e:
            log.error(f"Failed to save fact: {e}")


def get_memory_manager(user_id: str) -> MemoryManager:
    return MemoryManager(user_id)