"""Tests for code extraction and filtering logic."""

import pytest
import os
import sys
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))

from zest import Question


def test_decorator_removed_from_extraction(temp_test_file, valid_token, monkeypatch, server_url):
    """Test that @question decorators are removed from extracted code."""
    # Set server URL to test server
    monkeypatch.setenv('ZEST_SERVER_URL', server_url)
    
    code = '''from zest import question

@question('token123')
def my_solution(nums, target):
    return [0, 1]

my_solution.test()
'''
    
    temp_test_file.write_text(code)
    
    # Import the module and create Question instance
    spec = __import__('importlib').util.spec_from_file_location("test_module", temp_test_file)
    module = __import__('importlib').util.module_from_spec(spec)
    
    # We need to test the extraction logic directly
    # Create a simple function for testing
    def test_func(nums, target):
        return [0, 1]
    
    q = Question(1, func=test_func)
    
    # Mock the module to simulate extraction
    import types
    mock_module = types.ModuleType('test_module')
    mock_module.__file__ = str(temp_test_file)
    
    # Add the code as the module source
    import inspect
    original_getsource = inspect.getsource
    
    def mock_getsource(obj):
        if obj == mock_module:
            return code
        return original_getsource(obj)
    
    monkeypatch.setattr(inspect, 'getsource', mock_getsource)
    monkeypatch.setattr(inspect, 'getmodule', lambda x: mock_module)
    
    extracted = q._extract_user_code()
    
    # Verify decorator is removed
    assert '@question' not in extracted
    assert 'def my_solution' in extracted


def test_imports_removed_from_extraction(temp_test_file, monkeypatch, server_url):
    """Test that zest imports are removed from extracted code."""
    monkeypatch.setenv('ZEST_SERVER_URL', server_url)
    
    code = '''from zest import question
import zest

def my_solution(nums, target):
    return [0, 1]
'''
    
    temp_test_file.write_text(code)
    
    def test_func(nums, target):
        return [0, 1]
    
    q = Question(1, func=test_func)
    
    # Mock the extraction
    import types
    import inspect
    mock_module = types.ModuleType('test_module')
    
    def mock_getsource(obj):
        if obj == mock_module:
            return code
        return code
    
    monkeypatch.setattr(inspect, 'getsource', mock_getsource)
    monkeypatch.setattr(inspect, 'getmodule', lambda x: mock_module)
    
    extracted = q._extract_user_code()
    
    # Verify imports are removed
    assert 'from zest import' not in extracted
    assert 'import zest' not in extracted
    assert 'def my_solution' in extracted


def test_test_calls_removed_from_extraction(temp_test_file, monkeypatch, server_url):
    """Test that .test() and .submit() calls are removed from extracted code."""
    monkeypatch.setenv('ZEST_SERVER_URL', server_url)
    
    code = '''def my_solution(nums, target):
    return [0, 1]

my_solution.test()
my_solution.submit()
my_solution.test(no_cache=True)
my_solution.submit(skip_local=True)
'''
    
    def test_func(nums, target):
        return [0, 1]
    
    q = Question(1, func=test_func)
    
    # Mock the extraction
    import types
    import inspect
    mock_module = types.ModuleType('test_module')
    
    def mock_getsource(obj):
        if obj == mock_module:
            return code
        return code
    
    monkeypatch.setattr(inspect, 'getsource', mock_getsource)
    monkeypatch.setattr(inspect, 'getmodule', lambda x: mock_module)
    
    extracted = q._extract_user_code()
    
    # Verify test/submit calls are removed
    assert '.test(' not in extracted
    assert '.submit(' not in extracted
    assert 'def my_solution' in extracted


def test_helper_functions_preserved(temp_test_file, monkeypatch, server_url):
    """Test that helper functions are preserved in extraction."""
    monkeypatch.setenv('ZEST_SERVER_URL', server_url)
    
    code = '''def helper(x):
    return x * 2

def my_solution(nums, target):
    result = helper(5)
    return [0, result]
'''
    
    def test_func(nums, target):
        return [0, 1]
    
    q = Question(1, func=test_func)
    
    # Mock the extraction
    import types
    import inspect
    mock_module = types.ModuleType('test_module')
    
    def mock_getsource(obj):
        if obj == mock_module:
            return code
        return code
    
    monkeypatch.setattr(inspect, 'getsource', mock_getsource)
    monkeypatch.setattr(inspect, 'getmodule', lambda x: mock_module)
    
    extracted = q._extract_user_code()
    
    # Verify both functions are preserved
    assert 'def helper' in extracted
    assert 'def my_solution' in extracted


def test_multiple_functions_preserved(temp_test_file, monkeypatch, server_url):
    """Test that multiple solution functions are preserved (multi-question support)."""
    monkeypatch.setenv('ZEST_SERVER_URL', server_url)
    
    code = '''def solution_q1(nums, target):
    return [0, 1]

def solution_q2(nums, target):
    return [1, 2]

def shared_helper(x):
    return x + 1
'''
    
    def test_func(nums, target):
        return [0, 1]
    
    q = Question(1, func=test_func)
    
    # Mock the extraction
    import types
    import inspect
    mock_module = types.ModuleType('test_module')
    
    def mock_getsource(obj):
        if obj == mock_module:
            return code
        return code
    
    monkeypatch.setattr(inspect, 'getsource', mock_getsource)
    monkeypatch.setattr(inspect, 'getmodule', lambda x: mock_module)
    
    extracted = q._extract_user_code()
    
    # Verify all functions are preserved
    assert 'def solution_q1' in extracted
    assert 'def solution_q2' in extracted
    assert 'def shared_helper' in extracted


def test_comments_preserved(temp_test_file, monkeypatch, server_url):
    """Test that comments are preserved in extraction."""
    monkeypatch.setenv('ZEST_SERVER_URL', server_url)
    
    code = '''# This is a comment
def my_solution(nums, target):
    # Another comment
    return [0, 1]  # Inline comment
'''
    
    def test_func(nums, target):
        return [0, 1]
    
    q = Question(1, func=test_func)
    
    # Mock the extraction
    import types
    import inspect
    mock_module = types.ModuleType('test_module')
    
    def mock_getsource(obj):
        if obj == mock_module:
            return code
        return code
    
    monkeypatch.setattr(inspect, 'getsource', mock_getsource)
    monkeypatch.setattr(inspect, 'getmodule', lambda x: mock_module)
    
    extracted = q._extract_user_code()
    
    # Verify comments are preserved
    assert '# This is a comment' in extracted
    assert '# Another comment' in extracted
    assert '# Inline comment' in extracted
