#!/usr/bin/env python3
"""
mem0 Memory Server - Persistent Semantic Memory for Hermes Agent
Direct integration with llama-embed on port 4700
"""

import os
import json
import requests
from http.server import HTTPServer, BaseHTTPRequestHandler
from qdrant_client import QdrantClient, models

# Configuration
QDRANT_URL = os.environ.get("QDRANT_URL", "http://localhost:6333")
EMBEDDING_URL = os.environ.get("EMBEDDING_URL", "http://localhost:4700")
PORT = int(os.environ.get("MEM0_PORT", 8080))
USER_ID = "henry_hofmann"

# Initialize Qdrant client
qdrant_client = QdrantClient(url=QDRANT_URL)

# Create collection if it doesn't exist
try:
    qdrant_client.get_collection("hermes_memory")
except:
    qdrant_client.create_collection(
        collection_name="hermes_memory",
        vectors_config=models.VectorParams(size=1024, distance=models.Distance.COSINE)
    )

def get_embedding(text):
    """Get embedding from llama-embed server"""
    response = requests.post(
        f"{EMBEDDING_URL}/v1/embeddings",
        json={"input": text, "model": "BAAI/bge-m3"},
        timeout=30
    )
    response.raise_for_status()
    data = response.json()
    return data["data"][0]["embedding"]

class MemoryHandler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        pass  # Suppress logging
    
    def do_GET(self):
        if self.path == "/health":
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({"status": "ok", "service": "mem0", "user": USER_ID}).encode())
        elif self.path == "/memory":
            # Get recent memories for user
            try:
                records = qdrant_client.scroll(
                    collection_name="hermes_memory",
                    limit=10,
                    with_payload=True,
                    with_vectors=False
                )
                memories = []
                for record in records[0]:
                    if hasattr(record, 'payload'):
                        memories.append({
                            "id": record.id,
                            "text": record.payload.get("text", ""),
                            "timestamp": record.payload.get("timestamp", "")
                        })
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps(memories, default=str).encode())
            except Exception as e:
                self.send_response(500)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps({"error": str(e)}).encode())
        elif self.path.startswith("/memory/") and self.path.endswith("/search"):
            # Search memories by query
            query = self.path.split("/")[2]
            try:
                query_vector = get_embedding(query)
                results = qdrant_client.query_points(
                    collection_name="hermes_memory",
                    query=query_vector,
                    query_filter=models.Filter(
                        must=[models.FieldCondition(key="user_id", match=models.MatchValue(value=USER_ID))]
                    ),
                    limit=5,
                    with_payload=True
                )
                memories = []
                for result in results.points:
                    if hasattr(result, 'payload'):
                        memories.append({
                            "id": result.id,
                            "text": result.payload.get("text", ""),
                            "score": result.score,
                            "timestamp": result.payload.get("timestamp", "")
                        })
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps(memories, default=str).encode())
            except Exception as e:
                self.send_response(500)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps({"error": str(e)}).encode())
        else:
            self.send_response(404)
            self.end_headers()
    
    def do_POST(self):
        if self.path == "/memory":
            content_length = int(self.headers["Content-Length"])
            post_data = json.loads(self.rfile.read(content_length).decode())
            text = post_data.get("text", "")
            
            if text:
                try:
                    # Get embedding
                    embedding = get_embedding(text)
                    
                    # Store in Qdrant
                    qdrant_client.upsert(
                        collection_name="hermes_memory",
                        points=[
                            models.PointStruct(
                                id=hash(text) % 1000000,
                                vector=embedding,
                                payload={
                                    "text": text,
                                    "user_id": USER_ID,
                                    "timestamp": str(os.popen("date -Iseconds").read().strip())
                                }
                            )
                        ]
                    )
                    
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.end_headers()
                    self.wfile.write(json.dumps({"status": "ok", "text": text}).encode())
                except Exception as e:
                    self.send_response(500)
                    self.send_header("Content-Type", "application/json")
                    self.end_headers()
                    self.wfile.write(json.dumps({"error": str(e)}).encode())
            else:
                self.send_response(400)
                self.end_headers()
        else:
            self.send_response(404)
            self.end_headers()

if __name__ == "__main__":
    server = HTTPServer(("0.0.0.0", PORT), MemoryHandler)
    print(f"mem0 server running on port {PORT}")
    print(f"Qdrant: {QDRANT_URL}")
    print(f"Embedding: {EMBEDDING_URL}")
    server.serve_forever()