"""AST visitors for side-effect detection."""

import ast
from typing import Set

from guardrail.classifier.constants import (
    DB_LIBS,
    FILE_IO_FUNCTIONS,
    NETWORK_LIBS,
    SIDE_EFFECT_INDICATORS,
    SYSTEM_CALLS,
)


class SideEffectVisitor(ast.NodeVisitor):
    """AST visitor to detect side effects in function code."""

    def __init__(self, is_side_effect_module_func):
        """
        Initialize the visitor.

        Args:
            is_side_effect_module_func: Function to check if a module has side effects.
        """
        self._is_side_effect_module = is_side_effect_module_func
        self.has_side_effects = False
        self.reason = ""

    def visit_Call(self, node: ast.Call) -> None:
        """Visit a function call node."""
        # Check if it's a direct function call
        if isinstance(node.func, ast.Name):
            func_name = node.func.id
            if self._is_side_effect_call(func_name):
                self.has_side_effects = True
                self.reason = f"Calls side-effect function: {func_name}"
                return

        # Check if it's an attribute call (e.g., requests.get, open())
        if isinstance(node.func, ast.Attribute):
            # Check for open() calls
            if isinstance(node.func.value, ast.Name):
                if node.func.value.id == "open":
                    self.has_side_effects = True
                    self.reason = "Calls open() for file I/O"
                    return

            # Check for module.attribute calls
            if isinstance(node.func.value, ast.Name):
                module_name = node.func.value.id
                attr_name = node.func.attr

                # Check for subprocess calls
                if module_name == "subprocess":
                    self.has_side_effects = True
                    self.reason = f"Uses subprocess: {module_name}.{attr_name}"
                    return

                # Check for os.system, os.popen, etc.
                if module_name == "os" and attr_name in ["system", "popen", "exec", "spawn"]:
                    self.has_side_effects = True
                    self.reason = f"Uses os.{attr_name}"
                    return

                # Check for socket operations
                if module_name == "socket":
                    self.has_side_effects = True
                    self.reason = "Uses socket operations"
                    return

                # Check for database operations
                if module_name in DB_LIBS:
                    self.has_side_effects = True
                    self.reason = f"Uses database library: {module_name}"
                    return

                # Check for requests/httpx/aiohttp calls
                if module_name in NETWORK_LIBS:
                    self.has_side_effects = True
                    self.reason = f"Uses network library: {module_name}"
                    return

        self.generic_visit(node)

    def visit_Import(self, node: ast.Import) -> None:
        """Visit an import statement."""
        for alias in node.names:
            module_name = alias.name.split(".")[0]
            if self._is_side_effect_module(module_name):
                self.has_side_effects = True
                self.reason = f"Imports side-effect module: {module_name}"
                return
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        """Visit a from-import statement."""
        if node.module:
            module_name = node.module.split(".")[0]
            if self._is_side_effect_module(module_name):
                self.has_side_effects = True
                self.reason = f"Imports from side-effect module: {module_name}"
                return
        self.generic_visit(node)

    def _is_side_effect_call(self, func_name: str) -> bool:
        """Check if a function name indicates a side effect."""
        return func_name in (FILE_IO_FUNCTIONS | SIDE_EFFECT_INDICATORS)


class ImportCollector(ast.NodeVisitor):
    """AST visitor to collect import information."""

    def __init__(self):
        """Initialize the collector."""
        self.modules: Set[str] = set()
        self.names: Set[str] = set()

    def visit_Import(self, node: ast.Import) -> None:
        """Visit an import statement."""
        for alias in node.names:
            self.modules.add(alias.name.split(".")[0])
            self.names.add(alias.asname or alias.name.split(".")[-1])

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        """Visit a from-import statement."""
        if node.module:
            self.modules.add(node.module.split(".")[0])
            for alias in node.names:
                self.names.add(alias.asname or alias.name)

