"""
Memory Consolidation Plugin for Cheshire Cat

Phase 2: Sleep Consolidation Implementation

Implements human-like memory consolidation:
1. During the day: Store almost everything temporarily
2. On demand (or scheduled): Analyze conversations, keep important, delete trivial
3. Extract facts for declarative memory (per-user)

This mimics how human brains consolidate memories during REM sleep.
"""

from cat.mad_hatter.decorators import hook, tool
from datetime import datetime
import json
import os
from typing import List, Dict, Any

print("\U0001f319 [Consolidation Plugin] Loading...")

# Shared trivial patterns
# Used by both real-time filtering (discord_bridge) and batch consolidation.
# Keep this in sync with discord_bridge's skip_patterns.
TRIVIAL_PATTERNS = frozenset([
    'lol', 'k', 'ok', 'okay', 'haha', 'lmao', 'xd', 'rofl', 'lmfao',
    'brb', 'gtg', 'afk', 'ttyl', 'lmk', 'idk', 'tbh', 'imo', 'imho',
    'omg', 'wtf', 'fyi', 'btw', 'nvm', 'jk', 'ikr', 'smh',
    'hehe', 'heh', 'gg', 'wp', 'gz', 'gj', 'ty', 'thx', 'np', 'yw',
    'nice', 'cool', 'neat', 'wow', 'yep', 'nope', 'yeah', 'nah',
])

# Consolidation state
consolidation_state = {
    'last_run': None,
    'is_running': False,
    'stats': {
        'total_processed': 0,
        'kept': 0,
        'deleted': 0,
        'facts_learned': 0
    }
}


# ===================================================================
# HOOKS
# ===================================================================

@hook(priority=50)
def after_cat_bootstrap(cat):
    """Run after Cat starts up."""
    print("\U0001f319 [Memory Consolidation] Plugin loaded")
    print("   Manual consolidation available via 'consolidate now' command")


@hook(priority=10)
def agent_prompt_prefix(prefix, cat):
    """
    Runs AFTER miku_personality (priority 100) sets the base prompt.
    1. Search and inject declarative facts into the prompt.
    2. Handle 'consolidate now' command.
    """
    user_message_json = cat.working_memory.get('user_message_json', {})
    user_text = user_message_json.get('text', '').strip()

    # PART 1: Inject declarative facts
    try:
        if user_text and user_text.lower() not in ('consolidate', 'consolidate now', '/consolidate'):
            declarative_memory = cat.memory.vectors.declarative
            embedding = cat.embedder.embed_query(user_text)

            results = declarative_memory.recall_memories_from_embedding(
                embedding=embedding,
                metadata={"source": cat.user_id},
                k=5
            )

            if results:
                high_confidence_facts = [
                    item[0].page_content
                    for item in results
                    if item[1] > 0.5
                ]

                if high_confidence_facts:
                    facts_text = "\n\n## Personal Facts About the User:\n"
                    for fact in high_confidence_facts:
                        facts_text += f"- {fact}\n"
                    facts_text += "\n(Use these facts when answering the user's question)\n"
                    prefix += facts_text
                    print(f"[Declarative] Injected {len(high_confidence_facts)} facts into prompt")

    except Exception as e:
        print(f"[Declarative] Error: {e}")

    # PART 2: Handle consolidation command
    if user_text.lower() in ('consolidate', 'consolidate now', '/consolidate'):
        print("[Consolidation] Manual trigger command received!")
        trigger_consolidation_sync(cat)

        stats = consolidation_state['stats']
        cat.working_memory['consolidation_triggered'] = True
        cat.working_memory['consolidation_stats'] = stats

    return prefix


@hook(priority=10)
def before_cat_sends_message(message, cat):
    """
    1. Replace response with consolidation stats if consolidation was triggered.
    2. Store Miku's response in episodic memory (bidirectional memory).
    """

    # PART 1: Consolidation response replacement
    if cat.working_memory.get('consolidation_triggered', False):
        print("[Consolidation] Replacing message with stats")
        stats = cat.working_memory.get('consolidation_stats', {})
        output_str = (
            f"\U0001f319 **Memory Consolidation Complete!**\n\n"
            f"**Stats:**\n"
            f"- Total processed: {stats.get('total_processed', 0)}\n"
            f"- Kept: {stats.get('kept', 0)}\n"
            f"- Deleted: {stats.get('deleted', 0)}\n"
            f"- Facts learned: {stats.get('facts_learned', 0)}\n"
        )
        cat.working_memory['consolidation_triggered'] = False

        if hasattr(message, 'content'):
            message.content = output_str
        else:
            message['content'] = output_str

    # PART 2: Store Miku's response in episodic memory
    try:
        if hasattr(message, 'content'):
            miku_response = message.content
        elif isinstance(message, dict):
            miku_response = message.get('content', '')
        else:
            miku_response = str(message)

        if miku_response and len(miku_response.strip()) > 3:
            metadata = {
                'source': cat.user_id,
                'when': datetime.now().timestamp(),
                'stored_at': datetime.now().isoformat(),
                'speaker': 'miku',
                'consolidated': False,
                'guild_id': cat.working_memory.get('guild_id', 'dm'),
                'channel_id': cat.working_memory.get('channel_id'),
            }

            response_text = f"[Miku]: {miku_response}"
            vector = cat.embedder.embed_query(response_text)

            cat.memory.vectors.episodic.add_point(
                content=response_text,
                vector=vector,
                metadata=metadata
            )
            print(f"[Miku Memory] Stored response: {miku_response[:50]}...")

    except Exception as e:
        print(f"[Miku Memory] Error storing response: {e}")

    return message


