from math import prod

from libcst import Arg, Attribute, Call, CSTTransformer, Module, Name

from core.languages.python.cst_utils import evaluate_expression, make_number_node

# Functions where shape is passed as positional args (not as a tuple/list)
TORCH_SHAPE_AS_ARGS = frozenset(["rand", "randn", "ones", "zeros", "empty"])


class TensorLimit(CSTTransformer):
    def __init__(self, threshold: int) -> None:
        self.threshold = threshold

    def _calculate_tensor_size(self, shape: list[int]) -> int:
        """Calculates the total number of elements in a tensor."""
        return prod(shape) if shape else 0

    def _should_modify_tensor(self, shape: list[int]) -> tuple[bool, list[int]]:
        """Determines if a tensor should be modified based on its size and the threshold."""
        if not shape:
            return False, shape

        tensor_size = self._calculate_tensor_size(shape)

        if tensor_size > self.threshold:
            scaling_factor = (tensor_size // self.threshold) ** (1 / len(shape))
            new_shape = [max(int(dim // scaling_factor), 1) for dim in shape]
            return True, new_shape

        return False, shape

    def _extract_positional_shape_args(self, args: list[Arg]) -> tuple[list[int], list[Arg]]:
        """Extract shape values from positional args, separating them from keyword args.

        Returns (shape_values, keyword_args) where keyword_args are args with keywords.
        """
        shape_values = []
        keyword_args = []

        for arg in args:
            if arg.keyword is not None:
                # This is a keyword argument (dtype=, device=, etc.)
                keyword_args.append(arg)
            else:
                # This is a positional argument (shape dimension)
                arg_value = evaluate_expression(arg.value)
                if arg_value is not None:
                    shape_values.append(arg_value)
                else:
                    # Can't evaluate this arg, abort shape extraction
                    return [], list(args)

        return shape_values, keyword_args

    def leave_Call(self, original_node: Call, updated_node: Call) -> Call:
        """Modifies torch tensor creation calls if they exceed the threshold."""
        # Handle torch.{rand, randn, ones, zeros, empty}() calls
        if (
            isinstance(original_node.func, Attribute)
            and isinstance(original_node.func.value, Name)
            and original_node.func.value.value == "torch"
            and original_node.func.attr.value in TORCH_SHAPE_AS_ARGS
        ):
            # Extract shape arguments (positional) and preserve keyword arguments
            shape_values, keyword_args = self._extract_positional_shape_args(list(original_node.args))

            # If we successfully extracted shape dimensions
            if shape_values:
                should_modify, new_shape = self._should_modify_tensor(shape_values)

                if should_modify:
                    # Create new positional arguments with the modified shape values
                    new_positional_args = [Arg(value=make_number_node(dim)) for dim in new_shape]
                    # Combine new shape args with preserved keyword args (dtype, device, generator, etc.)
                    new_args = new_positional_args + keyword_args
                    # Return a modified call with the new arguments
                    return updated_node.with_changes(args=new_args)

        return updated_node


def modify_tensors(module: Module, threshold: int = 10000000) -> Module:
    """Applies the TensorLimit transformer to modify calls to torch tensor creation functions."""
    transformer = TensorLimit(threshold)
    return module.visit(transformer)
