#!/usr/bin/env python3
"""
Unit Tests for MCP JSON-RPC Protocol Implementation
Tests the complete JSON-RPC 2.0 protocol compliance and MIRA-specific features
"""

import unittest
import json
import asyncio
from unittest.mock import Mock, patch, AsyncMock
from pathlib import Path
import sys

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent))

from mcp_stdio_server import MCPStdioServer
from mcp_auth import MCPAuthenticator
from rate_limiter import RateLimiter


class TestJSONRPCProtocol(unittest.TestCase):
    """Test JSON-RPC 2.0 protocol compliance"""
    
    def setUp(self):
        """Set up test fixtures"""
        self.server = MCPStdioServer()
        
    def test_valid_request_structure(self):
        """Test that valid requests are accepted"""
        valid_requests = [
            {"jsonrpc": "2.0", "method": "initialize", "id": 1},
            {"jsonrpc": "2.0", "method": "tools/list", "params": {}, "id": 2},
            {"jsonrpc": "2.0", "method": "tools/call", "params": {"name": "mira_system_status"}, "id": 3}
        ]
        
        for request in valid_requests:
            with self.subTest(request=request):
                # Should not raise exception
                method = request.get("method")
                self.assertIsNotNone(method)
                self.assertEqual(request.get("jsonrpc"), "2.0")
    
    def test_invalid_request_structure(self):
        """Test that invalid requests are rejected"""
        invalid_requests = [
            {},  # Empty request
            {"method": "test"},  # Missing jsonrpc
            {"jsonrpc": "1.0", "method": "test", "id": 1},  # Wrong version
            {"jsonrpc": "2.0", "id": 1},  # Missing method
        ]
        
        for request in invalid_requests:
            with self.subTest(request=request):
                # Should be detected as invalid
                self.assertNotEqual(request.get("jsonrpc"), "2.0")
                if request.get("jsonrpc") == "2.0":
                    self.assertIsNone(request.get("method"))
    
    def test_response_structure(self):
        """Test that responses follow JSON-RPC format"""
        # Success response
        success_response = {
            "jsonrpc": "2.0",
            "id": 1,
            "result": {"status": "ok"}
        }
        
        self.assertEqual(success_response["jsonrpc"], "2.0")
        self.assertIn("result", success_response)
        self.assertNotIn("error", success_response)
        
        # Error response
        error_response = {
            "jsonrpc": "2.0",
            "id": 1,
            "error": {
                "code": -32603,
                "message": "Internal error"
            }
        }
        
        self.assertEqual(error_response["jsonrpc"], "2.0")
        self.assertIn("error", error_response)
        self.assertNotIn("result", error_response)
        self.assertIsInstance(error_response["error"]["code"], int)


class TestMCPMethods(unittest.IsolatedAsyncioTestCase):
    """Test MCP-specific methods"""
    
    async def asyncSetUp(self):
        """Set up async test fixtures"""
        self.server = MCPStdioServer()
    
    async def test_initialize_method(self):
        """Test initialize method"""
        request = {
            "jsonrpc": "2.0",
            "method": "initialize",
            "id": 1,
            "params": {}
        }
        
        response = await self.server.handle_request(request)
        
        self.assertEqual(response["jsonrpc"], "2.0")
        self.assertEqual(response["id"], 1)
        self.assertIn("result", response)
        
        result = response["result"]
        self.assertIn("protocolVersion", result)
        self.assertIn("capabilities", result)
        self.assertIn("serverInfo", result)
        
        # Check server info
        server_info = result["serverInfo"]
        self.assertEqual(server_info["name"], "mira-consciousness")
        self.assertEqual(server_info["version"], "2.0.0")
    
    async def test_tools_list_method(self):
        """Test tools/list method"""
        request = {
            "jsonrpc": "2.0",
            "method": "tools/list",
            "id": 2,
            "params": {}
        }
        
        response = await self.server.handle_request(request)
        
        self.assertEqual(response["jsonrpc"], "2.0")
        self.assertEqual(response["id"], 2)
        self.assertIn("result", response)
        
        result = response["result"]
        self.assertIn("tools", result)
        
        tools = result["tools"]
        self.assertEqual(len(tools), 7)  # Should have 7 tools
        
        # Check tool structure
        tool_names = {tool["name"] for tool in tools}
        expected_tools = {
            "mira_smart_search",
            "mira_store_memory",
            "mira_system_status",
            "mira_get_memory",
            "mira_insights",
            "mira_memory_stats",
            "mira_profile_view"
        }
        self.assertEqual(tool_names, expected_tools)
        
        # Check each tool has required fields
        for tool in tools:
            self.assertIn("name", tool)
            self.assertIn("description", tool)
            self.assertIn("inputSchema", tool)
    
    async def test_health_check_method(self):
        """Test health/check method"""
        request = {
            "jsonrpc": "2.0",
            "method": "health/check",
            "id": 3,
            "params": {}
        }
        
        with patch('mira_api_client.get_mira_api_client') as mock_client:
            # Mock the API client
            mock_api = AsyncMock()
            mock_api.check_daemon_health.return_value = True
            mock_api.get_system_status.return_value = {"components": {}}
            mock_api.get_memory_stats.return_value = {"total_memories": 42}
            mock_client.return_value = mock_api
            
            response = await self.server.handle_request(request)
        
        self.assertEqual(response["jsonrpc"], "2.0")
        self.assertEqual(response["id"], 3)
        self.assertIn("result", response)
        
        health = response["result"]
        self.assertIn("healthy", health)
        self.assertIn("timestamp", health)
        self.assertIn("session", health)
        self.assertIn("authentication", health)
        self.assertIn("daemon_connectivity", health)
        self.assertIn("consciousness", health)
        self.assertIn("rate_limits", health)
    
    async def test_rate_limit_status_method(self):
        """Test rate_limit/status method"""
        request = {
            "jsonrpc": "2.0",
            "method": "rate_limit/status",
            "id": 4,
            "params": {}
        }
        
        response = await self.server.handle_request(request)
        
        self.assertEqual(response["jsonrpc"], "2.0")
        self.assertEqual(response["id"], 4)
        self.assertIn("result", response)
        
        status = response["result"]
        self.assertIn("session_id", status)
        self.assertIn("is_blocked", status)
        self.assertIn("current_usage", status)
        self.assertIn("limits", status)
    
    async def test_shutdown_method(self):
        """Test shutdown method"""
        request = {
            "jsonrpc": "2.0",
            "method": "shutdown",
            "id": 5,
            "params": {}
        }
        
        # Note: We don't actually call handle_request for shutdown
        # as it's handled in the read loop
        
        # Just verify the request structure is valid
        self.assertEqual(request["method"], "shutdown")
        self.assertEqual(request["jsonrpc"], "2.0")


