"""
Integration test for vulnerability detection scenario.

This test validates the Brainary system using a real-world scenario:
detecting security vulnerabilities in code repositories and identifying CWE IDs.

Test Progression:
1. Basic primitive execution (perceive, think, action)
2. Memory system (working memory, prefetch, consolidation)
3. Resource management (allocation, quota checking)
4. Intelligent routing (experience cache, heuristics)
5. LLM integration (driver, cost tracking)
6. End-to-end workflow (full vulnerability scan)
"""

import pytest
import tempfile
import os
from pathlib import Path
from typing import List, Dict, Any

# Import Brainary components
from brainary import (
    # Core
    ExecutionContext,
    ExecutionMode,
    create_execution_context,
    CognitiveKernel,
    # Memory
    WorkingMemory,
    MemoryTier,
    PrefetchRequest,
    # Resource Management
    ResourceManager,
    ResourceQuota,
    # LLM
    ILLMDriver,
    OpenAIDriver,
    LLMRequest,
    create_driver,
    # Primitives
    Primitive,
    PrimitiveResult,
    CostMetrics,
    ConfidenceMetrics,
    # Errors
    BrainaryError,
    ResourceExhaustedError,
)

# Import security domain
from brainary.domains.security import (
    PerceiveCode,
    AnalyzeVulnerabilities,
    ValidateSecurityControls,
    VulnerabilityDetector,
    CWEDatabase,
)


# ===== Test Data =====

VULNERABLE_CODE_SAMPLES = {
    "sql_injection": {
        "code": '''
def get_user(username):
    query = "SELECT * FROM users WHERE username = '" + username + "'"
    cursor.execute(query)
    return cursor.fetchone()
''',
        "expected_cwe": ["CWE-89"],  # SQL Injection
        "severity": "HIGH",
    },
    "xss": {
        "code": '''
def render_comment(comment):
    return "<div>" + comment + "</div>"
''',
        "expected_cwe": ["CWE-79"],  # Cross-site Scripting
        "severity": "MEDIUM",
    },
    "hardcoded_credentials": {
        "code": '''
DATABASE_PASSWORD = "admin123"
API_KEY = "sk-1234567890abcdef"
''',
        "expected_cwe": ["CWE-798"],  # Use of Hard-coded Credentials
        "severity": "HIGH",
    },
    "path_traversal": {
        "code": '''
def read_file(filename):
    with open("/var/data/" + filename) as f:
        return f.read()
''',
        "expected_cwe": ["CWE-22"],  # Path Traversal
        "severity": "HIGH",
    },
    "insecure_random": {
        "code": '''
import random
token = random.randint(1000, 9999)
''',
        "expected_cwe": ["CWE-330"],  # Use of Insufficiently Random Values
        "severity": "MEDIUM",
    },
}

SAFE_CODE_SAMPLES = {
    "parameterized_query": {
        "code": '''
def get_user(username):
    query = "SELECT * FROM users WHERE username = ?"
    cursor.execute(query, (username,))
    return cursor.fetchone()
''',
        "expected_cwe": [],
    },
    "escaped_output": {
        "code": '''
from html import escape
def render_comment(comment):
    return "<div>" + escape(comment) + "</div>"
''',
        "expected_cwe": [],
    },
}


# ===== Test Fixtures =====

@pytest.fixture
def working_memory():
    """Create working memory instance."""
    return WorkingMemory(capacity=7, l2_capacity=100)


@pytest.fixture
def resource_manager():
    """Create resource manager instance."""
    return ResourceManager()


@pytest.fixture
def llm_driver():
    """Create LLM driver (mock in test mode)."""
    if os.getenv("BRAINARY_TEST_MODE"):
        # Return mock driver
        return None
    else:
        # Return real driver if API key available
        try:
            return create_driver("openai", model="gpt-4o-mini")
        except:
            pytest.skip("OpenAI API key not available")


