import ast

from libcst import parse_module

from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize
from core.languages.python.testgen.instrumentation.edit_generated_test import replace_definition_with_import


def test_replace_definition_with_import_top_level_function() -> None:
    """Test that top-level functions are replaced with imports correctly."""
    source_code = '''
import pytest

async def clone_agent_template(request, template_id):
    """Function to be replaced"""
    return "test"

def test_something():
    pass
'''

    function = FunctionToOptimize(
        function_name="clone_agent_template",
        file_path="/home/ubuntu/work/repo/backend/python/app/api/routes/agent.py",
        parents=[],
    )
    module_path = "app.api.routes.agent"

    result = replace_definition_with_import(parse_module(source_code), function, module_path).code

    expected = """
import pytest
from app.api.routes.agent import clone_agent_template

def test_something():
    pass
"""

    assert result.strip() == expected.strip()
    ast.parse(result)


def test_replace_definition_with_import_ignores_class_methods() -> None:
    source_code = '''
import pytest
from fastapi import HTTPException

async def clone_agent_template(request, template_id):
    """Function to be replaced"""
    return "test"

class MockArangoService:
    async def clone_agent_template(self, template_id):
        """This method should NOT be replaced"""
        return "mocked"

class FaultyArangoService:
    async def clone_agent_template(self, template_id):
        """This method should also NOT be replaced"""
        raise Exception("error")

def test_something():
    pass
'''

    function = FunctionToOptimize(
        function_name="clone_agent_template",
        file_path="/home/ubuntu/work/repo/backend/python/app/api/routes/agent.py",
        parents=[],
    )
    module_path = "app.api.routes.agent"

    result = replace_definition_with_import(parse_module(source_code), function, module_path).code

    expected = '''
import pytest
from fastapi import HTTPException
from app.api.routes.agent import clone_agent_template

class MockArangoService:
    async def clone_agent_template(self, template_id):
        """This method should NOT be replaced"""
        return "mocked"

class FaultyArangoService:
    async def clone_agent_template(self, template_id):
        """This method should also NOT be replaced"""
        raise Exception("error")

def test_something():
    pass
'''

    assert result.strip() == expected.strip()
    ast.parse(result)


def test_replace_definition_with_import_multiple_top_level_occurrences() -> None:
    source_code = '''
import pytest

# First occurrence - should be replaced
async def clone_agent_template(request, template_id):
    """First function to be replaced"""
    return "test1"

def some_other_function():
    pass

# Second occurrence - should also be replaced if it's top-level
async def clone_agent_template(request, template_id):
    """Second function - this is a redefinition"""
    return "test2"
'''

    function = FunctionToOptimize(
        function_name="clone_agent_template",
        file_path="/home/ubuntu/work/repo/backend/python/app/api/routes/agent.py",
        parents=[],
    )
    module_path = "app.api.routes.agent"

    result = replace_definition_with_import(parse_module(source_code), function, module_path).code

    expected = '''
import pytest
from app.api.routes.agent import clone_agent_template

def some_other_function():
    pass

# Second occurrence - should also be replaced if it's top-level
async def clone_agent_template(request, template_id):
    """Second function - this is a redefinition"""
    return "test2"
'''

    assert result.strip() == expected.strip()
    ast.parse(result)


def test_replace_definition_with_import_nested_classes() -> None:
    source_code = '''
import pytest

async def clone_agent_template(request, template_id):
    """Top-level function to be replaced"""
    return "test"

class OuterClass:
    class InnerClass:
        async def clone_agent_template(self, template_id):
            """Nested class method should not be replaced"""
            return "nested"

    async def clone_agent_template(self, template_id):
        """Outer class method should not be replaced"""
        return "outer"
'''

    function = FunctionToOptimize(
        function_name="clone_agent_template",
        file_path="/home/ubuntu/work/repo/backend/python/app/api/routes/agent.py",
        parents=[],
    )
    module_path = "app.api.routes.agent"

    result = replace_definition_with_import(parse_module(source_code), function, module_path).code

    expected = '''
import pytest
from app.api.routes.agent import clone_agent_template

class OuterClass:
    class InnerClass:
        async def clone_agent_template(self, template_id):
            """Nested class method should not be replaced"""
            return "nested"

    async def clone_agent_template(self, template_id):
        """Outer class method should not be replaced"""
        return "outer"
'''

    assert result.strip() == expected.strip()
    ast.parse(result)


def test_replace_definition_with_import_class_definition() -> None:
    """Test that top-level class definitions can be replaced."""
    source_code = """
import pytest

class SomeAgent:
    def __init__(self):
        pass

    def method(self):
        pass

def test_something():
    pass
"""

    function = FunctionToOptimize(
        function_name="SomeAgent",
        file_path="/home/ubuntu/work/repo/backend/python/app/api/routes/agent.py",
        parents=[FunctionParent(name="SomeAgent", type="ClassDef")],
    )
    module_path = "app.api.routes.agent"

    result = replace_definition_with_import(parse_module(source_code), function, module_path).code

    expected = """
import pytest
from app.api.routes.agent import SomeAgent

def test_something():
    pass
"""

    assert result.strip() == expected.strip()
    ast.parse(result)


def test_replace_definition_with_import_no_match() -> None:
    """Test that nothing is replaced when target doesn't exist."""
    source_code = """
import pytest

def some_function():
    pass

class SomeClass:
    def method(self):
        pass
"""

    function = FunctionToOptimize(
        function_name="nonexistent_function",
        file_path="/home/ubuntu/work/repo/backend/python/app/api/routes/agent.py",
        parents=[],
    )
    module_path = "app.api.routes.agent"

    result = replace_definition_with_import(parse_module(source_code), function, module_path).code

    expected = """
import pytest

def some_function():
    pass

class SomeClass:
    def method(self):
        pass
"""

    assert result.strip() == expected.strip()
    ast.parse(result)


def test_regression_syntax_error_from_import_in_class() -> None:
    source_code = '''
import pytest
from fastapi import HTTPException

async def clone_agent_template(request, template_id):
    """Function under test, simplified for testing."""
    return "test"

class FaultyArangoService:
    async def clone_agent_template(self, template_id):
        raise RuntimeError("Clone error")

class MockLogger:
    def __init__(self):
        self.errors = []

    def error(self, msg, exc_info=False):
        self.errors.append(msg)
'''

    function = FunctionToOptimize(
        function_name="clone_agent_template",
        file_path="/home/ubuntu/work/repo/backend/python/app/api/routes/agent.py",
        parents=[],
    )
    module_path = "app.api.routes.agent"

    result = replace_definition_with_import(parse_module(source_code), function, module_path).code

    expected = """
import pytest
from fastapi import HTTPException
from app.api.routes.agent import clone_agent_template

class FaultyArangoService:
    async def clone_agent_template(self, template_id):
        raise RuntimeError("Clone error")

class MockLogger:
    def __init__(self):
        self.errors = []

    def error(self, msg, exc_info=False):
        self.errors.append(msg)
"""

    assert result.strip() == expected.strip()
    ast.parse(result)
