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

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


class RangeModifierTransformer(CSTTransformer):
    def __init__(self, threshold):
        self.threshold = threshold
        self.in_random_sample = False

    def _calculate_range_size(self, start: int, stop: int, step: int) -> int:
        """Calculates the number of iterations a range will perform."""
        if step == 0:
            return 0
        if (step > 0 and start >= stop) or (step < 0 and start <= stop):
            return 0
        return abs((stop - start + (step - 1 if step > 0 else step + 1)) // step)

    def _should_modify_range(self, start: int, stop: int, step: int = 1) -> tuple[bool, int, int, int]:
        """Determines if a range should be modified based on its size and the threshold."""
        range_size = self._calculate_range_size(start, stop, step)

        if range_size > self.threshold:
            if step > 0:
                new_stop = start + (step * self.threshold)
                return True, start, new_stop, step
            new_stop = start + (step * self.threshold)
            return True, start, new_stop, step

        return False, start, stop, step

    def visit_Call(self, node: Call) -> None:
        # Check if this is a random.sample call
        if (
            isinstance(node.func, Attribute)
            and isinstance(node.func.value, Name)
            and node.func.value.value == "random"
            and node.func.attr.value == "sample"
        ):
            self.in_random_sample = True

    def leave_Call(self, original_node: Call, updated_node: Call) -> Call:
        """Modifies both range() calls and random.sample() calls if they exceed the threshold."""
        # Handle random.sample calls
        if (
            isinstance(original_node.func, Attribute)
            and isinstance(original_node.func.value, Name)
            and original_node.func.value.value == "random"
            and original_node.func.attr.value == "sample"
        ):
            self.in_random_sample = False

            if len(original_node.args) == 2:
                # Get the sample size (second argument)
                sample_size = evaluate_expression(original_node.args[1].value)
                if sample_size is not None and sample_size > self.threshold:
                    # Modify the sample size to match threshold
                    new_args = list(updated_node.args)
                    new_args[1] = Arg(value=make_number_node(self.threshold))
                    return updated_node.with_changes(args=new_args)
            return updated_node

        # Handle range() calls
        if isinstance(original_node.func, Name) and original_node.func.value == "range":
            if len(original_node.args) == 1:
                # range(stop)
                stop_val = evaluate_expression(original_node.args[0].value)
                if stop_val is not None:
                    if abs(stop_val) > self.threshold:
                        new_stop = self.threshold if stop_val > 0 else -self.threshold
                        new_arg = Arg(value=make_number_node(new_stop))
                        return updated_node.with_changes(args=[new_arg])

            elif len(original_node.args) == 2:
                # range(start, stop)
                start_val = evaluate_expression(original_node.args[0].value)
                stop_val = evaluate_expression(original_node.args[1].value)

                if start_val is not None and stop_val is not None:
                    range_size = abs(stop_val - start_val)
                    if range_size > self.threshold:
                        step = 1 if stop_val > start_val else -1
                        new_stop = start_val + (self.threshold * step)
                        new_args = [Arg(value=make_number_node(start_val)), Arg(value=make_number_node(new_stop))]
                        return updated_node.with_changes(args=new_args)

            elif len(original_node.args) == 3:
                # range(start, stop, step)
                start_val = evaluate_expression(original_node.args[0].value)
                stop_val = evaluate_expression(original_node.args[1].value)
                step_val = evaluate_expression(original_node.args[2].value)

                if start_val is None or stop_val is None or step_val is None:
                    return updated_node

                range_size = self._calculate_range_size(start_val, stop_val, step_val)
                if range_size > self.threshold:
                    new_stop = start_val + (step_val * self.threshold)
                    new_args = [
                        Arg(value=make_number_node(start_val)),
                        Arg(value=make_number_node(new_stop)),
                        Arg(value=make_number_node(step_val)),
                    ]
                    return updated_node.with_changes(args=new_args)

        return updated_node


def modify_large_loops(module: Module, threshold: int = 1000) -> Module:
    """Applies the RangeModifierTransformer to modify calls to the `range` function."""
    transformer = RangeModifierTransformer(threshold)
    return module.visit(transformer)
