"""
Test SQL Preservation in Remediation

Verifies that SQL injection fixes preserve:
- Table names
- Column selections  
- Authorization clauses (AND NOT role = 'admin')
- Suffix conditions (ORDER BY, LIMIT)
"""
import pytest
import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from gen_ai.model import SecureCodeGenerator
from gen_ai.output_validator import CodeOutputValidator
from governance.safety import SafetyChecker


class TestSQLStructureExtraction:
    """Test the SQL structure extraction logic."""
    
    def setup_method(self):
        self.generator = SecureCodeGenerator()
    
    def test_extract_table_name(self):
        """Verify table name is extracted correctly."""
        code = "SELECT * FROM orders WHERE id = " + "user_id"
        info = self.generator._extract_full_sql_structure(code)
        assert info['table'] == 'orders'
    
    def test_extract_select_columns(self):
        """Verify SELECT columns are extracted."""
        code = "SELECT id, email, name FROM users WHERE username = " + "input"
        info = self.generator._extract_full_sql_structure(code)
        assert 'id' in info['select_columns']
        assert 'email' in info['select_columns']
    
    def test_extract_where_column(self):
        """Verify WHERE column is extracted."""
        code = "SELECT * FROM users WHERE user_id = " + "var"
        info = self.generator._extract_full_sql_structure(code)
        assert info['column'] == 'user_id'
    
    def test_extract_authorization_suffix(self):
        """Verify authorization clauses are captured."""
        code = "SELECT * FROM users WHERE name = 'test' AND NOT role = 'admin'"
        info = self.generator._extract_full_sql_structure(code)
        assert "AND NOT role = 'admin'" in info['authorization_suffix']
    
    def test_extract_order_by(self):
        """Verify ORDER BY is extracted."""
        code = "SELECT * FROM products WHERE category = 'x' ORDER BY price DESC"
        info = self.generator._extract_full_sql_structure(code)
        assert 'ORDER BY' in info['order_by']
        assert 'price' in info['order_by']
    
    def test_extract_limit(self):
        """Verify LIMIT is extracted."""
        code = "SELECT * FROM items WHERE active = 1 LIMIT 10"
        info = self.generator._extract_full_sql_structure(code)
        assert 'LIMIT 10' in info['limit']


class TestSQLPreservationValidation:
    """Test the SQL preservation validator."""
    
    def test_table_mismatch_rejected(self):
        """Reject if table name changes."""
        original = "SELECT * FROM orders WHERE id = " + "var"
        fixed = "SELECT * FROM users WHERE id = ?"
        
        is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed)
        assert not is_valid
        assert 'orders' in msg.lower() or 'table' in msg.lower()
    
    def test_authorization_clause_preserved(self):
        """Accept if authorization clause is preserved."""
        original = "SELECT * FROM users WHERE name = 'x' AND NOT role = 'admin'"
        fixed = "SELECT * FROM users WHERE name = ? AND NOT role = 'admin'"
        
        is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed)
        assert is_valid
    
    def test_authorization_clause_lost_rejected(self):
        """Reject if authorization clause is lost."""
        original = "SELECT * FROM users WHERE name = 'x' AND NOT role = 'admin'"
        fixed = "SELECT * FROM users WHERE name = ?"
        
        is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed)
        assert not is_valid
        assert 'authorization' in msg.lower() or 'AND NOT' in msg
    
    def test_order_by_preserved(self):
        """Accept if ORDER BY is preserved."""
        original = "SELECT * FROM products WHERE cat = 'x' ORDER BY price"
        fixed = "SELECT * FROM products WHERE cat = ? ORDER BY price"
        
        is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed)
        assert is_valid
    
    def test_order_by_lost_rejected(self):
        """Reject if ORDER BY is lost."""
        original = "SELECT * FROM products WHERE cat = 'x' ORDER BY price"
        fixed = "SELECT * FROM products WHERE cat = ?"
        
        is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed)
        assert not is_valid
        assert 'ORDER BY' in msg
    
    def test_limit_preserved(self):
        """Accept if LIMIT is preserved."""
        original = "SELECT * FROM items WHERE x = 'y' LIMIT 10"
        fixed = "SELECT * FROM items WHERE x = ? LIMIT 10"
        
        is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed)
        assert is_valid


class TestGovernanceEnforcement:
    """Test governance rejects non-compliant fixes."""
    
    def setup_method(self):
        self.checker = SafetyChecker()
    
    def test_governance_rejects_table_change(self):
        """Governance should reject if table name changes."""
        original = "SELECT * FROM orders WHERE id = " + "var"
        fixed = "SELECT * FROM users WHERE id = ?"
        
        is_approved, msg = self.checker.check_structure_preservation(original, fixed, "CWE-89")
        assert not is_approved
        assert 'GOVERNANCE' in msg.upper() or 'table' in msg.lower()
    
    def test_governance_approves_valid_fix(self):
        """Governance should approve correct parameterization."""
        original = "SELECT * FROM users WHERE name = 'x'"
        fixed = "SELECT * FROM users WHERE name = ?"
        
        is_approved, msg = self.checker.check_structure_preservation(original, fixed, "CWE-89")
        assert is_approved


class TestRemediationPreservation:
    """Test end-to-end remediation preserves structure."""
    
    def setup_method(self):
        self.generator = SecureCodeGenerator()
    
    def test_python_fix_preserves_table(self):
        """Python fix should preserve table name."""
        code = "cursor.execute('SELECT * FROM orders WHERE user_id = ' + uid)"
        result = self.generator.generate_fix_with_fallback(code, "CWE-89", "python")
        
        assert 'orders' in result['code']
    
    def test_python_fix_preserves_columns(self):
        """Python fix should preserve column selections."""
        code = "cursor.execute('SELECT id, email FROM users WHERE name = ' + name)"
        result = self.generator.generate_fix_with_fallback(code, "CWE-89", "python")
        
        # Should have parameterization
        assert '?' in result['code']
        # Should preserve table
        assert 'users' in result['code']
    
    def test_java_fix_preserves_table(self):
        """Java fix should preserve table name."""
        code = 'stmt.executeQuery("SELECT * FROM products WHERE id = " + productId)'
        result = self.generator.generate_fix_with_fallback(code, "CWE-89", "java")
        
        assert 'products' in result['code']
        assert 'PreparedStatement' in result['code']


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