import gradio as gr
import json
import asyncio
from datetime import datetime
from typing import List, Dict, Any
import uuid
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
from pathlib import Path

class LongTermMemoryDemo:
    def __init__(self):
        self.db_path = "./memory_db"
        Path(self.db_path).mkdir(exist_ok=True)
        
        # Initialize ChromaDB
        self.chroma_client = chromadb.PersistentClient(
            path=self.db_path,
            settings=Settings(anonymized_telemetry=False)
        )
        
        # Get or create collection
        try:
            self.collection = self.chroma_client.get_collection("memories")
        except:
            self.collection = self.chroma_client.create_collection(
                name="memories",
                metadata={"description": "Long-term memory storage for conversations"}
            )
        
        # Initialize sentence transformer for embeddings
        print("Loading SentenceTransformer model...")
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        print("Model loaded successfully!")
    
    def save_memory(self, content: str, title: str, tags: str = "", context: str = "") -> str:
        """Save content to long-term memory."""
        if not content or not title:
            return "❌ Error: Content and title are required!"
        
        try:
            tags_list = [tag.strip() for tag in tags.split(',') if tag.strip()] if tags else []
            
            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            
            # Create embedding
            embedding = self.encoder.encode(f"{title} {content}").tolist()
            
            # Prepare metadata
            metadata = {
                "title": title,
                "timestamp": timestamp,
                "tags": json.dumps(tags_list),
                "context": context,
                "content_length": len(content)
            }
            
            # Save to ChromaDB
            self.collection.add(
                documents=[content],
                embeddings=[embedding],
                metadatas=[metadata],
                ids=[memory_id]
            )
            
            result = f"✅ **Memory saved successfully!**\n\n"
            result += f"**ID**: `{memory_id}`\n"
            result += f"**Title**: {title}\n"
            result += f"**Timestamp**: {timestamp}\n"
            if tags_list:
                result += f"**Tags**: {', '.join(tags_list)}\n"
            if context:
                result += f"**Context**: {context}\n"
            result += f"**Content Preview**: {content[:200]}{'...' if len(content) > 200 else ''}"
            
            return result
        except Exception as e:
            return f"❌ Error saving memory: {str(e)}"
    
    def search_memory(self, query: str, limit: int = 5, threshold: float = 0.3) -> str:
        """Search through memories."""
        if not query:
            return "❌ Error: Search query is required!"
        
        try:
            if self.collection.count() == 0:
                return "📭 No memories stored yet. Save some memories first!"
            
            # Create query embedding
            query_embedding = self.encoder.encode(query).tolist()
            
            # Search in ChromaDB
            results = self.collection.query(
                query_embeddings=[query_embedding],
                n_results=min(limit, self.collection.count())
            )
            
            if not results['documents'][0]:
                return "🔍 No relevant memories found."
            
            # Filter by threshold and format results
            response = f"🔍 **Search Results for**: \"{query}\"\n\n"
            
            found_relevant = False
            for i, (doc, metadata, distance) in enumerate(zip(
                results['documents'][0], 
                results['metadatas'][0], 
                results['distances'][0]
            )):
                similarity = 1 - distance
                if similarity >= threshold:
                    found_relevant = True
                    tags = json.loads(metadata.get('tags', '[]'))
                    
                    response += f"### {i+1}. {metadata['title']} (Similarity: {similarity:.2f})\n"
                    response += f"**Saved**: {metadata['timestamp']}\n"
                    if tags:
                        response += f"**Tags**: {', '.join(tags)}\n"
                    if metadata.get('context'):
                        response += f"**Context**: {metadata['context']}\n"
                    response += f"**Content**: {doc}\n\n"
                    response += "---\n\n"
            
            if not found_relevant:
                response += f"No memories found above similarity threshold of {threshold:.2f}"
            
            return response
        except Exception as e:
            return f"❌ Error searching memories: {str(e)}"
    
    def list_memories(self, limit: int = 10) -> str:
        """List all memories."""
        try:
            if self.collection.count() == 0:
                return "📭 No memories stored yet."
            
            # Get all memories
            results = self.collection.get()
            
            if not results['documents']:
                return "📭 No memories found."
            
            response = f"📚 **All Memories** (showing up to {limit})\n\n"
            
            # Sort by timestamp (newest first)
            memories = list(zip(results['ids'], results['documents'], results['metadatas']))
            memories.sort(key=lambda x: x[2]['timestamp'], reverse=True)
            
            for i, (memory_id, doc, metadata) in enumerate(memories[:limit]):
                tags = json.loads(metadata.get('tags', '[]'))
                
                response += f"### {i+1}. {metadata['title']}\n"
                response += f"**ID**: `{memory_id}`\n"
                response += f"**Saved**: {metadata['timestamp']}\n"
                if tags:
                    response += f"**Tags**: {', '.join(tags)}\n"
                response += f"**Preview**: {doc[:150]}{'...' if len(doc) > 150 else ''}\n\n"
                response += "---\n\n"
            
            if len(memories) > limit:
                response += f"... and {len(memories) - limit} more memories"
            
            return response
        except Exception as e:
            return f"❌ Error listing memories: {str(e)}"
    
    def get_memory_stats(self) -> str:
        """Get statistics about stored memories."""
        try:
            count = self.collection.count()
            if count == 0:
                return "📊 **Memory Statistics**: No memories stored yet."
            
            results = self.collection.get()
            
            # Calculate stats
            total_content_length = sum(metadata['content_length'] for metadata in results['metadatas'])
            avg_content_length = total_content_length / count if count > 0 else 0
            
            # Get all tags
            all_tags = []
            for metadata in results['metadatas']:
                tags = json.loads(metadata.get('tags', '[]'))
                all_tags.extend(tags)
            
            unique_tags = list(set(all_tags))
            
            stats = f"📊 **Memory Statistics**\n\n"
            stats += f"**Total Memories**: {count}\n"
            stats += f"**Total Content Length**: {total_content_length:,} characters\n"
            stats += f"**Average Content Length**: {avg_content_length:.0f} characters\n"
            stats += f"**Unique Tags**: {len(unique_tags)}\n"
            if unique_tags:
                stats += f"**Tags**: {', '.join(unique_tags[:10])}{'...' if len(unique_tags) > 10 else ''}\n"
            
            return stats
        except Exception as e:
            return f"❌ Error getting statistics: {str(e)}"

