"""Test MCP 1MB response size limits."""

import pytest
import json
from core.response_limiter import MCPResponseLimiter

class TestMCPResponseLimiter:
    
    def test_response_size_estimation(self):
        """Test response size estimation accuracy."""
        simple_response = {"status": "success", "message": "test"}
        estimated_size = MCPResponseLimiter.estimate_response_size(simple_response)
        
        # Should be roughly the size of the JSON string
        json_str = json.dumps(simple_response)
        actual_size = len(json_str.encode('utf-8'))
        
        # Estimation should be close to actual (within 10%)
        assert abs(estimated_size - actual_size) / actual_size < 0.1
    
    def test_content_truncation(self):
        """Test intelligent content truncation."""
        long_content = "A" * 1000
        
        # Test truncation preserving start
        truncated = MCPResponseLimiter.truncate_content(long_content, 100, preserve_start=True)
        assert len(truncated) == 100
        assert truncated.startswith("A")
        assert truncated.endswith("...")
        
        # Test truncation preserving end
        truncated = MCPResponseLimiter.truncate_content(long_content, 100, preserve_start=False)
        assert len(truncated) == 100
        assert truncated.startswith("...")
        assert truncated.endswith("A")
    
    def test_memory_content_truncation(self):
        """Test memory-specific content truncation."""
        memory = {
            "id": "test_id",
            "content": "A" * 500,
            "context_type": "code",
            "similarity": 0.95
        }
        
        truncated = MCPResponseLimiter.truncate_memory_content(memory, 100)
        
        assert len(truncated["content"]) <= 100
        assert truncated["_truncated"] is True
        assert truncated["_original_length"] == 500
        assert truncated["id"] == "test_id"  # Other fields preserved
    
    def test_memory_response_limiting(self):
        """Test memory response size limiting."""
        # Create memories with large content
        large_memories = []
        for i in range(10):
            large_memories.append({
                "id": f"memory_{i}",
                "content": "X" * 100000,  # 100KB each
                "context_type": "test",
                "similarity": 0.9 - (i * 0.1)
            })
        
        base_response = {
            "status": "success",
            "message": "Test response",
            "query": "test query"
        }
        
        limited_memories, was_truncated = MCPResponseLimiter.limit_memory_response(
            large_memories, base_response
        )
        
        # Should have truncated to fit in 1MB
        test_response = base_response.copy()
        test_response["results"] = limited_memories
        
        response_size = MCPResponseLimiter.estimate_response_size(test_response)
        assert response_size <= MCPResponseLimiter.CONTENT_LIMIT
        assert was_truncated is True  # Should have been truncated
        assert len(limited_memories) <= len(large_memories)  # Some memories removed or truncated
    
    def test_context_bridge_response_limiting(self):
        """Test context bridge response size limiting."""
        # Create large related content
        large_content = []
        for i in range(15):
            large_content.append({
                "type": "memory",
                "content": "Y" * 80000,  # 80KB each
                "session": f"session_{i}",
                "context_type": "project",
                "similarity": 0.95 - (i * 0.05)
            })
        
        base_response = {
            "status": "success",
            "new_topic": "test topic",
            "session_name": "test_session"
        }
        
        limited_content, was_truncated = MCPResponseLimiter.limit_context_bridge_response(
            large_content, base_response
        )
        
        # Verify response size compliance
        test_response = base_response.copy()
        test_response["related_content"] = limited_content
        
        response_size = MCPResponseLimiter.estimate_response_size(test_response)
        assert response_size <= MCPResponseLimiter.CONTENT_LIMIT
        assert was_truncated is True
        assert len(limited_content) <= len(large_content)
    
    def test_truncation_notice_addition(self):
        """Test truncation notice functionality."""
        response = {
            "status": "success",
            "results": []
        }
        
        # Test with truncation
        modified_response = MCPResponseLimiter.add_truncation_notice(
            response, was_truncated=True, total_available=10
        )
        
        assert modified_response["_mcp_truncated"] is True
        assert "1MB" in modified_response["_truncation_reason"]
        assert modified_response["_total_available"] == 10
        assert "partial results" in modified_response["message"]
        
        # Test without truncation
        clean_response = {"status": "success"}
        unmodified_response = MCPResponseLimiter.add_truncation_notice(
            clean_response, was_truncated=False
        )
        
        assert "_mcp_truncated" not in unmodified_response
    
    def test_edge_case_empty_data(self):
        """Test handling of empty or minimal data."""
        # Test empty memories list
        limited, truncated = MCPResponseLimiter.limit_memory_response([], {})
        assert limited == []
        assert truncated is False
        
        # Test very small content
        tiny_memories = [{"content": "small", "id": "1"}]
        limited, truncated = MCPResponseLimiter.limit_memory_response(
            tiny_memories, {"status": "success"}
        )
        assert len(limited) == 1
        assert truncated is False
    
    def test_exact_limit_boundary(self):
        """Test behavior at exactly the 1MB boundary."""
        # Create content that's exactly at the limit
        target_size = MCPResponseLimiter.CONTENT_LIMIT
        
        # Start with base response
        base_response = {"status": "success", "message": "test"}
        base_size = MCPResponseLimiter.estimate_response_size(base_response)
        
        # Calculate content size to hit exactly the limit
        available_space = target_size - base_size - 100  # Small buffer for JSON structure
        content_per_memory = available_space // 3  # 3 memories
        
        boundary_memories = []
        for i in range(3):
            boundary_memories.append({
                "id": f"boundary_{i}",
                "content": "B" * content_per_memory,
                "similarity": 0.9
            })
        
        limited, truncated = MCPResponseLimiter.limit_memory_response(
            boundary_memories, base_response
        )
        
        # Should fit without truncation or with minimal truncation
        test_response = base_response.copy()
        test_response["results"] = limited
        response_size = MCPResponseLimiter.estimate_response_size(test_response)
        
        assert response_size <= MCPResponseLimiter.CONTENT_LIMIT
        assert len(limited) > 0  # Should keep at least some memories
    
    def test_large_summary_response_limiting(self):
        """Test marm_summary response size limiting with many log entries."""
        # Simulate large summary response that would exceed 1MB
        large_summary_data = {
            "status": "success", 
            "session_name": "test_session",
            "total_entries": 500,
            "summary": "\n".join([
                f"**2025-01-15 10:{i:02d}:00** [milestone-{i}]: Very long summary text that contains detailed information about what was accomplished during this particular session milestone including technical details, decisions made, and outcomes achieved." 
                for i in range(5000)  # 5000 entries with long summaries
            ])
        }
        
        # Verify this would exceed 1MB
        response_size = MCPResponseLimiter.estimate_response_size(large_summary_data)
        print(f"Large summary size: {response_size / 1024:.1f} KB")
        
        # This should be larger than our limit
        assert response_size > MCPResponseLimiter.CONTENT_LIMIT, "Test summary should exceed 1MB limit"
        
        # Test that our new summary endpoint would handle this properly
        # (simulating the progressive entry addition logic)
        base_response = {
            "status": "success",
            "session_name": "test_session",
            "total_entries": 500
        }
        
        # Simulate adding entries one by one until size limit hit
        summary_lines = ["# MARM Session Summary: test_session", "Generated: 2025-01-15", ""]
        entries_included = 0
        
        # Start with the base summary
        current_summary_text = "\n".join(summary_lines)

        for i in range(5000):
            entry_line = f"**2025-01-15 10:{i:02d}:00** [milestone-{i}]: {'A' * 500}" # Use a much larger string
            
            # Test adding the NEXT line
            test_summary = current_summary_text + "\n" + entry_line
            
            test_response = base_response.copy()
            test_response["summary"] = test_summary
            test_response["entry_count"] = i + 1
            
            response_size = MCPResponseLimiter.estimate_response_size(test_response)
            
            if response_size > MCPResponseLimiter.CONTENT_LIMIT:
                entries_included = i  # We successfully included i entries (0-indexed)
                break
            
            # If it fits, update the current summary
            current_summary_text = test_summary
        else:
            # If the loop completes without breaking, all entries were included
            entries_included = 5000

        print(f"Would include {entries_included} entries before hitting size limit")
        
        # Should have stopped before including all entries
        assert entries_included < 5000, "Should have truncated entries to stay under limit"
        assert entries_included > 0, "Should have included at least some entries"

if __name__ == "__main__":
    # Quick test run
    test = TestMCPResponseLimiter()
    
    print("Testing response size estimation...")
    test.test_response_size_estimation()
    print("✅ Passed")
    
    print("Testing content truncation...")
    test.test_content_truncation()
    print("✅ Passed")
    
    print("Testing memory response limiting...")
    test.test_memory_response_limiting()
    print("✅ Passed")
    
    print("Testing edge cases...")
    test.test_edge_case_empty_data()
    print("✅ Passed")
    
    print("Testing large summary response limiting...")
    test.test_large_summary_response_limiting()
    print("✅ Passed")
    
    print("\n🎉 All MCP size limit tests passed!")