from math import prod

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

from aiservice.common.cst_utils import evaluate_expression, make_number_node


class TensorLimit(CSTTransformer):
    def __init__(self, threshold) -> 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)

    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."""
        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 leave_Call(self, original_node: Call, updated_node: Call) -> Call:
        """Modifies torch tensor creation calls if they exceed the threshold."""
        # Handle torch.{rand, ones, zeros}() 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 ["rand", "ones", "zeros"]
        ):
            # Extract shape arguments
            shape_values = []
            for arg in original_node.args:
                arg_value = evaluate_expression(arg.value)
                if arg_value is not None:
                    shape_values.append(arg_value)

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

                if should_modify:
                    # Create new arguments with the modified shape values
                    new_args = [Arg(value=make_number_node(dim)) for dim in new_shape]
                    # 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)
