import asyncio
import json
import logging
from datetime import datetime
from typing import Any, Dict, List, Optional, Sequence
from pathlib import Path
import uuid

import chromadb
from chromadb.config import Settings
import numpy as np
from sentence_transformers import SentenceTransformer

# MCP imports
try:
    from mcp.server import Server
    from mcp.server.models import InitializationOptions
    from mcp.server.stdio import stdio_server
    from mcp.types import (
        Resource,
        Tool,
        TextContent,
        ImageContent,
        EmbeddedResource,
        LoggingLevel
    )
    MCP_AVAILABLE = True
except ImportError:
    print("MCP library not available, running in demo mode only")
    MCP_AVAILABLE = False
    
    # Mock classes for demo mode
    class Server:
        def __init__(self, name): 
            self.name = name
        def list_tools(self): return lambda: None
        def call_tool(self): return lambda: None
    
    class TextContent:
        def __init__(self, type, text):
            self.type = type
            self.text = text

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("ltm-mcp-server")

class LongTermMemoryServer:
    def __init__(self):
        self.server = Server("long-term-memory")
        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
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        
        self.setup_handlers()
    
    def setup_handlers(self):
        @self.server.list_tools()
        async def handle_list_tools() -> List[Tool]:
            """List available tools."""
            return [
                Tool(
                    name="save_memory",
                    description="Save important insights, conclusions, or context from conversation to long-term memory",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "content": {
                                "type": "string",
                                "description": "The content/insight to save to memory"
                            },
                            "title": {
                                "type": "string", 
                                "description": "A short title/summary for this memory"
                            },
                            "tags": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "Optional tags to categorize this memory",
                                "default": []
                            },
                            "context": {
                                "type": "string",
                                "description": "Additional context about when/why this was saved",
                                "default": ""
                            }
                        },
                        "required": ["content", "title"]
                    }
                ),
                Tool(
                    name="search_memory",
                    description="Search through long-term memory for relevant information",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Search query to find relevant memories"
                            },
                            "limit": {
                                "type": "integer",
                                "description": "Maximum number of results to return",
                                "default": 5
                            },
                            "threshold": {
                                "type": "number",
                                "description": "Similarity threshold (0-1, higher = more similar)",
                                "default": 0.3
                            }
                        },
                        "required": ["query"]
                    }
                ),
                Tool(
                    name="list_memories",
                    description="List all stored memories with basic info",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "limit": {
                                "type": "integer",
                                "description": "Maximum number of memories to return",
                                "default": 10
                            }
                        }
                    }
                ),
                Tool(
                    name="delete_memory",
                    description="Delete a specific memory by ID",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "memory_id": {
                                "type": "string",
                                "description": "The ID of the memory to delete"
                            }
                        },
                        "required": ["memory_id"]
                    }
                )
            ]

        @self.server.call_tool()
        async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> Sequence[TextContent]:
            """Handle tool calls."""
            try:
                if name == "save_memory":
                    return await self._save_memory(**arguments)
                elif name == "search_memory":
                    return await self._search_memory(**arguments)
                elif name == "list_memories":
                    return await self._list_memories(**arguments)
                elif name == "delete_memory":
                    return await self._delete_memory(**arguments)
                else:
                    raise ValueError(f"Unknown tool: {name}")
            except Exception as e:
                logger.error(f"Error in tool {name}: {e}")
                return [TextContent(type="text", text=f"Error: {str(e)}")]

    async def _save_memory(self, content: str, title: str, tags: List[str] = None, context: str = "") -> Sequence[TextContent]:
        """Save content to long-term memory."""
        if tags is None:
            tags = []
            
        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),
            "context": context,
            "content_length": len(content)
        }
        
        # Save to ChromaDB
        self.collection.add(
            documents=[content],
            embeddings=[embedding],
            metadatas=[metadata],
            ids=[memory_id]
        )
        
        logger.info(f"Saved memory: {title} (ID: {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:
            result += f"**Tags**: {', '.join(tags)}\n"
        if context:
            result += f"**Context**: {context}\n"
        result += f"**Content Preview**: {content[:200]}{'...' if len(content) > 200 else ''}"
        
        return [TextContent(type="text", text=result)]

    async def _search_memory(self, query: str, limit: int = 5, threshold: float = 0.3) -> Sequence[TextContent]:
        """Search through memories."""
        if self.collection.count() == 0:
            return [TextContent(type="text", text="No memories stored yet.")]
        
        # 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 [TextContent(type="text", text="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 [TextContent(type="text", text=response)]

    async def _list_memories(self, limit: int = 10) -> Sequence[TextContent]:
        """List all memories."""
        if self.collection.count() == 0:
            return [TextContent(type="text", text="No memories stored yet.")]
        
        # Get all memories (ChromaDB doesn't have a direct "get all" with limit)
        results = self.collection.get()
        
        if not results['documents']:
            return [TextContent(type="text", text="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 [TextContent(type="text", text=response)]

    async def _delete_memory(self, memory_id: str) -> Sequence[TextContent]:
        """Delete a memory by ID."""
        try:
            # Check if memory exists
            result = self.collection.get(ids=[memory_id])
            if not result['documents']:
                return [TextContent(type="text", text=f"❌ Memory with ID {memory_id} not found.")]
            
            # Get memory info before deletion
            metadata = result['metadatas'][0]
            title = metadata['title']
            
            # Delete from ChromaDB
            self.collection.delete(ids=[memory_id])
            
            logger.info(f"Deleted memory: {title} (ID: {memory_id})")
            
            return [TextContent(type="text", text=f"✅ Memory deleted successfully!\n\n**Title**: {title}\n**ID**: {memory_id}")]
            
        except Exception as e:
            logger.error(f"Error deleting memory {memory_id}: {e}")
            return [TextContent(type="text", text=f"❌ Error deleting memory: {str(e)}")]

    async def run(self):
        """Run the MCP server."""
        if not MCP_AVAILABLE:
            print("MCP not available, cannot run server")
            return
            
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name="long-term-memory",
                    server_version="1.0.0",
                    capabilities=self.server.get_capabilities(
                        notification_options=None,
                        experimental_capabilities=None,
                    ),
                ),
            )

# Main entry point
async def main():
    if not MCP_AVAILABLE:
        print("MCP server cannot run without MCP library")
        return
        
    server = LongTermMemoryServer()
    await server.run()

if __name__ == "__main__":
    asyncio.run(main())