@pytest.fixture
def execution_context():
    """Create test execution context."""
    return create_execution_context(
        program_name="test_vulnerability_detection",
        execution_mode=ExecutionMode.ADAPTIVE,
        quality_threshold=0.7,
        criticality=0.8,
        time_pressure=0.3,
        token_budget=10000,
        domain="security"
    )


# ===== Level 1: Basic Primitive Execution Tests =====

@pytest.mark.level1
class TestLevel1BasicPrimitives:
    """Test basic primitive execution without full system integration."""
    
    def test_perceive_code_execution(self, working_memory, execution_context):
        """Test code perception primitive."""
        # Set test mode
        os.environ["BRAINARY_TEST_MODE"] = "1"
        
        primitive = PerceiveCode()
        
        # Test with SQL injection code
        code = VULNERABLE_CODE_SAMPLES["sql_injection"]["code"]
        
        # Execute primitive directly
        result = primitive.execute(
            context=execution_context,
            working_memory=working_memory,
            code=code
        )
        
        # Verify result
        assert result.success
        assert result.content is not None
        assert "patterns" in result.content
        assert result.content["has_suspicious_patterns"]
        assert result.confidence.overall > 0.8
    
    def test_analyze_vulnerabilities_execution(self, working_memory, execution_context, llm_driver):
        """Test vulnerability analysis primitive."""
        os.environ["BRAINARY_TEST_MODE"] = "1"
        
        primitive = AnalyzeVulnerabilities()
        
        # Test with SQL injection code
        code = VULNERABLE_CODE_SAMPLES["sql_injection"]["code"]
        
        # Execute (with LLM disabled for faster testing)
        result = primitive.execute(
            context=execution_context,
            working_memory=working_memory,
            code=code,
            enable_llm=False
        )
        
        # Verify result
        assert result.success
        assert result.content["has_vulnerabilities"]
        assert len(result.content["vulnerabilities"]) > 0
        # Check that CWE-89 (SQL Injection) was detected
        cwe_ids = [v["cwe_id"] for v in result.content["vulnerabilities"]]
        assert "CWE-89" in cwe_ids


# ===== Level 2: Memory System Tests =====

@pytest.mark.level2
class TestLevel2MemorySystem:
    """Test memory system with L1/L2/L3 tiers."""
    
    def test_memory_storage_and_retrieval(self, working_memory):
        """Test basic memory operations."""
        # Store in L1
        item_id = working_memory.store(
            content={"code": "test", "vuln": True},
            tier=MemoryTier.L1_WORKING,
            importance=0.8,
            tags=["test", "vulnerability"],
        )
        
        assert item_id is not None
        
        # Retrieve from L1
        results = working_memory.retrieve(tags=["vulnerability"], top_k=1)
        assert len(results) == 1
        assert results[0].content["vuln"] is True
    
    def test_memory_tiers(self, working_memory):
        """Test L1/L2/L3 tier operations."""
        # Store in L1
        l1_id = working_memory.store(
            content={"tier": "L1"},
            tier=MemoryTier.L1_WORKING,
            importance=0.9,
        )
        
        # Store in L2
        l2_id = working_memory.store(
            content={"tier": "L2"},
            tier=MemoryTier.L2_EPISODIC,
            importance=0.6,
        )
        
        # Store in L3
        l3_id = working_memory.store(
            content={"tier": "L3"},
            tier=MemoryTier.L3_SEMANTIC,
            importance=0.4,
        )
        
        # Verify storage
        assert len(working_memory._l1_items) == 1
        assert len(working_memory._l2_items) == 1
        assert len(working_memory._l3_items) == 1
        
        # Test retrieval across tiers
        all_items = working_memory.retrieve(tier=None, top_k=3)
        assert len(all_items) == 3
    
    def test_memory_consolidation(self, working_memory):
        """Test L1 to L2/L3 consolidation."""
        # Fill L1 beyond capacity
        for i in range(10):
            working_memory.store(
                content={"index": i},
                tier=MemoryTier.L1_WORKING,
                importance=0.5 + i * 0.05,  # Varying importance
                tags=[f"item_{i}"],
            )
        
        # Trigger consolidation
        results = working_memory.consolidate()
        
        # Verify consolidation
        assert results["to_l2"] + results["to_l3"] + results["discarded"] > 0
        assert len(working_memory._l1_items) <= working_memory.capacity
    
    def test_memory_prefetch(self, working_memory):
        """Test intelligent prefetching."""
        # Store items in L2 and L3
        for i in range(5):
            working_memory.store(
                content={"vuln_id": i},
                tier=MemoryTier.L2_EPISODIC,
                importance=0.7,
                tags=["vulnerability", "sql_injection"],
            )
        
        # Create prefetch request
        request = PrefetchRequest(
            context_tags=["vulnerability", "sql_injection"],
            primitive_name="analyze_vulnerabilities",
            max_items=3,
            min_salience=0.3,
        )
        
        # Prefetch
        prefetched_ids = working_memory.prefetch(request)
        
        # Verify prefetch
        assert len(prefetched_ids) > 0
        assert len(prefetched_ids) <= 3
    
    def test_memory_statistics(self, working_memory):
        """Test memory statistics collection."""
        # Add some items
        for i in range(5):
            working_memory.store(
                content={"index": i},
                tier=MemoryTier.L1_WORKING,
                importance=0.5,
            )
        
        # Get statistics
        stats = working_memory.get_statistics()
        
        # Verify stats
        assert stats.l1_size == 5
        assert stats.l1_capacity == 7
        assert stats.l1_utilization > 0
        assert stats.total_stores == 5


