"""
This module provides iterator-related variable tracking functionality for Dynamo.
It implements variable classes for handling Python iterators and itertools functions
during symbolic execution and tracing.

The module includes:
- Base iterator variable classes for tracking iterator state
- Implementations of built-in iterators (zip, map, filter)
- Support for itertools functions (product, accumulate, combinations, etc.)
- Mutation tracking and reconstruction capabilities for iterator operations

These classes integrate with Dynamo's variable tracking system to enable proper
handling of iterator operations during code transformation and optimization.
"""

import itertools
import sys
from typing import Any, TYPE_CHECKING

from .. import graph_break_hints, polyfills, variables
from ..bytecode_transformation import (
    create_call_function,
    create_call_function_ex,
    create_instruction,
)
from ..exc import (
    handle_observed_exception,
    ObservedUserStopIteration,
    raise_observed_exception,
    raise_value_error,
    unimplemented,
)
from ..utils import unpack_iterable
from .base import ValueMutationNew, VariableTracker
from .constant import ConstantVariable
from .hashable import HashableTracker
from .object_protocol import generic_iternext


if TYPE_CHECKING:
    from torch._dynamo.codegen import PyCodegen
    from torch._dynamo.symbolic_convert import InstructionTranslatorBase


MAX_ITERATOR_LIMIT = 100 * 1024  # 100k


def is_iterator_exhausted(
    tx: "InstructionTranslatorBase", iterator: VariableTracker
) -> bool:
    try:
        generic_iternext(tx, iterator)
        return False
    except ObservedUserStopIteration:
        handle_observed_exception(tx)
        return True


