#!/usr/bin/env python3
"""
Test suite for LangMem API
"""

import asyncio
import json
import pytest
import httpx
from uuid import uuid4

# Configuration
API_BASE_URL = "http://localhost:8765"
API_KEY = "langmem_api_key_2025"

class TestLangMemAPI:
    """Test suite for LangMem API endpoints"""
    
    def setup_method(self):
        """Setup test client"""
        self.client = httpx.AsyncClient(base_url=API_BASE_URL)
        self.headers = {"Authorization": f"Bearer {API_KEY}"}
        self.test_user_id = f"test_user_{uuid4()}"
    
    async def teardown_method(self):
        """Cleanup test client"""
        await self.client.aclose()
    
    @pytest.mark.asyncio
    async def test_root_endpoint(self):
        """Test root endpoint"""
        response = await self.client.get("/")
        assert response.status_code == 200
        data = response.json()
        assert data["message"] == "LangMem API - Long-term Memory System"
        assert data["version"] == "1.0.0"
    
    @pytest.mark.asyncio
    async def test_health_check(self):
        """Test health check endpoint"""
        response = await self.client.get("/health")
        assert response.status_code == 200
        data = response.json()
        assert "status" in data
        assert "services" in data
        assert "timestamp" in data
    
    @pytest.mark.asyncio
    async def test_store_memory(self):
        """Test storing a memory"""
        memory_data = {
            "content": "This is a test memory about Python programming",
            "user_id": self.test_user_id,
            "session_id": "test_session_1",
            "metadata": {
                "category": "programming",
                "language": "python",
                "importance": "high"
            }
        }
        
        response = await self.client.post(
            "/v1/memories/store",
            json=memory_data,
            headers=self.headers
        )
        
        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "stored"
        assert data["user_id"] == self.test_user_id
        assert "id" in data
        assert "created_at" in data
        
        return data["id"]
    
    @pytest.mark.asyncio
    async def test_search_memories(self):
        """Test searching memories"""
        # First store a memory
        memory_id = await self.test_store_memory()
        
        # Wait a moment for indexing
        await asyncio.sleep(1)
        
        # Search for the memory
        search_data = {
            "query": "Python programming",
            "user_id": self.test_user_id,
            "limit": 10,
            "threshold": 0.5,
            "include_graph": True
        }
        
        response = await self.client.post(
            "/v1/memories/search",
            json=search_data,
            headers=self.headers
        )
        
        assert response.status_code == 200
        data = response.json()
        assert "memories" in data
        assert "context" in data
        assert "total_count" in data
        assert data["total_count"] > 0
        
        # Check first memory result
        if data["memories"]:
            memory = data["memories"][0]
            assert "id" in memory
            assert "content" in memory
            assert "similarity" in memory
            assert memory["user_id"] == self.test_user_id
    
    @pytest.mark.asyncio
    async def test_retrieve_memories(self):
        """Test retrieving memories for conversation context"""
        # Store a memory first
        await self.test_store_memory()
        
        # Wait a moment for indexing
        await asyncio.sleep(1)
        
        # Retrieve memories based on conversation
        retrieve_data = {
            "messages": [
                {"role": "user", "content": "I want to learn about Python"},
                {"role": "assistant", "content": "Python is a great programming language"},
                {"role": "user", "content": "Tell me more about Python programming"}
            ],
            "user_id": self.test_user_id,
            "session_id": "test_session_1"
        }
        
        response = await self.client.post(
            "/v1/memories/retrieve",
            json=retrieve_data,
            headers=self.headers
        )
        
        assert response.status_code == 200
        data = response.json()
        assert "memories" in data
        assert "context" in data
        assert "total_count" in data
    
    @pytest.mark.asyncio
    async def test_get_user_memories(self):
        """Test getting all memories for a user"""
        # Store a memory first
        await self.test_store_memory()
        
        response = await self.client.get(
            f"/v1/memories/users/{self.test_user_id}",
            headers=self.headers
        )
        
        assert response.status_code == 200
        data = response.json()
        assert "memories" in data
        assert "total_count" in data
        assert data["total_count"] > 0
        
        # Check memory structure
        if data["memories"]:
            memory = data["memories"][0]
            assert "id" in memory
            assert "content" in memory
            assert "user_id" in memory
            assert "created_at" in memory
    
    @pytest.mark.asyncio
    async def test_delete_memory(self):
        """Test deleting a memory"""
        # Store a memory first
        memory_id = await self.test_store_memory()
        
        # Delete the memory
        response = await self.client.delete(
            f"/v1/memories/{memory_id}",
            headers=self.headers
        )
        
        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "deleted"
        assert data["id"] == memory_id
    
    @pytest.mark.asyncio
    async def test_authentication_required(self):
        """Test that authentication is required"""
        response = await self.client.get("/v1/memories/users/test_user")
        assert response.status_code == 401
    
    @pytest.mark.asyncio
    async def test_invalid_api_key(self):
        """Test invalid API key"""
        headers = {"Authorization": "Bearer invalid_key"}
        response = await self.client.get("/v1/memories/users/test_user", headers=headers)
        assert response.status_code == 401

if __name__ == "__main__":
    pytest.main([__file__, "-v"])