# ===== Level 3: Resource Management Tests =====

@pytest.mark.level3
class TestLevel3ResourceManagement:
    """Test resource allocation and quota management."""
    
    def test_resource_allocation(self, resource_manager, execution_context):
        """Test resource allocation with quota."""
        # Create quota
        quota = ResourceQuota(
            token_budget=1000,
            timeout_ms=5000,
            memory_slots=5,
            priority=0.5,
        )
        
        # Estimate cost
        cost = CostMetrics(
            tokens=100,
            latency_ms=500,
            memory_slots=1,
            provider_cost_usd=0.001,
        )
        
        # Allocate resources
        allocation = resource_manager.allocate(execution_context, cost)
        
        # Verify allocation
        assert allocation is not None
        assert allocation.allocation_id is not None
        assert allocation.quota.token_budget >= cost.tokens
    
    def test_quota_checking(self, resource_manager, execution_context):
        """Test quota availability checking."""
        # Check initial quota
        available = resource_manager.check_quota(execution_context)
        
        assert "tokens" in available
        assert "timeout_ms" in available
        assert "memory_slots" in available
    
    def test_resource_exhaustion(self, resource_manager, execution_context):
        """Test resource exhaustion error."""
        # Create very large cost estimate
        huge_cost = CostMetrics(
            tokens=1_000_000_000,  # 1 billion tokens
            latency_ms=100,
            memory_slots=1,
            provider_cost_usd=100.0,
        )
        
        # Should raise ResourceExhaustedError
        with pytest.raises(ResourceExhaustedError):
            resource_manager.allocate(execution_context, huge_cost)
    
    def test_usage_tracking(self, resource_manager, execution_context):
        """Test actual usage tracking."""
        # Allocate
        cost = CostMetrics(tokens=100, latency_ms=500, memory_slots=1)
        allocation = resource_manager.allocate(execution_context, cost)
        
        # Track actual usage
        actual_cost = CostMetrics(tokens=80, latency_ms=450, memory_slots=1)
        resource_manager.update_usage(allocation.allocation_id, actual_cost)
        
        # Release
        resource_manager.release(allocation.allocation_id)
        
        # Verify tracking worked (no exceptions)
        assert True