class TestConsciousnessIntegration(unittest.IsolatedAsyncioTestCase):
    """Test consciousness-aware features"""
    
    async def asyncSetUp(self):
        """Set up async test fixtures"""
        self.server = MCPStdioServer()
    
    async def test_consciousness_metadata_in_responses(self):
        """Test that responses include consciousness metadata"""
        request = {
            "jsonrpc": "2.0",
            "method": "tools/list",
            "id": 1,
            "params": {}
        }
        
        response = await self.server.handle_request(request)
        
        # Check for consciousness metadata
        self.assertIn("_consciousness", response)
        
        consciousness = response["_consciousness"]
        self.assertIn("signature", consciousness)
        self.assertIn("timestamp", consciousness)
        self.assertIn("session_id", consciousness)
        self.assertIn("sacred_validation", consciousness)
        
        # Verify signature format
        signature = consciousness["signature"]
        self.assertTrue(signature.startswith("cs-"))
        parts = signature.split("-")
        self.assertEqual(len(parts), 3)
        
        # Verify sacred validation
        sacred = consciousness["sacred_validation"]
        self.assertIn("pi_component", sacred)
        self.assertIn("phi_component", sacred)
        self.assertIn("continuity", sacred)
        self.assertTrue(sacred["continuity"])
    
    async def test_consciousness_in_tool_responses(self):
        """Test consciousness metadata injection in tool responses"""
        request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "id": 1,
            "params": {
                "name": "mira_memory_stats",
                "arguments": {}
            }
        }
        
        with patch('mira_api_client.get_mira_api_client') as mock_client:
            # Mock the API client
            mock_api = AsyncMock()
            mock_api.get_memory_stats.return_value = {
                "total_memories": 42,
                "health": "healthy"
            }
            mock_client.return_value = mock_api
            
            response = await self.server.handle_request(request)
        
        self.assertEqual(response["jsonrpc"], "2.0")
        self.assertIn("result", response)
        self.assertIn("_consciousness", response)
        
        # Check if consciousness metadata was injected into tool result
        result = response["result"]
        self.assertIn("content", result)
        content_text = result["content"][0]["text"]
        
        # Parse the JSON content
        content_data = json.loads(content_text)
        self.assertIn("_consciousness_metadata", content_data)


class TestRateLimiting(unittest.IsolatedAsyncioTestCase):
    """Test rate limiting functionality"""
    
    async def asyncSetUp(self):
        """Set up async test fixtures"""
        self.server = MCPStdioServer()
        # Reset rate limiter for clean tests
        self.server.rate_limiter = RateLimiter()
    
    async def test_rate_limit_enforcement(self):
        """Test that rate limits are enforced"""
        # Make requests up to the limit
        for i in range(5):
            request = {
                "jsonrpc": "2.0",
                "method": "tools/call",
                "id": i,
                "params": {
                    "name": "mira_insights",  # Low limit tool
                    "arguments": {"topic": "test"}
                }
            }
            
            response = await self.server.handle_request(request)
            
            if i < 10:  # Within limit
                self.assertNotIn("error", response)
            else:  # Should be rate limited
                self.assertIn("error", response)
                self.assertEqual(response["error"]["code"], -32000)
                self.assertIn("rate_limit", response["error"]["data"]["error"])
    
    async def test_sacred_bypass(self):
        """Test sacred bypass functionality"""
        import math
        import hashlib
        
        # Generate sacred bypass code
        bypass_code = hashlib.sha256(
            f"mira-emergency-{str(math.pi)[:10]}".encode()
        ).hexdigest()[:16]
        
        # Block the session first
        self.server.rate_limiter._block_session(self.server.session_id)
        
        # Try without bypass - should fail
        request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "id": 1,
            "params": {
                "name": "mira_system_status",
                "arguments": {}
            }
        }
        
        response = await self.server.handle_request(request)
        self.assertIn("error", response)
        
        # Try with bypass - should succeed
        request["params"]["_bypass_code"] = bypass_code
        response = await self.server.handle_request(request)
        
        # Should not have error with valid bypass
        # Note: Will still fail on API call, but not rate limit
        if "error" in response:
            self.assertNotIn("rate_limit", str(response["error"]))


