vi/services/memory/memory_service.py

"""
Memory Service - Three-layer memory system with modular architecture.

Refactored into storage backends, operations, and handlers for maintainability.
"""
import asyncio
from pathlib import Path
from typing import Dict, Any

from core.config import SHORT_TERM_DB
from core.logger import setup_logger
from core.nats_event_bus import nats_bus as event_bus
from core.base_service import BaseService
from core.service_registry import ServiceManifest

# Import refactored components
from .storage.sqlite_store import SQLiteStore
from .storage.chroma_store import ChromaStore
from .storage.migrations import archive_old_database
from .operations.short_term_ops import ShortTermOperations
from .operations.long_term_ops import LongTermOperations
from .operations.facts_ops import FactsOperations
from .handlers.store_handler import StoreHandler
from .handlers.search_handler import SearchHandler
from .handlers.reset_handler import ResetHandler
from .handlers.short_memory_handler import ShortMemoryHandler
from .handlers.long_memory_handler import LongMemoryHandler
from .handlers.facts_handler import FactsHandler
from .handlers.save_fact_handler import SaveFactHandler
from .handlers.update_fact_handler import UpdateFactHandler
from .utils.embeddings import get_model

logger = setup_logger('memory_service', service_name='memory_service')


class MemoryService(BaseService):
    """Memory service with three-layer architecture and modular design"""

    def __init__(self):
        super().__init__('memory')

        # Initialize storage backends
        self.sqlite_store = SQLiteStore()
        self.chroma_store = ChromaStore()

        # Initialize operations (will be set up after storage connects)
        self.short_term_ops = None
        self.long_term_ops = None
        self.facts_ops = None

        # Initialize handlers (will be set up after operations)
        self.store_handler = None
        self.search_handler = None
        self.reset_handler = None
        self.short_memory_handler = None
        self.long_memory_handler = None
        self.facts_handler = None
        self.save_fact_handler = None
        self.update_fact_handler = None

        logger.info("[μ] Memory Service initialized with modular architecture")

    def get_service_manifest(self) -> ServiceManifest:
        """Return service manifest with operations and metadata"""
        operations = [
            # Legacy operations (backward compatibility)
            self.create_service_operation(
                "store",
                "Store a memory (routes to short-term)",
                timeout_ms=5000
            ),
            self.create_service_operation(
                "search",
                "Search memories (legacy, redirects to short_memory)",
                timeout_ms=3000
            ),
            self.create_service_operation(
                "reset",
                "Reset/clear memory database for debugging",
                timeout_ms=10000
            ),
            # New three-layer memory operations
            self.create_service_operation(
                "short_memory",
                "Get recent literal memories with offset support",
                timeout_ms=3000
            ),
            self.create_service_operation(
                "long_memory",
                "Semantic search in long-term summarized memories",
                timeout_ms=5000
            ),
            self.create_service_operation(
                "facts",
                "Search factual memory by category or semantic query",
                timeout_ms=3000
            ),
            self.create_service_operation(
                "save_fact",
                "Store a new fact in factual memory",
                timeout_ms=2000
            ),
            self.create_service_operation(
                "update_fact",
                "Update an existing fact (if mutable)",
                timeout_ms=2000
            )
        ]

        return ServiceManifest(
            service_id=self.service_id,
            name="Memory Service",
            description="Three-layer memory system: short-term (literal), long-term (summarized), factual (exact)",
            version="3.0.0",
            operations=operations,
            dependencies=[],
            health_check_topic=f"vi.services.{self.service_id}.health",
            heartbeat_interval=30,
            metadata={
                "storage_type": "hybrid",
                "short_term_storage": "sqlite",
                "long_term_storage": "chromadb",
                "facts_storage": "chromadb",
                "embedding_model": "all-MiniLM-L6-v2",
                "vector_search": True,
                "urgency": 0.8
            }
        )

    async def initialize_service(self):
        """Initialize service-specific resources and register handlers"""
        # Archive old database if it exists (one-time migration)
        archive_old_database(Path(SHORT_TERM_DB))

        # Initialize storage backends
        self.sqlite_store.connect()
        self.chroma_store.connect()

        # Initialize embedding model
        get_model()  # Loads model on first call

        # Initialize operations
        self.short_term_ops = ShortTermOperations(self.sqlite_store)
        self.long_term_ops = LongTermOperations(self.chroma_store)
        self.facts_ops = FactsOperations(self.chroma_store)

        # Initialize handlers
        self.store_handler = StoreHandler(self.sqlite_store)
        self.search_handler = SearchHandler(self.short_term_ops)
        self.reset_handler = ResetHandler(self.sqlite_store, self.chroma_store)
        self.short_memory_handler = ShortMemoryHandler(self.short_term_ops)
        self.long_memory_handler = LongMemoryHandler(self.long_term_ops)
        self.facts_handler = FactsHandler(self.facts_ops)
        self.save_fact_handler = SaveFactHandler(self.facts_ops)
        self.update_fact_handler = UpdateFactHandler(self.facts_ops)

        # Register handlers using new topic patterns
        await self.register_handler("store", self.store_handler.handle)
        await self.register_handler("search", self.search_handler.handle)
        await self.register_handler("reset", self.reset_handler.handle)
        await self.register_handler("short_memory", self.short_memory_handler.handle)
        await self.register_handler("long_memory", self.long_memory_handler.handle)
        await self.register_handler("facts", self.facts_handler.handle)
        await self.register_handler("save_fact", self.save_fact_handler.handle)
        await self.register_handler("update_fact", self.update_fact_handler.handle)

        # Also register legacy topic handlers for backward compatibility
        await self.event_bus.on("vi.memory.store", self.store_handler.handle)
        await self.event_bus.on("vi.memory.search", self.search_handler.handle)
        await self.event_bus.on("vi.memory.debug.reset", self.reset_handler.handle)

        self.logger.info("[μ] Memory Service initialized with three-layer memory system")

    async def cleanup_service(self):
        """Cleanup service-specific resources"""
        # Unregister event handlers
        await self.event_bus.off("vi.memory.store")
        await self.event_bus.off("vi.memory.search")
        await self.event_bus.off("vi.memory.debug.reset")

        # Close storage connections
        if self.sqlite_store:
            self.sqlite_store.close()

        self.logger.info("[μ] Memory Service cleanup completed")

    async def perform_health_check(self) -> Dict[str, Any]:
        """Perform service-specific health check"""
        health_data = {
            'healthy': True,
            'checks': {
                'running': self._running,
                'event_bus': self.event_bus is not None,
                'database_connected': self.sqlite_store.conn is not None,
                'embedding_model': get_model() is not None
            }
        }

        # Check database connectivity
        try:
            if self.sqlite_store.conn:
                cursor = self.sqlite_store.conn.cursor()
                cursor.execute("SELECT COUNT(*) FROM short_term_memory")
                short_term_count = cursor.fetchone()[0]
                health_data['checks']['short_term_count'] = short_term_count
                health_data['checks']['long_term_count'] = self.chroma_store.get_long_term_collection().count()
                health_data['checks']['facts_count'] = self.chroma_store.get_facts_collection().count()
                health_data['checks']['database_accessible'] = True
            else:
                health_data['checks']['database_accessible'] = False
                health_data['healthy'] = False
        except Exception as e:
            health_data['checks']['database_accessible'] = False
            health_data['checks']['database_error'] = str(e)
            health_data['healthy'] = False

        return health_data


async def main():
    """Main entry point for memory service"""
    memory_service = MemoryService()

    try:
        await event_bus.connect()
        await memory_service.start(event_bus)

        logger.info("[μ] Memory service running. Press Ctrl+C to stop.")

        # Keep running
        while True:
            await asyncio.sleep(1)

    except KeyboardInterrupt:
        logger.info("[μ] Shutdown requested")
    except Exception as e:
        logger.exception(f"[μ] Unexpected error: {e}")
    finally:
        await memory_service.stop()
        await event_bus.close()


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