# ===== Level 4: End-to-End Vulnerability Detection Tests =====

@pytest.mark.level4
class TestLevel4EndToEndDetection:
    """Test complete vulnerability detection workflow."""
    
    def test_detect_sql_injection(self, execution_context, llm_driver):
        """Test SQL injection detection."""
        os.environ["BRAINARY_TEST_MODE"] = "1"
        
        # Step 1: Perceive code
        perceive = PerceiveCode()
        context = execution_context
        context.inputs = {"code": VULNERABLE_CODE_SAMPLES["sql_injection"]["code"]}
        
        perceive_result = perceive.execute(context)
        assert perceive_result.success
        
        # Step 2: Analyze vulnerabilities
        analyze = AnalyzeVulnerabilities(llm_driver)
        context.inputs = {
            "code": VULNERABLE_CODE_SAMPLES["sql_injection"]["code"],
            "parsed_code": perceive_result.output,
        }
        
        analyze_result = analyze.execute(context)
        
        # Verify detection
        assert analyze_result.success
        assert analyze_result.output["has_vulnerabilities"]
        
        vulnerabilities = analyze_result.output["vulnerabilities"]
        cwe_ids = [v["cwe_id"] for v in vulnerabilities]
        
        assert any("CWE-89" in cwe_id for cwe_id in cwe_ids)
    
    def test_detect_xss(self, execution_context, llm_driver):
        """Test XSS detection."""
        os.environ["BRAINERY_TEST_MODE"] = "1"
        
        # Full workflow
        perceive = PerceiveCode()
        analyze = AnalyzeVulnerabilities(llm_driver)
        
        context = execution_context
        context.inputs = {"code": VULNERABLE_CODE_SAMPLES["xss"]["code"]}
        
        perceive_result = perceive.execute(context)
        context.inputs = {
            "code": VULNERABLE_CODE_SAMPLES["xss"]["code"],
            "parsed_code": perceive_result.output,
        }
        
        analyze_result = analyze.execute(context)
        
        # Verify
        assert analyze_result.success
        vulnerabilities = analyze_result.output["vulnerabilities"]
        assert any("CWE-79" in v["cwe_id"] for v in vulnerabilities)
    
    def test_detect_hardcoded_credentials(self, execution_context, llm_driver):
        """Test hardcoded credentials detection."""
        os.environ["BRAINERY_TEST_MODE"] = "1"
        
        perceive = PerceiveCode()
        analyze = AnalyzeVulnerabilities(llm_driver)
        
        context = execution_context
        context.inputs = {"code": VULNERABLE_CODE_SAMPLES["hardcoded_credentials"]["code"]}
        
        perceive_result = perceive.execute(context)
        context.inputs = {
            "code": VULNERABLE_CODE_SAMPLES["hardcoded_credentials"]["code"],
            "parsed_code": perceive_result.output,
        }
        
        analyze_result = analyze.execute(context)
        
        # Verify
        assert analyze_result.success
        vulnerabilities = analyze_result.output["vulnerabilities"]
        assert any("CWE-798" in v["cwe_id"] for v in vulnerabilities)
    
    def test_safe_code_no_vulnerabilities(self, execution_context, llm_driver):
        """Test that safe code is correctly identified."""
        os.environ["BRAINARY_TEST_MODE"] = "1"
        
        perceive = PerceiveCode()
        analyze = AnalyzeVulnerabilities(llm_driver)
        
        context = execution_context
        context.inputs = {"code": SAFE_CODE_SAMPLES["parameterized_query"]["code"]}
        
        perceive_result = perceive.execute(context)
        context.inputs = {
            "code": SAFE_CODE_SAMPLES["parameterized_query"]["code"],
            "parsed_code": perceive_result.output,
        }
        
        analyze_result = analyze.execute(context)
        
        # Verify no vulnerabilities detected
        assert analyze_result.success
        assert not analyze_result.output["has_vulnerabilities"]
    
    def test_multiple_vulnerabilities(self, execution_context, llm_driver):
        """Test detection of multiple vulnerabilities in one code sample."""
        os.environ["BRAINERY_TEST_MODE"] = "1"
        
        # Code with multiple vulnerabilities
        multi_vuln_code = '''
def process_request(username, comment):
    # SQL Injection
    query = "SELECT * FROM users WHERE name = '" + username + "'"
    cursor.execute(query)
    
    # XSS
    output = "<div>" + comment + "</div>"
    
    # Hardcoded password
    db_password = "admin123"
    
    return output
'''
        
        perceive = PerceiveCode()
        analyze = AnalyzeVulnerabilities(llm_driver)
        
        context = execution_context
        context.inputs = {"code": multi_vuln_code}
        
        perceive_result = perceive.execute(context)
        context.inputs = {
            "code": multi_vuln_code,
            "parsed_code": perceive_result.output,
        }
        
        analyze_result = analyze.execute(context)
        
        # Verify multiple vulnerabilities detected
        assert analyze_result.success
        assert analyze_result.output["has_vulnerabilities"]
        assert analyze_result.output["count"] >= 2