# ===================================================================
# CONSOLIDATION ENGINE
# ===================================================================

def trigger_consolidation_sync(cat):
    """
    Synchronous consolidation for use in hooks.
    Processes ALL unconsolidated memories across all users.
    """
    from qdrant_client import QdrantClient

    print("[Consolidation] Starting synchronous consolidation...")

    qdrant_host = os.getenv('QDRANT_HOST', 'localhost')
    qdrant_port = int(os.getenv('QDRANT_PORT', 6333))
    client = QdrantClient(host=qdrant_host, port=qdrant_port)

    # Query all unconsolidated memories
    result = client.scroll(
        collection_name='episodic',
        scroll_filter={
            "must_not": [
                {"key": "metadata.consolidated", "match": {"value": True}}
            ]
        },
        limit=10000,
        with_payload=True,
        with_vectors=False
    )

    memories = result[0]
    print(f"[Consolidation] Found {len(memories)} unconsolidated memories")

    if not memories:
        consolidation_state['stats'] = {
            'total_processed': 0, 'kept': 0, 'deleted': 0, 'facts_learned': 0
        }
        return

    # Classify memories
    to_delete = []
    to_mark_consolidated = []
    # Group user messages by source (user_id) for per-user fact extraction
    user_messages_by_source = {}

    for point in memories:
        content = point.payload.get('page_content', '').strip()
        content_lower = content.lower()
        metadata = point.payload.get('metadata', {})

        is_miku_message = (
            metadata.get('speaker') == 'miku'
            or content.startswith('[Miku]:')
        )

        # Check if trivial
        is_trivial = content_lower in TRIVIAL_PATTERNS

        if is_trivial:
            to_delete.append(point.id)
        else:
            to_mark_consolidated.append(point.id)
            # Only user messages go to fact extraction, grouped by user
            if not is_miku_message:
                source = metadata.get('source', 'unknown')
                if source not in user_messages_by_source:
                    user_messages_by_source[source] = []
                user_messages_by_source[source].append(point.id)

    # Delete trivial memories
    if to_delete:
        client.delete(collection_name='episodic', points_selector=to_delete)
        print(f"[Consolidation] Deleted {len(to_delete)} trivial memories")

    # Mark kept memories as consolidated
    if to_mark_consolidated:
        for point_id in to_mark_consolidated:
            client.set_payload(
                collection_name='episodic',
                payload={"metadata.consolidated": True},
                points=[point_id]
            )
        print(f"[Consolidation] Marked {len(to_mark_consolidated)} memories as consolidated")

    # Extract facts per user
    total_facts = 0
    for source_user_id, memory_ids in user_messages_by_source.items():
        print(f"[Consolidation] Extracting facts for user '{source_user_id}' from {len(memory_ids)} messages...")
        facts = extract_and_store_facts(client, memory_ids, cat, source_user_id)
        total_facts += facts
        print(f"[Consolidation] Extracted {facts} facts for user '{source_user_id}'")

    consolidation_state['stats'] = {
        'total_processed': len(memories),
        'kept': len(to_mark_consolidated),
        'deleted': len(to_delete),
        'facts_learned': total_facts
    }

    print("[Consolidation] Synchronous consolidation complete!")
    return True


# ===================================================================
# FACT EXTRACTION
# ===================================================================