# Initialize the demo
print("Initializing Long Term Memory Demo...")
ltm_demo = LongTermMemoryDemo()

# Create Gradio interface
with gr.Blocks(title="Long Term Memory MCP Server Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🧠 Long Term Memory MCP Server Demo
    
    This is a demonstration of an MCP (Model Context Protocol) Server that provides long-term memory capabilities for LLM conversations.
    
    ## Features:
    - 💾 **Save Memory**: Store important insights, conclusions, or context
    - 🔍 **Search Memory**: Find relevant information using semantic search
    - 📚 **List Memories**: Browse all stored memories
    - 📊 **Statistics**: View memory usage statistics
    
    ## How it works:
    1. **Embeddings**: Uses SentenceTransformers to create semantic embeddings
    2. **Vector Storage**: ChromaDB for efficient similarity search
    3. **MCP Protocol**: Exposes tools that any MCP-compatible client can use
    """)
    
    with gr.Tabs():
        # Save Memory Tab
        with gr.Tab("💾 Save Memory"):
            gr.Markdown("### Save important insights or context to long-term memory")
            
            with gr.Row():
                with gr.Column():
                    save_title = gr.Textbox(
                        label="Title",
                        placeholder="Brief title for this memory...",
                        lines=1
                    )
                    save_content = gr.Textbox(
                        label="Content",
                        placeholder="The insight, conclusion, or context you want to remember...",
                        lines=5
                    )
                    save_tags = gr.Textbox(
                        label="Tags (optional)",
                        placeholder="quantum physics, consciousness, philosophy",
                        lines=1
                    )
                    save_context = gr.Textbox(
                        label="Context (optional)",
                        placeholder="Why is this important? When was it discussed?",
                        lines=2
                    )
                    save_btn = gr.Button("💾 Save Memory", variant="primary")
                
                with gr.Column():
                    save_output = gr.Markdown()
            
            save_btn.click(
                ltm_demo.save_memory,
                inputs=[save_content, save_title, save_tags, save_context],
                outputs=[save_output]
            )
        
        # Search Memory Tab
        with gr.Tab("🔍 Search Memory"):
            gr.Markdown("### Search through your memories using semantic similarity")
            
            with gr.Row():
                with gr.Column():
                    search_query = gr.Textbox(
                        label="Search Query",
                        placeholder="quantum consciousness, reality nature, philosophical insights...",
                        lines=2
                    )
                    with gr.Row():
                        search_limit = gr.Slider(
                            label="Max Results",
                            minimum=1,
                            maximum=20,
                            value=5,
                            step=1
                        )
                        search_threshold = gr.Slider(
                            label="Similarity Threshold",
                            minimum=0.0,
                            maximum=1.0,
                            value=0.3,
                            step=0.05
                        )
                    search_btn = gr.Button("🔍 Search Memories", variant="primary")
                
                with gr.Column():
                    search_output = gr.Markdown()
            
            search_btn.click(
                ltm_demo.search_memory,
                inputs=[search_query, search_limit, search_threshold],
                outputs=[search_output]
            )
        
        # List Memories Tab
        with gr.Tab("📚 Browse Memories"):
            gr.Markdown("### Browse all stored memories")
            
            with gr.Row():
                with gr.Column(scale=1):
                    list_limit = gr.Slider(
                        label="Number of memories to show",
                        minimum=5,
                        maximum=50,
                        value=10,
                        step=5
                    )
                    list_btn = gr.Button("📚 List Memories", variant="primary")
                    stats_btn = gr.Button("📊 Show Statistics", variant="secondary")
                
                with gr.Column(scale=3):
                    list_output = gr.Markdown()
            
            list_btn.click(
                ltm_demo.list_memories,
                inputs=[list_limit],
                outputs=[list_output]
            )
            
            stats_btn.click(
                ltm_demo.get_memory_stats,
                outputs=[list_output]
            )
    
    gr.Markdown("""
    ---
    
    ## 🔧 MCP Server Usage
    
    This Gradio app is also an MCP Server! You can connect to it from MCP-compatible clients like:
    - Claude Desktop
    - Cursor IDE
    - Other MCP clients
    
    ### Available MCP Tools:
    - `save_memory` - Save content to long-term memory
    - `search_memory` - Search through memories
    - `list_memories` - List all memories
    - `delete_memory` - Delete a specific memory
    
    ### Example Usage in Claude Desktop:
    ```
    "Save this insight to memory: 'Consciousness might be a quantum phenomenon 
    that emerges from the collapse of wave functions in microtubules.' 
    Title: 'Quantum Consciousness Theory', Tags: 'quantum, consciousness, microtubules'"
    ```
    
    Then later:
    ```
    "Search my memories for information about consciousness and quantum physics"
    ```
    """)

print("Gradio interface created successfully!")