enhanced_database_manager.py

import asyncio
import os
import json
import logging
import numpy as np
import pickle
import gzip
import asyncpg
from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime
import uuid
import base64

class EnhancedDatabaseManager:
    """Enhanced Database Manager that stores everything in PostgreSQL + Vercel Blob"""
    
    def __init__(self, database_url: str):
        self.database_url = database_url
        self.pool = None
        self.logger = logging.getLogger(__name__)
        
    async def connect(self):
        """Initialize database connection pool"""
        try:
            self.pool = await asyncpg.create_pool(
                self.database_url,
                min_size=2,
                max_size=20,
                command_timeout=60
            )
            self.logger.info("Enhanced database connection pool created successfully")
            
            # Create all necessary tables
            await self._create_all_tables()
            
        except Exception as e:
            self.logger.error(f"Database connection failed: {e}")
            raise

    async def _create_all_tables(self):
        """Create all tables for comprehensive storage"""
        async with self.pool.acquire() as conn:
            await conn.execute("""
                -- RAG instances metadata
                CREATE TABLE IF NOT EXISTS rag_instances (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    ai_type VARCHAR(50) NOT NULL,
                    user_id VARCHAR(100),
                    ai_id VARCHAR(100),
                    name VARCHAR(255) NOT NULL,
                    description TEXT,
                    
                    -- Storage references
                    blob_url TEXT,
                    config_json JSONB,
                    
                    -- Statistics
                    total_chunks INTEGER DEFAULT 0,
                    total_tokens INTEGER DEFAULT 0,
                    file_count INTEGER DEFAULT 0,
                    
                    -- Timestamps
                    created_at TIMESTAMP DEFAULT NOW(),
                    updated_at TIMESTAMP DEFAULT NOW(),
                    last_accessed_at TIMESTAMP DEFAULT NOW(),
                    
                    -- Status
                    status VARCHAR(20) DEFAULT 'active',
                    
                    UNIQUE(ai_type, user_id, ai_id)
                );
                
                -- Knowledge files metadata
                CREATE TABLE IF NOT EXISTS knowledge_files (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    rag_instance_id UUID REFERENCES rag_instances(id) ON DELETE CASCADE,
                    filename VARCHAR(255) NOT NULL,
                    original_filename VARCHAR(255),
                    file_type VARCHAR(50),
                    file_size INTEGER,
                    
                    -- Content storage
                    content_text TEXT,
                    content_blob BYTEA,
                    
                    -- Processing info
                    processed_at TIMESTAMP DEFAULT NOW(),
                    processing_status VARCHAR(20) DEFAULT 'pending',
                    token_count INTEGER DEFAULT 0,
                    
                    -- Timestamps
                    created_at TIMESTAMP DEFAULT NOW(),
                    updated_at TIMESTAMP DEFAULT NOW()
                );
                
                -- RAG graph data (for large graphs, store in chunks)
                CREATE TABLE IF NOT EXISTS rag_graph_data (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    rag_instance_id UUID REFERENCES rag_instances(id) ON DELETE CASCADE,
                    data_type VARCHAR(20) NOT NULL, -- 'nodes', 'edges', 'attrs'
                    chunk_index INTEGER DEFAULT 0,
                    chunk_data JSONB,
                    created_at TIMESTAMP DEFAULT NOW()
                );
                
                -- RAG vector data (for large embeddings, store in chunks)
                CREATE TABLE IF NOT EXISTS rag_vector_data (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    rag_instance_id UUID REFERENCES rag_instances(id) ON DELETE CASCADE,
                    data_type VARCHAR(20) NOT NULL, -- 'embeddings', 'metadata'
                    chunk_index INTEGER DEFAULT 0,
                    chunk_data JSONB,
                    created_at TIMESTAMP DEFAULT NOW()
                );
                
                -- User conversations
                CREATE TABLE IF NOT EXISTS conversations (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    user_id VARCHAR(100) NOT NULL,
                    rag_instance_id UUID REFERENCES rag_instances(id) ON DELETE CASCADE,
                    title VARCHAR(255),
                    created_at TIMESTAMP DEFAULT NOW(),
                    updated_at TIMESTAMP DEFAULT NOW(),
                    is_active BOOLEAN DEFAULT TRUE
                );
                
                -- Conversation messages
                CREATE TABLE IF NOT EXISTS conversation_messages (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    conversation_id UUID REFERENCES conversations(id) ON DELETE CASCADE,
                    role VARCHAR(20) NOT NULL, -- 'user', 'assistant'
                    content TEXT NOT NULL,
                    metadata JSONB DEFAULT '{}',
                    created_at TIMESTAMP DEFAULT NOW()
                );
                
                -- System statistics
                CREATE TABLE IF NOT EXISTS system_stats (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    stat_date DATE DEFAULT CURRENT_DATE,
                    total_rag_instances INTEGER DEFAULT 0,
                    total_conversations INTEGER DEFAULT 0,
                    total_messages INTEGER DEFAULT 0,
                    total_knowledge_files INTEGER DEFAULT 0,
                    created_at TIMESTAMP DEFAULT NOW(),
                    UNIQUE(stat_date)
                );
                
                -- Create indexes for performance
                CREATE INDEX IF NOT EXISTS idx_rag_instances_lookup ON rag_instances(ai_type, user_id, ai_id);
                CREATE INDEX IF NOT EXISTS idx_rag_instances_status ON rag_instances(status);
                CREATE INDEX IF NOT EXISTS idx_rag_instances_user ON rag_instances(user_id);
                CREATE INDEX IF NOT EXISTS idx_knowledge_files_rag ON knowledge_files(rag_instance_id);
                CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id);
                CREATE INDEX IF NOT EXISTS idx_conversation_messages_conv ON conversation_messages(conversation_id);
                CREATE INDEX IF NOT EXISTS idx_rag_graph_data_rag ON rag_graph_data(rag_instance_id);
                CREATE INDEX IF NOT EXISTS idx_rag_vector_data_rag ON rag_vector_data(rag_instance_id);
            """)
            
            self.logger.info("Enhanced database tables created/verified successfully")

    async def save_complete_rag_instance(
        self,
        ai_type: str,
        user_id: Optional[str],
        ai_id: Optional[str],
        name: str,
        description: Optional[str],
        rag_state: Dict[str, Any],
        blob_url: Optional[str] = None
    ) -> str:
        """Save complete RAG instance with all data to database"""
        
        async with self.pool.acquire() as conn:
            async with conn.transaction():
                # Save main RAG instance
                rag_instance_id = await conn.fetchval("""
                    INSERT INTO rag_instances (
                        ai_type, user_id, ai_id, name, description, blob_url,
                        config_json, total_chunks, total_tokens, file_count
                    ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
                    ON CONFLICT (ai_type, user_id, ai_id) DO UPDATE SET
                        name = EXCLUDED.name,
                        description = EXCLUDED.description,
                        blob_url = EXCLUDED.blob_url,
                        config_json = EXCLUDED.config_json,
                        total_chunks = EXCLUDED.total_chunks,
                        total_tokens = EXCLUDED.total_tokens,
                        file_count = EXCLUDED.file_count,
                        updated_at = NOW()
                    RETURNING id;
                """, 
                    ai_type, user_id, ai_id, name, description, blob_url,
                    json.dumps(rag_state.get('config', {})),
                    len(rag_state.get('vectors', {}).get('embeddings', [])),
                    self._estimate_tokens(rag_state),
                    0
                )
                
                # Clear existing graph and vector data
                await conn.execute("""
                    DELETE FROM rag_graph_data WHERE rag_instance_id = $1
                """, rag_instance_id)
                
                await conn.execute("""
                    DELETE FROM rag_vector_data WHERE rag_instance_id = $1
                """, rag_instance_id)
                
                # Save graph data in chunks
                graph_data = rag_state.get('graph', {})
                await self._save_graph_data(conn, rag_instance_id, graph_data)
                
                # Save vector data in chunks
                vector_data = rag_state.get('vectors', {})
                await self._save_vector_data(conn, rag_instance_id, vector_data)
                
                return str(rag_instance_id)

    async def _save_graph_data(self, conn, rag_instance_id: str, graph_data: Dict[str, Any]):
        """Save graph data in chunks to avoid size limits"""
        
        # Save nodes in chunks
        nodes = graph_data.get('nodes', [])
        if nodes:
            chunk_size = 1000  # Adjust based on your needs
            for i in range(0, len(nodes), chunk_size):
                chunk = nodes[i:i + chunk_size]
                await conn.execute("""
                    INSERT INTO rag_graph_data (rag_instance_id, data_type, chunk_index, chunk_data)
                    VALUES ($1, $2, $3, $4)
                """, rag_instance_id, 'nodes', i // chunk_size, json.dumps(chunk))
        
        # Save edges in chunks
        edges = graph_data.get('edges', [])
        if edges:
            chunk_size = 1000
            for i in range(0, len(edges), chunk_size):
                chunk = edges[i:i + chunk_size]
                await conn.execute("""
                    INSERT INTO rag_graph_data (rag_instance_id, data_type, chunk_index, chunk_data)
                    VALUES ($1, $2, $3, $4)
                """, rag_instance_id, 'edges', i // chunk_size, json.dumps(chunk))
        
        # Save graph attributes
        graph_attrs = graph_data.get('graph_attrs', {})
        if graph_attrs:
            await conn.execute("""
                INSERT INTO rag_graph_data (rag_instance_id, data_type, chunk_index, chunk_data)
                VALUES ($1, $2, $3, $4)
            """, rag_instance_id, 'attrs', 0, json.dumps(graph_attrs))

    async def _save_vector_data(self, conn, rag_instance_id: str, vector_data: Dict[str, Any]):
        """Save vector data in chunks to avoid size limits"""
        
        # Save embeddings in chunks
        embeddings = vector_data.get('embeddings', [])
        if embeddings:
            chunk_size = 100  # Smaller chunks for embeddings
            for i in range(0, len(embeddings), chunk_size):
                chunk = embeddings[i:i + chunk_size]
                await conn.execute("""
                    INSERT INTO rag_vector_data (rag_instance_id, data_type, chunk_index, chunk_data)
                    VALUES ($1, $2, $3, $4)
                """, rag_instance_id, 'embeddings', i // chunk_size, json.dumps(chunk))
        
        # Save metadata
        metadata = vector_data.get('metadata', [])
        if metadata:
            await conn.execute("""
                INSERT INTO rag_vector_data (rag_instance_id, data_type, chunk_index, chunk_data)
                VALUES ($1, $2, $3, $4)
            """, rag_instance_id, 'metadata', 0, json.dumps(metadata))

    async def load_complete_rag_instance(
        self,
        ai_type: str,
        user_id: Optional[str] = None,
        ai_id: Optional[str] = None
    ) -> Optional[Dict[str, Any]]:
        """Load complete RAG instance from database"""
        
        async with self.pool.acquire() as conn:
            # Get main RAG instance
            rag_instance = await conn.fetchrow("""
                SELECT id, ai_type, user_id, ai_id, name, description, blob_url,
                       config_json, total_chunks, total_tokens, file_count,
                       created_at, updated_at, last_accessed_at, status
                FROM rag_instances
                WHERE ai_type = $1 AND user_id = $2 AND ai_id = $3 AND status = 'active'
            """, ai_type, user_id, ai_id)
            
            if not rag_instance:
                return None
            
            # Update last accessed time
            await conn.execute("""
                UPDATE rag_instances SET last_accessed_at = NOW() WHERE id = $1
            """, rag_instance['id'])
            
            # Load graph data
            graph_data = await self._load_graph_data(conn, rag_instance['id'])
            
            # Load vector data
            vector_data = await self._load_vector_data(conn, rag_instance['id'])
            
            return {
                "metadata": dict(rag_instance),
                "rag_state": {
                    "graph": graph_data,
                    "vectors": vector_data,
                    "config": rag_instance['config_json'] or {},
                    "version": "1.0"
                }
            }

    async def _load_graph_data(self, conn, rag_instance_id: str) -> Dict[str, Any]:
        """Load graph data from chunks"""
        
        # Load nodes
        nodes_chunks = await conn.fetch("""
            SELECT chunk_index, chunk_data FROM rag_graph_data
            WHERE rag_instance_id = $1 AND data_type = 'nodes'
            ORDER BY chunk_index
        """, rag_instance_id)
        
        nodes = []
        for chunk_row in nodes_chunks:
            nodes.extend(chunk_row['chunk_data'])
        
        # Load edges
        edges_chunks = await conn.fetch("""
            SELECT chunk_index, chunk_data FROM rag_graph_data
            WHERE rag_instance_id = $1 AND data_type = 'edges'
            ORDER BY chunk_index
        """, rag_instance_id)
        
        edges = []
        for chunk_row in edges_chunks:
            edges.extend(chunk_row['chunk_data'])
        
        # Load graph attributes
        attrs_row = await conn.fetchrow("""
            SELECT chunk_data FROM rag_graph_data
            WHERE rag_instance_id = $1 AND data_type = 'attrs'
        """, rag_instance_id)
        
        graph_attrs = attrs_row['chunk_data'] if attrs_row else {}
        
        return {
            "nodes": nodes,
            "edges": edges,
            "graph_attrs": graph_attrs
        }

    async def _load_vector_data(self, conn, rag_instance_id: str) -> Dict[str, Any]:
        """Load vector data from chunks"""
        
        # Load embeddings
        embeddings_chunks = await conn.fetch("""
            SELECT chunk_index, chunk_data FROM rag_vector_data
            WHERE rag_instance_id = $1 AND data_type = 'embeddings'
            ORDER BY chunk_index
        """, rag_instance_id)
        
        embeddings = []
        for chunk_row in embeddings_chunks:
            embeddings.extend(chunk_row['chunk_data'])
        
        # Load metadata
        metadata_row = await conn.fetchrow("""
            SELECT chunk_data FROM rag_vector_data
            WHERE rag_instance_id = $1 AND data_type = 'metadata'
        """, rag_instance_id)
        
        metadata = metadata_row['chunk_data'] if metadata_row else []
        
        return {
            "embeddings": embeddings,
            "metadata": metadata,
            "dimension": 1024
        }

    async def save_knowledge_file(
        self,
        rag_instance_id: str,
        filename: str,
        original_filename: str,
        file_type: str,
        file_size: int,
        content_text: str,
        content_blob: Optional[bytes] = None
    ) -> str:
        """Save knowledge file to database"""
        
        async with self.pool.acquire() as conn:
            file_id = await conn.fetchval("""
                INSERT INTO knowledge_files (
                    rag_instance_id, filename, original_filename, file_type, 
                    file_size, content_text, content_blob, processing_status
                ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                RETURNING id
            """, rag_instance_id, filename, original_filename, file_type, 
                file_size, content_text, content_blob, 'processed')
            
            return str(file_id)

    async def get_knowledge_files(self, rag_instance_id: str) -> List[Dict[str, Any]]:
        """Get all knowledge files for a RAG instance"""
        
        async with self.pool.acquire() as conn:
            files = await conn.fetch("""
                SELECT id, filename, original_filename, file_type, file_size,
                       content_text, processing_status, token_count,
                       created_at, updated_at
                FROM knowledge_files
                WHERE rag_instance_id = $1
                ORDER BY created_at DESC
            """, rag_instance_id)
            
            return [dict(file) for file in files]

    async def list_user_rag_instances(self, user_id: str) -> List[Dict[str, Any]]:
        """List all RAG instances for a user"""
        async with self.pool.acquire() as conn:
            results = await conn.fetch("""
                SELECT id, ai_type, ai_id, name, description, total_chunks,
                       total_tokens, file_count, created_at, updated_at,
                       last_accessed_at, status
                FROM rag_instances
                WHERE user_id = $1 AND status = 'active'
                ORDER BY created_at DESC
            """, user_id)
            
            return [dict(row) for row in results]

    async def save_conversation(
        self,
        user_id: str,
        rag_instance_id: str,
        title: Optional[str] = None
    ) -> str:
        """Save conversation to database"""
        
        async with self.pool.acquire() as conn:
            conversation_id = await conn.fetchval("""
                INSERT INTO conversations (user_id, rag_instance_id, title)
                VALUES ($1, $2, $3)
                RETURNING id
            """, user_id, rag_instance_id, title)
            
            return str(conversation_id)

    async def save_conversation_message(
        self,
        conversation_id: str,
        role: str,
        content: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """Save conversation message to database"""
        
        async with self.pool.acquire() as conn:
            message_id = await conn.fetchval("""
                INSERT INTO conversation_messages (conversation_id, role, content, metadata)
                VALUES ($1, $2, $3, $4)
                RETURNING id
            """, conversation_id, role, content, json.dumps(metadata or {}))
            
            # Update conversation timestamp
            await conn.execute("""
                UPDATE conversations SET updated_at = NOW() WHERE id = $1
            """, conversation_id)
            
            return str(message_id)

    async def get_conversation_messages(
        self,
        conversation_id: str,
        limit: int = 50
    ) -> List[Dict[str, Any]]:
        """Get conversation messages from database"""
        
        async with self.pool.acquire() as conn:
            messages = await conn.fetch("""
                SELECT id, role, content, metadata, created_at
                FROM conversation_messages
                WHERE conversation_id = $1
                ORDER BY created_at DESC
                LIMIT $2
            """, conversation_id, limit)
            
            return [dict(msg) for msg in reversed(messages)]

    async def get_user_conversations(self, user_id: str) -> List[Dict[str, Any]]:
        """Get all conversations for a user"""
        
        async with self.pool.acquire() as conn:
            conversations = await conn.fetch("""
                SELECT c.id, c.title, c.created_at, c.updated_at,
                       r.name as ai_name, r.ai_type,
                       (SELECT content FROM conversation_messages 
                        WHERE conversation_id = c.id 
                        ORDER BY created_at DESC LIMIT 1) as last_message
                FROM conversations c
                JOIN rag_instances r ON c.rag_instance_id = r.id
                WHERE c.user_id = $1 AND c.is_active = TRUE
                ORDER BY c.updated_at DESC
            """, user_id)
            
            return [dict(conv) for conv in conversations]

    async def update_system_stats(self):
        """Update system statistics"""
        
        async with self.pool.acquire() as conn:
            # Get current counts
            stats = await conn.fetchrow("""
                SELECT 
                    (SELECT COUNT(*) FROM rag_instances WHERE status = 'active') as rag_count,
                    (SELECT COUNT(*) FROM conversations WHERE is_active = TRUE) as conv_count,
                    (SELECT COUNT(*) FROM conversation_messages) as msg_count,
                    (SELECT COUNT(*) FROM knowledge_files) as file_count
            """)
            
            # Update stats for today
            await conn.execute("""
                INSERT INTO system_stats (
                    stat_date, total_rag_instances, total_conversations, 
                    total_messages, total_knowledge_files
                ) VALUES (CURRENT_DATE, $1, $2, $3, $4)
                ON CONFLICT (stat_date) DO UPDATE SET
                    total_rag_instances = EXCLUDED.total_rag_instances,
                    total_conversations = EXCLUDED.total_conversations,
                    total_messages = EXCLUDED.total_messages,
                    total_knowledge_files = EXCLUDED.total_knowledge_files
            """, stats['rag_count'], stats['conv_count'], stats['msg_count'], stats['file_count'])

    async def get_system_stats(self) -> Dict[str, Any]:
        """Get system statistics"""
        
        async with self.pool.acquire() as conn:
            stats = await conn.fetchrow("""
                SELECT * FROM system_stats
                ORDER BY stat_date DESC
                LIMIT 1
            """)
            
            return dict(stats) if stats else {}

    async def delete_rag_instance(self, rag_instance_id: str):
        """Soft delete a RAG instance"""
        
        async with self.pool.acquire() as conn:
            await conn.execute("""
                UPDATE rag_instances 
                SET status = 'deleted', updated_at = NOW() 
                WHERE id = $1
            """, rag_instance_id)

    async def cleanup_old_data(self, days_old: int = 30):
        """Clean up old data from database"""
        
        async with self.pool.acquire() as conn:
            # Clean up old deleted RAG instances
            await conn.execute("""
                DELETE FROM rag_instances 
                WHERE status = 'deleted' AND updated_at < NOW() - INTERVAL '%s days'
            """, days_old)
            
            # Clean up old system stats (keep last 90 days)
            await conn.execute("""
                DELETE FROM system_stats 
                WHERE stat_date < CURRENT_DATE - INTERVAL '90 days'
            """)

    def _estimate_tokens(self, rag_state: Dict[str, Any]) -> int:
        """Estimate token count from RAG state"""
        try:
            # Simple estimation based on serialized size
            content_size = len(json.dumps(rag_state))
            return content_size // 4  # Rough estimate: 4 chars per token
        except:
            return 0

    async def get_database_size(self) -> Dict[str, Any]:
        """Get database size information"""
        
        async with self.pool.acquire() as conn:
            size_info = await conn.fetchrow("""
                SELECT 
                    pg_size_pretty(pg_database_size(current_database())) as total_size,
                    (SELECT COUNT(*) FROM rag_instances) as rag_instances,
                    (SELECT COUNT(*) FROM knowledge_files) as knowledge_files,
                    (SELECT COUNT(*) FROM conversations) as conversations,
                    (SELECT COUNT(*) FROM conversation_messages) as messages,
                    (SELECT COUNT(*) FROM rag_graph_data) as graph_chunks,
                    (SELECT COUNT(*) FROM rag_vector_data) as vector_chunks
            """)
            
            return dict(size_info)

    async def test_connection(self) -> bool:
        """Test database connection"""
        try:
            async with self.pool.acquire() as conn:
                await conn.fetchval("SELECT 1")
            return True
        except Exception as e:
            self.logger.error(f"Database connection test failed: {e}")
            return False

    async def close(self):
        """Close database connection pool"""
        if self.pool:
            await self.pool.close()
            self.logger.info("Database connection pool closed")