class ItertoolsVariable(VariableTracker):
    def __init__(self, value: Any, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.value = value

    def richcompare_impl(
        self, tx: "InstructionTranslatorBase", other: VariableTracker, op: str
    ) -> VariableTracker:
        from .object_protocol import python_constant_richcompare_impl

        return python_constant_richcompare_impl(self, tx, other, op)

    def __repr__(self) -> str:
        return f"ItertoolsVariable({self.value})"

    def as_python_constant(self) -> Any:
        return self.value

    def get_real_python_backed_value(self) -> Any:
        return self.value

    def call_function(
        self,
        tx: "InstructionTranslatorBase",
        args: list["VariableTracker"],
        kwargs: "dict[str, VariableTracker]",
    ) -> "VariableTracker":
        # See also: module `torch._dynamo.polyfills.itertools`

        if self.value is itertools.product:
            if any(kw != "repeat" for kw in kwargs):
                unimplemented(
                    gb_type="Unsupported kwargs for itertools.product",
                    context=f"call_function {self} {args} {kwargs}",
                    explanation=f"Expected kwargs: 'repeat', but got "
                    f"{','.join(set(kwargs.keys()) - {'repeat'})}",
                    hints=[*graph_break_hints.USER_ERROR],
                )

            if "repeat" in kwargs:
                r = kwargs["repeat"].as_python_constant()
            else:
                r = 1
            seqs = [unpack_iterable(tx, arg) for arg in args]
            items = [
                variables.TupleVariable(list(item))
                for item in itertools.product(*seqs, repeat=r)
            ]
            return variables.ListIteratorVariable(
                items,  # type: ignore[arg-type]
                mutation_type=ValueMutationNew(),
            )
        elif (
            self.value is itertools.combinations
            and not kwargs
            and len(args) == 2
            and args[1].is_python_constant()
        ):
            iterable = unpack_iterable(tx, args[0])
            r = args[1].as_python_constant()

            items = []
            for item in itertools.combinations(iterable, r):
                items.append(variables.TupleVariable(list(item)))
            return variables.ListIteratorVariable(
                items,  # type: ignore[arg-type]
                mutation_type=ValueMutationNew(),
            )
        elif self.value is itertools.groupby:
            if any(kw != "key" for kw in kwargs):
                unimplemented(
                    gb_type="Unsupported kwargs for itertools.groupby",
                    context=f"call_function {self} {args} {kwargs}",
                    explanation=f"Expected kwargs: 'key', but got "
                    f"{','.join(set(kwargs.keys()) - {'key'})}",
                    hints=[*graph_break_hints.USER_ERROR],
                )

            def retrieve_const_key(key: VariableTracker) -> Any:
                from ..utils import specialize_symnode

                # Unwrap LazyVariableTracker to get the underlying variable
                key = specialize_symnode(key)
                if isinstance(key, variables.SymNodeVariable):
                    return key.evaluate_expr()
                elif key.is_python_constant():
                    return key.as_python_constant()
                else:
                    unimplemented(
                        gb_type="Unsupported key type for itertools.groupby",
                        context=f"call_function {self} {args} {kwargs}",
                        explanation="Dynamo does not know how to trace "
                        f"itertools.groupby with key type: {str(type(key))}. "
                        "We only support grouping keys that are constants (int, float, str, etc.)",
                        hints=[*graph_break_hints.SUPPORTABLE],
                    )

            if len(args) != 1:
                unimplemented(
                    gb_type="Unsupported arguments for itertools.groupby",
                    context=f"call_function {self} {args} {kwargs}",
                    explanation="Dynamo does not know how to trace "
                    f"itertools.groupby with args: {args} and kwargs: {kwargs}. "
                    "itertools.groupby expects an iterable to group and an "
                    "optional key function to determine groupings.",
                    hints=[
                        "Make sure the arguments to itertools.groupby are correct.",
                        *graph_break_hints.SUPPORTABLE,
                    ],
                )
            seq = unpack_iterable(tx, args[0])

            if "key" in kwargs:

                def keyfunc(x: VariableTracker) -> Any:
                    return retrieve_const_key(
                        kwargs.get("key").call_function(tx, [x], {})  # type: ignore[union-attr]
                    )

            else:

                def keyfunc(x: VariableTracker) -> Any:
                    return retrieve_const_key(x)

            result = []
            try:
                for k, v in itertools.groupby(seq, key=keyfunc):
                    result.append(
                        variables.TupleVariable(
                            [
                                (
                                    variables.ConstantVariable.create(k)
                                    if variables.ConstantVariable.is_literal(k)
                                    else k
                                ),
                                variables.ListIteratorVariable(
                                    list(v), mutation_type=ValueMutationNew()
                                ),
                            ],
                            mutation_type=ValueMutationNew(),
                        )
                    )
            except Exception as e:
                unimplemented(
                    gb_type="Unexpected failure during itertools.groupby() iteration",
                    context=f"call_function {self} {args} {kwargs}",
                    explanation="Unexpected failure in invoking function during groupby",
                    hints=[*graph_break_hints.SUPPORTABLE],
                    from_exc=e,
                )
            return variables.ListIteratorVariable(
                result,  # type: ignore[arg-type]
                mutation_type=ValueMutationNew(),
            )
        elif self.value is itertools.repeat:
            if len(args) < 2:
                return RepeatIteratorVariable(*args, mutation_type=ValueMutationNew())

            return tx.inline_user_function_return(
                VariableTracker.build(tx, polyfills.repeat), args, kwargs
            )
        elif self.value is itertools.count and not kwargs:
            if len(args) == 0:
                return variables.CountIteratorVariable(mutation_type=ValueMutationNew())
            if len(args) == 1:
                return variables.CountIteratorVariable(
                    item=args[0], mutation_type=ValueMutationNew()
                )
            if len(args) == 2:
                return variables.CountIteratorVariable(
                    item=args[0],
                    step=args[1],
                    mutation_type=ValueMutationNew(),
                )
            return super().call_function(tx, args, kwargs)
        elif (
            self.value is itertools.permutations
            and (len(args) == 1 or (len(args) == 2 and args[1].is_python_constant()))
            and not kwargs
        ):
            if len(args) == 2:
                r = args[1].as_python_constant()
            else:
                r = None
            items = [
                variables.TupleVariable(list(item))
                for item in itertools.permutations(unpack_iterable(tx, args[0]), r)
            ]
            return variables.ListIteratorVariable(
                items,  # type: ignore[arg-type]
                mutation_type=ValueMutationNew(),
            )
        else:
            return super().call_function(tx, args, kwargs)


class IteratorVariable(VariableTracker):
    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)

    def richcompare_impl(
        self, tx: "InstructionTranslatorBase", other: VariableTracker, op: str
    ) -> VariableTracker:
        from .object_protocol import object_richcompare

        return object_richcompare(self, tx, other, op)

    def tp_iternext_impl(self, tx: "InstructionTranslatorBase") -> VariableTracker:
        unimplemented(
            gb_type="Unimplemented next() call",
            context=f"next({self})",
            explanation="This abstract method must be implemented",
            hints=[*graph_break_hints.DYNAMO_BUG],
        )

    def call_obj_hasattr(
        self, tx: "InstructionTranslatorBase", name: str
    ) -> "ConstantVariable":
        if name == "__iter__" or name == "__next__":
            return variables.ConstantVariable.create(True)
        return super().call_obj_hasattr(tx, name)

    def tp_iter_impl(self, tx: "InstructionTranslatorBase") -> "VariableTracker":
        """Iterators are their own iterator."""
        return self