# ===== Level 5: Performance and Scale Tests =====

@pytest.mark.level5
class TestLevel5Performance:
    """Test performance characteristics and scalability."""
    
    def test_memory_access_performance(self, working_memory):
        """Test L1 memory access is < 1ms."""
        import time
        
        # Store item
        item_id = working_memory.store(
            content={"test": "data"},
            tier=MemoryTier.L1_WORKING,
        )
        
        # Measure retrieval time
        start = time.time()
        result = working_memory.retrieve(query=item_id)
        elapsed_ms = (time.time() - start) * 1000
        
        # L1 should be < 1ms
        assert elapsed_ms < 1.0
        assert len(result) == 1
    
    def test_cost_estimation_accuracy(self, execution_context):
        """Test cost estimation is reasonable."""
        os.environ["BRAINARY_TEST_MODE"] = "1"
        
        primitive = PerceiveCode()
        
        context = execution_context
        context.inputs = {"code": VULNERABLE_CODE_SAMPLES["sql_injection"]["code"]}
        
        # Estimate
        estimated_cost = primitive.estimate_cost(context)
        
        # Execute
        result = primitive.execute(context)
        
        # Verify estimate was reasonable (within 2x)
        assert estimated_cost.tokens > 0
        assert estimated_cost.latency_ms > 0
    
    def test_batch_vulnerability_detection(self, execution_context, llm_driver):
        """Test detecting vulnerabilities across multiple code samples."""
        os.environ["BRAINERY_TEST_MODE"] = "1"
        
        perceive = PerceiveCode()
        analyze = AnalyzeVulnerabilities(llm_driver)
        
        results = []
        
        for name, sample in VULNERABLE_CODE_SAMPLES.items():
            context = execution_context
            context.inputs = {"code": sample["code"]}
            
            perceive_result = perceive.execute(context)
            context.inputs = {
                "code": sample["code"],
                "parsed_code": perceive_result.output,
            }
            
            analyze_result = analyze.execute(context)
            results.append({
                "name": name,
                "has_vulns": analyze_result.output["has_vulnerabilities"],
                "count": analyze_result.output.get("count", 0),
            })
        
        # Verify all vulnerable samples detected
        assert all(r["has_vulns"] for r in results)
        assert sum(r["count"] for r in results) >= len(VULNERABLE_CODE_SAMPLES)


# ===== Pytest Configuration =====

def pytest_configure(config):
    """Configure pytest."""
    config.addinivalue_line("markers", "slow: marks tests as slow")
    config.addinivalue_line("markers", "integration: marks tests as integration tests")
    config.addinivalue_line("markers", "requires_llm: marks tests requiring LLM API")


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