def extract_and_store_facts(client, memory_ids, cat, user_id):
    """
    Extract declarative facts from user memories using LLM and store them.
    Facts are scoped to the specific user_id.
    Uses Cat's embedder to ensure vector compatibility.
    Deduplicates against existing facts before storing.
    """
    import uuid

    if not memory_ids:
        return 0

    memories = client.retrieve(collection_name='episodic', ids=memory_ids)
    facts_stored = 0

    # Process in batches of 5
    batch_size = 5
    for i in range(0, len(memories), batch_size):
        batch = memories[i:i + batch_size]

        conversation_context = "\n".join([
            f"- {mem.payload.get('page_content', '')}"
            for mem in batch
        ])

        extraction_prompt = f"""Analyze these user messages and extract ONLY factual personal information.

User messages:
{conversation_context}

Extract facts in this exact format (one per line):
- The user's name is [name]
- The user is [age] years old
- The user lives in [location]
- The user works as [job]
- The user is allergic to [allergen]
- The user's favorite color is [color]
- The user enjoys [hobby/activity]
- The user prefers [preference]
- The user's birthday is [date]
- The user graduated from [school/university]

IMPORTANT:
- Only include facts that are CLEARLY stated in the messages
- Use the EXACT format shown above
- If no facts found, respond with: "No facts found"
- Do not include greetings, questions, or opinions
- Do not invent or assume facts not explicitly stated
"""

        try:
            response = cat.llm(extraction_prompt)
            print(f"[LLM Extract] Response:\n{response[:200]}...")

            lines = response.strip().split('\n')
            for line in lines:
                line = line.strip()

                if not line or line.lower().startswith(('no facts', '#', 'user messages:', '```')):
                    continue

                if line.startswith('- The user'):
                    fact_text = line[2:].strip()  # Remove "- " prefix

                    # Determine fact type
                    fact_type = 'general'
                    fact_value = fact_text

                    if "'s name is" in fact_text:
                        fact_type = 'name'
                        fact_value = fact_text.split("'s name is")[-1].strip()
                    elif " is " in fact_text and " years old" in fact_text:
                        fact_type = 'age'
                        fact_value = fact_text.split(" is ")[1].split(" years")[0].strip()
                    elif "lives in" in fact_text:
                        fact_type = 'location'
                        fact_value = fact_text.split("lives in")[-1].strip()
                    elif "works as" in fact_text:
                        fact_type = 'job'
                        fact_value = fact_text.split("works as")[-1].strip()
                    elif "allergic to" in fact_text:
                        fact_type = 'allergy'
                        fact_value = fact_text.split("allergic to")[-1].strip()
                    elif "favorite color is" in fact_text:
                        fact_type = 'favorite_color'
                        fact_value = fact_text.split("favorite color is")[-1].strip()
                    elif "enjoys" in fact_text:
                        fact_type = 'hobby'
                        fact_value = fact_text.split("enjoys")[-1].strip()
                    elif "prefers" in fact_text:
                        fact_type = 'preference'
                        fact_value = fact_text.split("prefers")[-1].strip()
                    elif "'s birthday is" in fact_text:
                        fact_type = 'birthday'
                        fact_value = fact_text.split("'s birthday is")[-1].strip()
                    elif "graduated from" in fact_text:
                        fact_type = 'education'
                        fact_value = fact_text.split("graduated from")[-1].strip()

                    # Duplicate detection
                    if _is_duplicate_fact(client, cat, fact_text, fact_type, user_id):
                        print(f"[Fact Skip] Duplicate: {fact_text}")
                        continue

                    # Store fact using Cat's embedder
                    fact_embedding = cat.embedder.embed_query(fact_text)
                    point_id = str(uuid.uuid4())

                    client.upsert(
                        collection_name='declarative',
                        points=[{
                            'id': point_id,
                            'vector': fact_embedding,
                            'payload': {
                                'page_content': fact_text,
                                'metadata': {
                                    'source': user_id,
                                    'when': datetime.now().timestamp(),
                                    'fact_type': fact_type,
                                    'fact_value': fact_value,
                                }
                            }
                        }]
                    )

                    facts_stored += 1
                    print(f"[Fact Stored] [{user_id}] {fact_text}")

        except Exception as e:
            print(f"[LLM Extract] Error: {e}")
            import traceback
            traceback.print_exc()

    return facts_stored


def _is_duplicate_fact(client, cat, fact_text, fact_type, user_id):
    """
    Check if a similar fact already exists for this user.
    Uses vector similarity to detect semantic duplicates.
    """
    try:
        fact_embedding = cat.embedder.embed_query(fact_text)

        # Search existing facts for this user with same fact_type
        results = client.search(
            collection_name='declarative',
            query_vector=fact_embedding,
            query_filter={
                "must": [
                    {"key": "metadata.source", "match": {"value": user_id}},
                    {"key": "metadata.fact_type", "match": {"value": fact_type}},
                ]
            },
            limit=1,
            score_threshold=0.85  # High threshold = very similar
        )

        if results:
            existing = results[0].payload.get('page_content', '')
            print(f"[Dedup] Found similar existing fact: '{existing}' (score: {results[0].score:.2f})")
            return True

        return False

    except Exception as e:
        print(f"[Dedup] Error checking duplicates: {e}")
        return False  # On error, allow storing


# ===================================================================
# TOOL (for Cat's tool system)
# ===================================================================

@tool(return_direct=True)
def consolidate_memories(tool_input, cat):
    """Use this tool to consolidate memories. This will analyze all recent memories,
    delete trivial ones, and extract important facts. Input is always an empty string."""

    print("[Consolidation] Tool called!")
    trigger_consolidation_sync(cat)

    stats = consolidation_state['stats']
    return (
        f"\U0001f319 **Memory Consolidation Complete!**\n\n"
        f"**Stats:**\n"
        f"- Total processed: {stats['total_processed']}\n"
        f"- Kept: {stats['kept']}\n"
        f"- Deleted: {stats['deleted']}\n"
        f"- Facts learned: {stats['facts_learned']}\n"
    )


# ===================================================================
# PLUGIN METADATA
# ===================================================================

__version__ = "2.0.0"
__description__ = "Sleep consolidation - analyze memories, keep important, delete trivial, extract per-user facts"

print("[Consolidation Plugin] All hooks registered")