class RepeatIteratorVariable(IteratorVariable):
    def __init__(self, item: VariableTracker, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.item = item

    def python_type(self) -> type:
        return itertools.repeat

    # Repeat needs no mutation, clone self
    def tp_iternext_impl(self, tx: "InstructionTranslatorBase") -> VariableTracker:
        # ref: https://github.com/python/cpython/blob/3.13/Modules/itertoolsmodule.c#L4332-L4340
        # TODO(dynamo-team): Missing `times` argument handling
        return self.item

    def reconstruct(self, codegen: "PyCodegen") -> None:
        codegen.add_push_null(
            lambda: codegen.extend_output(
                [
                    codegen.create_load_python_module(itertools),
                    codegen.create_load_attr("repeat"),
                ]
            )
        )
        codegen(self.item)
        codegen.extend_output(create_call_function(1, False))


class CountIteratorVariable(IteratorVariable):
    # advance_count tracks how many next() calls were made during tracing,
    # used by side_effects.py to replay them on the real iterator post-execution.
    _nonvar_fields = {
        "advance_count",
        *IteratorVariable._nonvar_fields,
    }

    def python_type(self) -> type:
        return itertools.count

    def __init__(
        self,
        item: int | VariableTracker = 0,
        step: int | VariableTracker = 1,
        advance_count: int = 0,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        if not isinstance(item, VariableTracker):
            item = ConstantVariable.create(item)
        if not isinstance(step, VariableTracker):
            step = ConstantVariable.create(step)
        self.item = item
        self.step = step
        self.advance_count = advance_count

    def tp_iternext_impl(self, tx: "InstructionTranslatorBase") -> VariableTracker:
        # ref: https://github.com/python/cpython/blob/3.13/Modules/itertoolsmodule.c#L4189-L4216
        if not self.is_mutable():
            raise AssertionError("CountIteratorVariable must be mutable for next()")
        old_item = self.item
        tx.output.side_effects.mutation(self)
        self.item = self.item.call_method(tx, "__add__", [self.step], {})
        self.advance_count += 1
        return old_item

    def reconstruct(self, codegen: "PyCodegen") -> None:
        codegen.add_push_null(
            lambda: codegen.extend_output(
                [
                    codegen.create_load_python_module(itertools),
                    codegen.create_load_attr("count"),
                ]
            )
        )
        codegen(self.item)
        codegen(self.step)
        codegen.extend_output(create_call_function(2, False))


class ZipVariable(IteratorVariable):
    """
    Represents zip(*iterables)
    """

    # PyZip_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L3011
    _cpython_type = zip

    _nonvar_fields = {
        "strict",
        *IteratorVariable._nonvar_fields,
    }

    def __init__(
        self,
        iterable: VariableTracker,
        strict: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        if not isinstance(iterable, variables.TupleVariable):
            raise AssertionError(f"Expected a tuple of iterables, got {type(iterable)}")
        # can be list[Variable] or VariableTracker (with next_variable implemented)
        self.iterable = iterable
        self.strict = strict

    def python_type(self) -> type[zip]:  # type: ignore[type-arg]
        return zip

    def tp_iternext_impl(self, tx: "InstructionTranslatorBase") -> VariableTracker:
        # ref: https://github.com/python/cpython/blob/v3.13.3/Python/bltinmodule.c#L2906-L2994
        if not self.is_mutable():
            raise AssertionError("ZipVariable must be mutable for next()")
        tuplesize = len(self.iterable.items)

        if tuplesize == 0:
            raise_observed_exception(StopIteration, tx)

        items = []
        for i in range(tuplesize):
            it = self.iterable.items[i]
            try:
                items.append(generic_iternext(tx, it))
            except ObservedUserStopIteration:
                if not self.strict:
                    raise

                if i > 0:
                    raise_value_error(
                        tx, f"zip() argument {i} shorter than previous arguments"
                    )

                # In strict mode, if any iterable is exhausted, all must be exhausted
                for j in range(i + 1, tuplesize):
                    it_j = self.iterable.items[j]
                    if is_iterator_exhausted(tx, it_j):
                        continue
                    break
                else:
                    # all iterables exhausted, raise StopIteration
                    raise

                handle_observed_exception(tx)  # StopIteration
                raise_value_error(
                    tx, f"zip() argument {i} is longer than previous arguments"
                )

        return variables.TupleVariable(items)

    def reconstruct(self, codegen: "PyCodegen") -> None:
        codegen.add_push_null(
            lambda: codegen.load_import_from("builtins", "zip"), call_function_ex=True
        )
        codegen(self.iterable)
        codegen.extend_output(
            [
                codegen.create_load_const("strict"),
                codegen.create_load_const(self.strict),
                create_instruction("BUILD_MAP", arg=1),
                *create_call_function_ex(True, False),
            ]
        )


class MapVariable(IteratorVariable):
    """
    Represents map(fn, *iterables)
    """

    # PyMap_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L1484
    _cpython_type = map

    _nonvar_fields = {
        "strict",
        *IteratorVariable._nonvar_fields,
    }

    def __init__(
        self,
        fn: VariableTracker,
        iterables: VariableTracker,
        strict: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.fn = fn
        if not isinstance(iterables, variables.TupleVariable):
            raise AssertionError(
                f"Expected a tuple of iterables, got {type(iterables)}"
            )
        self.iterable = iterables
        self.strict = strict

    def python_type(self) -> type:
        return map

    def tp_iternext_impl(self, tx: "InstructionTranslatorBase") -> VariableTracker:
        # ref: https://github.com/python/cpython/blob/v3.13.3/Python/bltinmodule.c#L1409-L1450
        if not self.is_mutable():
            raise AssertionError("MapVariable must be mutable for next()")

        tuplesize = len(self.iterable.items)

        items = []
        for i in range(tuplesize):
            it = self.iterable.items[i]
            try:
                items.append(generic_iternext(tx, it))
            except ObservedUserStopIteration:
                if not self.strict:
                    raise

                handle_observed_exception(tx)  # StopIteration

                if i:
                    raise_value_error(
                        tx,
                        f"map() argument {i + 1} shorter than argument {i}",
                    )

                # In strict mode, if any iterable is exhausted, all must be exhausted.
                for j in range(1, tuplesize):
                    it_j = self.iterable.items[j]
                    if not is_iterator_exhausted(tx, it_j):
                        raise_value_error(
                            tx, f"map() argument {j + 1} is longer than argument {j}"
                        )

                raise_observed_exception(StopIteration, tx)

        # type: ignore[attr-defined]
        return self.fn.call_function(tx, items, {})

    def reconstruct(self, codegen: "PyCodegen") -> None:
        codegen.add_push_null(
            lambda: codegen.load_import_from("builtins", "map"), call_function_ex=True
        )
        codegen(variables.TupleVariable([self.fn, *self.iterable.items]))
        if self.strict:
            if sys.version_info < (3, 14):
                raise AssertionError(
                    "Unexpected bug: map(strict=True) requires Python 3.14+"
                )
            codegen.extend_output(
                [
                    codegen.create_load_const("strict"),
                    codegen.create_load_const(self.strict),
                    create_instruction("BUILD_MAP", arg=1),
                    *create_call_function_ex(True, False),
                ]
            )
        else:
            codegen.extend_output(create_call_function_ex(False, False))


class FilterVariable(IteratorVariable):
    """
    Represents filter(fn, iterable)
    """

    # PyFilter_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L630
    _cpython_type = filter

    def __init__(
        self,
        fn: VariableTracker,
        iterable: VariableTracker,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.fn = fn
        self.iterable = iterable

    def python_type(self) -> type:
        return filter

    def tp_iternext_impl(self, tx: "InstructionTranslatorBase") -> VariableTracker:
        # ref: https://github.com/python/cpython/blob/v3.13.3/Python/bltinmodule.c#L573-L606
        # A do-while loop to find elements that make fn return true
        while True:
            item = generic_iternext(tx, self.iterable)
            if self.fn.is_constant_none():
                res = item
            else:
                res = self.fn.call_function(tx, [item], {})
            pred_res = variables.UserFunctionVariable(
                polyfills.predicate  # type: ignore[arg-type]
            ).call_function(tx, [res], {})
            if pred_res.as_python_constant():
                return item

    def reconstruct(self, codegen: "PyCodegen") -> None:
        codegen.add_push_null(lambda: codegen.load_import_from("builtins", "filter"))
        codegen(self.fn)
        codegen(self.iterable)
        codegen.extend_output(create_call_function(2, False))


class DictViewIterator(IteratorVariable):
    """Base class for dict view iterators (keys, values, or items)."""

    _nonvar_fields = {
        "view_type",
        *IteratorVariable._nonvar_fields,
    }

    view_type: str = "keys"

    def __init__(
        self,
        items: dict[HashableTracker, VariableTracker],
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        if self.view_type == "keys":
            self._iter = iter(items)
        elif self.view_type == "values":
            self._iter = iter(items.values())  # type: ignore[bad-assignment]
        else:
            if self.view_type != "items":
                raise AssertionError(
                    f"Expected view_type 'items', got {self.view_type!r}"
                )
            self._iter = iter(items.items())  # type: ignore[bad-assignment]

    def tp_iternext_impl(self, tx: "InstructionTranslatorBase") -> VariableTracker:
        # dictiter_iternextitem: https://github.com/python/cpython/blob/v3.13.3/Objects/dictobject.c#L5538-L5578
        # dictiter_iternextkey: https://github.com/python/cpython/blob/v3.13.3/Objects/dictobject.c#L5125-L5144
        # dictiter_iternextvalue: https://github.com/python/cpython/blob/v3.13.3/Objects/dictobject.c#L5248-L5267
        try:
            item = next(self._iter)

            if self.view_type == "keys":
                return item.vt
            elif self.view_type == "values":
                return item  # type: ignore[bad-return]
            else:  # items
                k, v = item  # type: ignore[not-iterable]
                return VariableTracker.build(tx, (k.vt, v))
        except (RuntimeError, StopIteration) as e:
            raise_observed_exception(
                type(e),
                tx,
                args=[VariableTracker.build(tx, a) for a in e.args],
            )

    def python_type(self) -> type:
        if self.view_type == "keys":
            return type(iter({}))
        elif self.view_type == "values":
            return type(iter({}.values()))
        else:  # items
            return type(iter({}.items()))


class DictIterator(DictViewIterator):
    _cpython_type = type(iter({}))
    view_type = "keys"


class DictKeysIterator(DictViewIterator):
    _cpython_type = type(iter({}.keys()))
    view_type = "keys"


class DictValuesIterator(DictViewIterator):
    _cpython_type = type(iter({}.values()))
    view_type = "values"


class DictItemsIterator(DictViewIterator):
    _cpython_type = type(iter({}.items()))
    view_type = "items"


class SetIterator(DictViewIterator):
    _cpython_type = type(iter(set()))
    view_type = "keys"