#!/usr/bin/env python3
"""
Simple test of LangMem API core functionality
"""

import asyncio
import asyncpg
import httpx
import json
from datetime import datetime
from uuid import uuid4

# Configuration
OLLAMA_URL = "http://localhost:11434"
SUPABASE_DB_URL = "postgresql://postgres:CzkaYmRvc26Y@localhost:5435/postgres"
API_KEY = "langmem_api_key_2025"

async def get_embedding(text: str):
    """Generate embedding using Ollama"""
    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{OLLAMA_URL}/api/embeddings",
                json={
                    "model": "nomic-embed-text",
                    "prompt": text
                },
                timeout=30.0
            )
            response.raise_for_status()
            data = response.json()
            return data["embedding"]
    except Exception as e:
        print(f"❌ Error generating embedding: {e}")
        return None

async def test_database_connection():
    """Test database connection"""
    try:
        conn = await asyncpg.connect(SUPABASE_DB_URL)
        result = await conn.fetchval("SELECT 1")
        await conn.close()
        print(f"✅ Database connection successful: {result}")
        return True
    except Exception as e:
        print(f"❌ Database connection failed: {e}")
        return False

async def test_ollama_connection():
    """Test Ollama connection"""
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{OLLAMA_URL}/api/tags")
            if response.status_code == 200:
                data = response.json()
                models = [model["name"] for model in data.get("models", [])]
                print(f"✅ Ollama connection successful, models: {models}")
                return True
            else:
                print(f"❌ Ollama connection failed: {response.status_code}")
                return False
    except Exception as e:
        print(f"❌ Ollama connection failed: {e}")
        return False

async def test_embedding_generation():
    """Test embedding generation"""
    test_text = "This is a test memory for the LangMem system"
    print(f"🧪 Testing embedding generation for: '{test_text}'")
    
    embedding = await get_embedding(test_text)
    if embedding:
        print(f"✅ Embedding generated successfully, dimension: {len(embedding)}")
        return embedding
    else:
        print("❌ Failed to generate embedding")
        return None

async def test_vector_storage():
    """Test vector storage in Supabase"""
    try:
        conn = await asyncpg.connect(SUPABASE_DB_URL)
        
        # Create test table
        await conn.execute("""
            CREATE TABLE IF NOT EXISTS test_langmem_documents (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                content TEXT NOT NULL,
                embedding vector(768),
                user_id TEXT NOT NULL,
                created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
            );
        """)
        
        # Generate test embedding
        test_content = "FastAPI is a modern web framework for building APIs with Python"
        embedding = await get_embedding(test_content)
        
        if not embedding:
            print("❌ Cannot test vector storage without embedding")
            return False
        
        # Store test document
        doc_id = uuid4()
        await conn.execute("""
            INSERT INTO test_langmem_documents (id, content, embedding, user_id)
            VALUES ($1, $2, $3, $4)
        """, doc_id, test_content, str(embedding), "test_user")
        
        # Test similarity search
        query_embedding = await get_embedding("Python web framework")
        if query_embedding:
            results = await conn.fetch("""
                SELECT id, content, 1 - (embedding <=> $1) as similarity
                FROM test_langmem_documents
                WHERE user_id = 'test_user'
                ORDER BY embedding <=> $1
                LIMIT 5
            """, str(query_embedding))
            
            print(f"✅ Vector storage and similarity search successful")
            for row in results:
                print(f"   - Content: {row['content'][:50]}...")
                print(f"   - Similarity: {row['similarity']:.3f}")
        
        # Cleanup
        await conn.execute("DROP TABLE IF EXISTS test_langmem_documents")
        await conn.close()
        return True
        
    except Exception as e:
        print(f"❌ Vector storage test failed: {e}")
        return False

async def main():
    """Run all tests"""
    print("🚀 Starting LangMem API Tests")
    print("=" * 50)
    
    # Test 1: Database connection
    print("\n1. Testing Database Connection...")
    db_ok = await test_database_connection()
    
    # Test 2: Ollama connection
    print("\n2. Testing Ollama Connection...")
    ollama_ok = await test_ollama_connection()
    
    # Test 3: Embedding generation
    print("\n3. Testing Embedding Generation...")
    embedding_ok = await test_embedding_generation()
    
    # Test 4: Vector storage
    print("\n4. Testing Vector Storage...")
    vector_ok = await test_vector_storage()
    
    # Summary
    print("\n" + "=" * 50)
    print("📊 Test Results Summary:")
    print(f"   Database Connection: {'✅ PASS' if db_ok else '❌ FAIL'}")
    print(f"   Ollama Connection: {'✅ PASS' if ollama_ok else '❌ FAIL'}")
    print(f"   Embedding Generation: {'✅ PASS' if embedding_ok else '❌ FAIL'}")
    print(f"   Vector Storage: {'✅ PASS' if vector_ok else '❌ FAIL'}")
    
    all_passed = all([db_ok, ollama_ok, embedding_ok, vector_ok])
    print(f"\n🎯 Overall Status: {'✅ ALL TESTS PASSED' if all_passed else '❌ SOME TESTS FAILED'}")
    
    if all_passed:
        print("\n🎉 LangMem API core functionality is working!")
        print("   Ready to proceed with full API deployment.")
    else:
        print("\n⚠️  Some tests failed. Please check the configuration.")

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