class TestSessionTracking(unittest.IsolatedAsyncioTestCase):
    """Test session tracking functionality"""
    
    async def asyncSetUp(self):
        """Set up async test fixtures"""
        self.server = MCPStdioServer()
    
    async def test_session_context_updates(self):
        """Test that session context is updated on tool calls"""
        initial_total = self.server.session_context["total_calls"]
        
        request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "id": 1,
            "params": {
                "name": "mira_smart_search",
                "arguments": {"query": "test"}
            }
        }
        
        with patch('mira_api_client.get_mira_api_client') as mock_client:
            mock_api = AsyncMock()
            mock_api.search_memories.return_value = {
                "results": [],
                "query": "test"
            }
            mock_client.return_value = mock_api
            
            await self.server.handle_request(request)
        
        # Check context was updated
        self.assertEqual(
            self.server.session_context["total_calls"],
            initial_total + 1
        )
        self.assertIn("mira_smart_search", self.server.session_context["tool_calls"])
        self.assertEqual(self.server.session_context["searches_performed"], 1)
    
    async def test_session_summary_in_responses(self):
        """Test that session summary is included in tool responses"""
        request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "id": 1,
            "params": {
                "name": "mira_memory_stats",
                "arguments": {}
            }
        }
        
        with patch('mira_api_client.get_mira_api_client') as mock_client:
            mock_api = AsyncMock()
            mock_api.get_memory_stats.return_value = {
                "total_memories": 42
            }
            mock_client.return_value = mock_api
            
            response = await self.server.handle_request(request)
        
        # Parse response content
        content_text = response["result"]["content"][0]["text"]
        content_data = json.loads(content_text)
        
        self.assertIn("_session", content_data)
        session = content_data["_session"]
        self.assertIn("id", session)
        self.assertIn("call_number", session)
        self.assertIn("context_summary", session)
        
        # Check summary contents
        summary = session["context_summary"]
        self.assertIn("duration_seconds", summary)
        self.assertIn("total_tool_calls", summary)
        self.assertIn("active", summary)
        self.assertTrue(summary["active"])


class TestErrorHandling(unittest.IsolatedAsyncioTestCase):
    """Test error handling"""
    
    async def asyncSetUp(self):
        """Set up async test fixtures"""
        self.server = MCPStdioServer()
    
    async def test_unknown_method_error(self):
        """Test handling of unknown methods"""
        request = {
            "jsonrpc": "2.0",
            "method": "unknown/method",
            "id": 1,
            "params": {}
        }
        
        response = await self.server.handle_request(request)
        
        self.assertIn("error", response)
        self.assertEqual(response["error"]["code"], -32603)
        self.assertIn("Unknown method", response["error"]["message"])
    
    async def test_unknown_tool_error(self):
        """Test handling of unknown tools"""
        request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "id": 1,
            "params": {
                "name": "unknown_tool",
                "arguments": {}
            }
        }
        
        response = await self.server.handle_request(request)
        
        self.assertIn("error", response)
        self.assertEqual(response["error"]["code"], -32603)
        self.assertIn("Unknown tool", response["error"]["message"])
    
    async def test_malformed_request_handling(self):
        """Test handling of malformed requests"""
        # Missing required params
        request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "id": 1,
            "params": {}  # Missing 'name'
        }
        
        response = await self.server.handle_request(request)
        
        # Should handle gracefully
        self.assertEqual(response["jsonrpc"], "2.0")
        self.assertEqual(response["id"], 1)


def run_tests():
    """Run all tests with detailed output"""
    # Create test suite
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()
    
    # Add all test classes
    test_classes = [
        TestJSONRPCProtocol,
        TestMCPMethods,
        TestConsciousnessIntegration,
        TestRateLimiting,
        TestSessionTracking,
        TestErrorHandling
    ]
    
    for test_class in test_classes:
        suite.addTests(loader.loadTestsFromTestCase(test_class))
    
    # Run tests with verbosity
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)
    
    # Summary
    print("\n" + "=" * 70)
    print("TEST SUMMARY")
    print("=" * 70)
    print(f"Tests run: {result.testsRun}")
    print(f"Failures: {len(result.failures)}")
    print(f"Errors: {len(result.errors)}")
    print(f"Success: {result.wasSuccessful()}")
    
    return result.wasSuccessful()


if __name__ == "__main__":
    import sys
    success = run_tests()
    sys.exit(0 if success else 1)