"""Mem0-powered memory system for intelligent semantic retrieval."""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Any

from loguru import logger

if TYPE_CHECKING:
    from nanobot.providers.base import LLMProvider
    from nanobot.session.manager import Session

try:
    from mem0 import Memory
    from mem0.configs.base import MemoryConfig
    HAS_MEM0 = True
except ImportError:
    HAS_MEM0 = False
    MemoryConfig = None  # type: ignore


class Mem0MemoryStore:
    """
    Enhanced memory store using mem0 for semantic search and automatic extraction.

    Features:
    - Multi-level memory (user, session, agent)
    - Semantic search with embeddings
    - Automatic memory extraction from conversations
    - 90% token reduction vs full-context
    - 91% faster responses
    """

    def __init__(self, workspace: Path, config: dict[str, Any] | None = None):
        if not HAS_MEM0:
            raise ImportError(
                "mem0 not installed. Install with: pip install mem0ai"
            )

        self.workspace = workspace
        self.memory_dir = workspace / "memory"
        self.memory_dir.mkdir(parents=True, exist_ok=True)

        # Build custom extraction prompt tuned for nanobot conversations
        from datetime import datetime

        today = datetime.now().strftime("%Y-%m-%d")
        custom_prompt = f"Extract dated facts from this conversation as JSON: {{\"facts\": [...]}}. Today is {today}.\n\n"
        self.custom_prompt = custom_prompt

        # Initialize mem0 with optional config + custom prompt
        # Extract only MemoryConfig-relevant fields
        raw_config = config if config else {}
        logger.debug(f"Mem0MemoryStore received config keys: {list(raw_config.keys())}")
        mem0_cfg_dict = {}
        for key in ("vector_store", "llm", "embedder", "graph_store", "version"):
            if key in raw_config:
                mem0_cfg_dict[key] = raw_config[key]
        logger.debug(f"Extracted for MemoryConfig: {list(mem0_cfg_dict.keys())}")
        mem0_config = MemoryConfig(**mem0_cfg_dict)
        logger.debug(f"MemoryConfig created: vector_store={mem0_config.vector_store.provider if mem0_config.vector_store else None}")
        self.memory = Memory(config=mem0_config)

        logger.info("Mem0 memory system initialized with custom nanobot prompt")

    def search_memories(
        self,
        query: str,
        user_id: str,
        limit: int = 5,
        session_id: str | None = None,
    ) -> list[dict[str, Any]]:
        """
        Search for relevant memories using semantic search.

        Args:
            query: Search query (user's current message)
            user_id: User identifier (e.g., "telegram_12345")
            limit: Max number of memories to return
            session_id: Optional session-specific memories

        Returns:
            List of memory dicts with 'memory' and 'score' keys
        """
        try:
            # Search user-level memories
            user_memories = self.memory.search(
                query=query,
                user_id=user_id,
                limit=limit
            )

            results = []
            if user_memories and "results" in user_memories:
                results.extend(user_memories["results"])

            # Optionally search session-level memories
            if session_id:
                session_memories = self.memory.search(
                    query=query,
                    user_id=user_id,
                    metadata={"session_id": session_id},
                    limit=limit // 2  # Reserve half for session context
                )
                if session_memories and "results" in session_memories:
                    results.extend(session_memories["results"])

            logger.debug(
                f"Mem0 search: query='{query[:50]}...', found {len(results)} memories"
            )
            return results[:limit]  # Limit total results

        except Exception as e:
            logger.error(f"Mem0 search failed: {e}")
            return []

    def add_conversation(
        self,
        messages: list[dict[str, Any]],
        user_id: str,
        session_id: str | None = None,
    ) -> None:
        """
        Add conversation messages to memory for automatic extraction.

        Args:
            messages: List of message dicts with 'role' and 'content'
            user_id: User identifier
            session_id: Optional session identifier for session-level memories
        """
        try:
            metadata = {}
            if session_id:
                metadata["session_id"] = session_id

            # mem0 automatically extracts and stores relevant facts
            result = self.memory.add(
                messages,
                user_id=user_id,
                metadata=metadata if metadata else None
            )

            facts_count = len(result.get("results", [])) if result else 0
            logger.debug(
                f"Mem0 add: {len(messages)} messages for user {user_id}, extracted {facts_count} facts"
            )
        except Exception as e:
            logger.error(f"Mem0 add failed: {e}")

    async def extract_facts(
        self,
        messages: list[dict[str, Any]],
        provider: Any,
        model: str,
    ) -> list[str]:
        """Extract facts from conversation using the main agent's LLM provider."""
        import json as _json

        conv_text = ""
        for msg in messages:
            role = msg.get("role", "unknown")
            content_val = msg.get("content", "")
            if isinstance(content_val, str) and content_val.strip():
                conv_text += f"{role}: {content_val}\n\n"

        if not conv_text.strip():
            return []

        extraction_messages = [
            {"role": "user", "content": self.custom_prompt + conv_text}
        ]

        try:
            response = await provider.chat(
                messages=extraction_messages,
                model=model,
                max_tokens=2000,
                temperature=0.3,
                thinking_budget=0,
            )
            text = (response.content or "").strip()
            if text.startswith("```"):
                text = text.split("```")[1]
                if text.startswith("json"):
                    text = text[4:]
                text = text.strip()
            data = _json.loads(text)
            facts = data.get("facts", [])
            if not isinstance(facts, list):
                logger.warning(f"LLM returned non-list facts: {type(facts)}")
                return []
            logger.debug(f"Extracted {len(facts)} facts using {model}")
            return facts
        except Exception as e:
            logger.error(f"Fact extraction failed: {e}")
            return []

    def store_facts(
        self,
        facts: list[str],
        user_id: str,
        session_id: str | None = None,
    ) -> None:
        """Store pre-extracted facts in mem0 with infer=False."""
        if not facts:
            return

        metadata = {}
        if session_id:
            metadata["session_id"] = session_id

        stored = 0
        for fact in facts:
            try:
                self.memory.add(
                    fact,
                    user_id=user_id,
                    infer=False,
                    metadata=metadata if metadata else None,
                )
                stored += 1
            except Exception as e:
                logger.error(f"Failed to store fact '{fact[:50]}...': {e}")

        logger.info(f"Stored {stored}/{len(facts)} facts for user {user_id}")

    def get_memory_context(
        self,
        query: str,
        user_id: str,
        limit: int = 5
    ) -> str:
        """
        Get formatted memory context for inclusion in system prompt.

        Args:
            query: Current user query
            user_id: User identifier
            limit: Max memories to include

        Returns:
            Formatted memory context string
        """
        memories = self.search_memories(query, user_id, limit=limit)

        if not memories:
            return ""

        lines = ["## Relevant Memories"]
        for i, mem in enumerate(memories, 1):
            memory_text = mem.get("memory", "")
            # Include score if available for debugging
            score = mem.get("score", "")
            score_str = f" (relevance: {score:.2f})" if score else ""
            lines.append(f"{i}. {memory_text}{score_str}")

        return "\n".join(lines)

    def update_memory(self, memory_id: str, data: dict[str, Any]) -> None:
        """Update a specific memory by ID."""
        try:
            self.memory.update(memory_id, data)
            logger.debug(f"Mem0 update: memory_id={memory_id}")
        except Exception as e:
            logger.error(f"Mem0 update failed: {e}")

    def delete_memory(self, memory_id: str) -> None:
        """Delete a specific memory by ID."""
        try:
            self.memory.delete(memory_id)
            logger.debug(f"Mem0 delete: memory_id={memory_id}")
        except Exception as e:
            logger.error(f"Mem0 delete failed: {e}")

    def get_all_memories(self, user_id: str) -> list[dict[str, Any]]:
        """Get all memories for a user."""
        try:
            result = self.memory.get_all(user_id=user_id)
            return result.get("results", []) if result else []
        except Exception as e:
            logger.error(f"Mem0 get_all failed: {e}")
            return []

    async def consolidate(
        self,
        session: Session,
        provider: LLMProvider,
        model: str,
        *,
        archive_all: bool = False,
        memory_window: int = 50,
    ) -> bool:
        """
        Consolidate session messages into mem0 memory.

        Unlike the original MemoryStore, mem0 handles extraction automatically,
        so this just needs to feed recent messages to mem0.

        Returns True on success.
        """
        try:
            # Extract user_id from session key (e.g., "telegram:12345" -> "telegram_12345")
            user_id = session.key.replace(":", "_")

            # Determine which messages to consolidate
            if archive_all:
                messages_to_add = session.messages
                logger.info(
                    f"Mem0 consolidation (archive_all): {len(messages_to_add)} messages"
                )
            else:
                keep_count = memory_window // 2
                if len(session.messages) <= keep_count:
                    return True

                # Consolidate messages except the most recent (kept for context)
                start_idx = 0
                end_idx = len(session.messages) - keep_count

                if end_idx <= start_idx:
                    return True

                messages_to_add = session.messages[start_idx:end_idx]

                if not messages_to_add:
                    return True

                logger.info(
                    f"Mem0 consolidation: {len(messages_to_add)} to consolidate, "
                    f"{keep_count} keep"
                )

            # Convert to mem0 format with intelligent filtering
            mem0_messages = []
            for msg in messages_to_add:
                role = msg.get("role")
                content = msg.get("content")

                # Skip tool results — raw bash output, file contents, and JSON
                # get misinterpreted by the extraction LLM as user interests
                if role == "tool":
                    continue

                # Skip system messages — they're boilerplate instructions, not facts
                if role == "system":
                    continue

                # Skip messages with no content
                if not content:
                    continue

                # Normalize assistant message content: extract text from Anthropic list format
                if role == "assistant" and isinstance(content, list):
                    # Anthropic format: list of {type: "text"|"tool_use", text: "..."} blocks
                    text_parts = [
                        block.get("text", "")
                        for block in content
                        if isinstance(block, dict) and block.get("type") == "text"
                    ]
                    content = " ".join(text_parts).strip()
                    if not content:
                        continue  # Skip if assistant only called tools with no text explanation

                # Normalize user message content (could also be a list in some formats)
                if isinstance(content, list):
                    text_parts = [
                        block.get("text", "") if isinstance(block, dict) else str(block)
                        for block in content
                    ]
                    content = " ".join(text_parts).strip()
                    if not content:
                        continue

                # Skip trivially short messages (commands like "/new")
                if len(content.strip()) < 10:
                    continue

                mem0_messages.append({
                    "role": role,
                    "content": content
                })

            if mem0_messages:
                # Extract facts using the main agent's LLM (already paid for),
                # then store with infer=False to bypass mem0's GPT-nano
                facts = await self.extract_facts(mem0_messages, provider, model)
                self.store_facts(facts, user_id=user_id, session_id=session.key)

            logger.info(
                f"Mem0 consolidation done: {len(session.messages)} messages total"
            )
            return True

        except Exception:
            logger.exception("Mem0 consolidation failed")
            return False
