"""Shared AST utilities for test instrumentation.

Provides call-position matching and import-alias detection used by
both sync and async instrumentation paths.
"""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING

from .._model import (
    FunctionParent,
    FunctionToOptimize,
)

if TYPE_CHECKING:
    from ..test_discovery.models import CodePosition

log = logging.getLogger(__name__)


def node_in_call_position(
    node: ast.AST, call_positions: list[CodePosition]
) -> bool:
    """Return True if the AST node overlaps any of the given call positions.

    When *call_positions* is empty, every ``ast.Call`` node matches.
    """
    if not call_positions:
        return isinstance(node, ast.Call)
    if isinstance(node, ast.Call):
        node_lineno = getattr(node, "lineno", None)
        node_col_offset = getattr(node, "col_offset", None)
        node_end_lineno = getattr(node, "end_lineno", None)
        node_end_col_offset = getattr(node, "end_col_offset", None)
        if (
            node_lineno is not None
            and node_col_offset is not None
            and node_end_lineno is not None
        ):
            for pos in call_positions:
                pos_line = pos.line_no
                if (
                    pos_line is not None
                    and node_lineno <= pos_line <= node_end_lineno
                ):
                    if (
                        pos_line == node_lineno
                        and node_col_offset <= pos.col_no
                    ):
                        return True
                    if (
                        pos_line == node_end_lineno
                        and node_end_col_offset is not None
                        and node_end_col_offset >= pos.col_no
                    ):
                        return True
                    if node_lineno < pos_line < node_end_lineno:
                        return True
    return False


class FunctionImportedAsVisitor(ast.NodeVisitor):
    """Check if a function was imported as an alias.

    from numpy import array as np_array
    np_array is what we want
    """

    def __init__(self, function: FunctionToOptimize) -> None:
        """Initialize with the target function to look for import aliases."""
        assert len(function.parents) <= 1, (  # noqa: S101
            "Only support functions with one or less parent"
        )
        self.imported_as: FunctionToOptimize = function
        self.function = function
        if function.parents:
            self.to_match = function.parents[0].name
        else:
            self.to_match = function.function_name

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        """Detect import aliases for the target function."""
        for alias in node.names:
            if (
                alias.name == self.to_match
                and hasattr(alias, "asname")
                and alias.asname is not None
            ):
                if self.function.parents:
                    self.imported_as = FunctionToOptimize(
                        function_name=self.function.function_name,
                        parents=(FunctionParent(alias.asname, "ClassDef"),),
                        file_path=self.function.file_path,
                        starting_line=self.function.starting_line,
                        ending_line=self.function.ending_line,
                        is_async=self.function.is_async,
                    )
                else:
                    self.imported_as = FunctionToOptimize(
                        function_name=alias.asname,
                        parents=(),
                        file_path=self.function.file_path,
                        starting_line=self.function.starting_line,
                        ending_line=self.function.ending_line,
                        is_async=self.function.is_async,
                    )


def detect_frameworks_from_code(
    code: str,
) -> dict[str, str]:
    """Detect GPU/device frameworks used in code.

    Analyzes imports for torch, tensorflow, and jax.

    Returns:
        A dictionary mapping framework names to their import aliases.
        For example: {"torch": "th", "tensorflow": "tf", "jax": "jax"}

    """
    frameworks: dict[str, str] = {}
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return frameworks

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                module_name = alias.name.split(".")[0]
                if module_name == "torch":
                    frameworks["torch"] = alias.asname or module_name
                elif module_name == "tensorflow":
                    frameworks["tensorflow"] = alias.asname or module_name
                elif module_name == "jax":
                    frameworks["jax"] = alias.asname or module_name
        elif isinstance(node, ast.ImportFrom) and node.module:
            module_name = node.module.split(".")[0]
            if module_name == "torch" and "torch" not in frameworks:
                frameworks["torch"] = module_name
            elif (
                module_name == "tensorflow" and "tensorflow" not in frameworks
            ):
                frameworks["tensorflow"] = module_name
            elif module_name == "jax" and "jax" not in frameworks:
                frameworks["jax"] = module_name

    return frameworks
