import collections
from collections.abc import Callable
from typing import Any, Optional

import torch
from torch._dynamo.variables.dicts import ConstDictVariable
from torch._dynamo.variables.lists import TupleVariable
from torch.fx import has_side_effect, Proxy

from .. import graph_break_hints
from ..bytecode_transformation import create_call_function
from ..exc import TYPE_CHECKING, unimplemented
from ..graph_bytecode_inputs import (
    CURRENT_STREAM_INDEX,
    get_external_object_by_index,
    register_graph_created_object,
    register_user_object,
    reset_user_object_tracking,
)
from ..source import CurrentStreamSource
from .base import VariableTracker
from .constant import ConstantVariable
from .ctx_manager import FxTracebackAnnotateVariable
from .lazy import LazyVariableTracker


if TYPE_CHECKING:
    from torch._dynamo.symbolic_convert import InstructionTranslator

    from ..codegen import PyCodegen

from torch._library.custom_ops import custom_op


Tensor = torch.Tensor


def new_event(*args: Any, **kwargs: Any) -> int:
    event = torch.Event(*args, **kwargs)
    return register_graph_created_object(
        event,
        EventVariable.make_construct_in_graph_event_fn(
            TupleVariable([]), ConstDictVariable({})
        ),
    )


def new_stream(*args: tuple[Any], **kwargs: Any) -> int:
    stream = torch.Stream(*args, **kwargs)  # type: ignore[no-matching-overload,call-overload]
    return register_graph_created_object(
        stream,
        StreamVariable.make_construct_in_graph_stream_fn(
            TupleVariable([]), ConstDictVariable({})
        ),
    )


def _codegen_current_stream(device: torch.device, cg: "PyCodegen") -> None:
    cg.add_push_null(
        lambda: cg.load_import_from(
            torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
            "stash_graph_created_object",
        )
    )
    cg(CurrentStreamSource(device))
    cg.extend_output(create_call_function(1, False))


def get_current_stream(device: torch.device) -> int:
    stream = torch.accelerator.current_stream(device)
    return register_graph_created_object(
        stream, lambda _, cg: _codegen_current_stream(device, cg)
    )


def _get_stream_by_index(index: int) -> torch.Stream:
    stream = get_external_object_by_index(index)
    assert isinstance(stream, torch.Stream), (
        f"Fork/join stream expected a stream object at index {index}"
    )
    return stream


def _get_event_by_index(index: int) -> torch.Event:
    event = get_external_object_by_index(index)
    assert isinstance(event, torch.Event), (
        f"Record/wait event expected an event object at index {index}"
    )
    return event


@custom_op("streams::fork", mutates_args=())
def fork_stream(
    from_index: int,  # kept to make stream transitions clearer
    to_index: int,
) -> None:
    torch.accelerator.set_stream(_get_stream_by_index(to_index))


@fork_stream.register_fake
def _(
    from_index: int,  # kept to make stream transitions clearer
    to_index: int,
) -> None:
    pass


has_side_effect(torch.ops.streams.fork.default)


@custom_op("streams::join", mutates_args=())
def join_stream(from_index: int, to_index: int) -> None:
    torch.accelerator.set_stream(_get_stream_by_index(to_index))


@join_stream.register_fake
def _(
    from_index: int,
    to_index: int,
) -> None:
    pass


has_side_effect(torch.ops.streams.join.default)


@custom_op("streams::record_event", mutates_args=())
def record_event(event_index: int, stream_index: int) -> None:
    event = _get_event_by_index(event_index)
    stream = _get_stream_by_index(stream_index)
    event.record(stream)


@record_event.register_fake
def _(
    event_index: int,
    stream_index: int,
) -> None:
    pass


has_side_effect(torch.ops.streams.record_event.default)


@custom_op("streams::wait_event", mutates_args=())
def wait_event(event_index: int, stream_index: int) -> None:
    event = _get_event_by_index(event_index)
    stream = _get_stream_by_index(stream_index)
    event.wait(stream)


@wait_event.register_fake
def _(
    event_index: int,
    stream_index: int,
) -> None:
    pass


has_side_effect(torch.ops.streams.wait_event.default)


@custom_op("streams::synchronize_event", mutates_args=())
def synchronize_event(event_index: int) -> None:
    event = _get_event_by_index(event_index)
    event.synchronize()


@synchronize_event.register_fake
def _(event_index: int) -> None:
    pass


has_side_effect(torch.ops.streams.synchronize_event.default)


@custom_op("streams::synchronize_device", mutates_args=())
def synchronize_device(device_type: str, device_index: int) -> None:
    torch.accelerator.synchronize(torch.device(device_type, device_index))


@synchronize_device.register_fake
def _(device_type: str, device_index: int) -> None:
    pass


has_side_effect(torch.ops.streams.synchronize_device.default)


@custom_op("streams::synchronize_stream", mutates_args=())
def synchronize_stream(stream_index: int) -> None:
    stream = _get_stream_by_index(stream_index)
    stream.synchronize()


@synchronize_stream.register_fake
def _(stream_index: int) -> None:
    pass


has_side_effect(torch.ops.streams.synchronize_stream.default)


@custom_op("streams::wait_stream", mutates_args=())
def wait_stream(waiting_stream_index: int, waited_on_stream_index: int) -> None:
    waiting = _get_stream_by_index(waiting_stream_index)
    waited_on = _get_stream_by_index(waited_on_stream_index)
    waiting.wait_stream(waited_on)


@wait_stream.register_fake
def _(
    event_index: int,
    stream_index: int,
) -> None:
    pass


has_side_effect(torch.ops.streams.wait_stream.default)


@custom_op("streams::sync_dealloc", mutates_args=())
def sync_dealloc(
    wait_event_index: int, src_stream_index: int, to_dealloc: torch.Tensor
) -> None:
    """An op which waits on an event and moves the last usage of to_dealloc
    after the wait, so that after the sync occurs, the deallocation or
    subsequent reuse of the tensor's memory will be guaranteed to happen
    after a side stream is finished using it.
    See https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html#torch.Tensor.record_stream
    for more details"""
    torch.ops.streams.wait_event.default(wait_event_index, src_stream_index)


@sync_dealloc.register_fake
def _(
    wait_event_index: int,
    src_stream_index: int,
    to_dealloc: torch.Tensor,
) -> None:
    pass


has_side_effect(torch.ops.streams.sync_dealloc.default)


@custom_op("streams::record_stream", mutates_args=())
def record_stream(tensor: torch.Tensor, stream_index: int) -> None:
    tensor.record_stream(_get_stream_by_index(stream_index))


@record_stream.register_fake
def _(
    tensor: torch.Tensor,
    stream_index: int,
) -> None:
    pass


has_side_effect(torch.ops.streams.record_stream.default)


class SymbolicStreamState:
    """Track the currently entered stream if any"""

    def __init__(self) -> None:
        from ..source import CurrentStreamSource

        cur_stack: list[StreamVariable] = []
        if torch.accelerator.is_available():
            # Reset the registry so the current stream is guaranteed index 0.
            reset_user_object_tracking()
            stream = torch.accelerator.current_stream()
            source = CurrentStreamSource(stream.device)
            # Register the current stream so it gets index 0 (registry is
            # fresh at tracing start).  The inductor wrapper updates this
            # entry at runtime so cudagraph capture uses the capture stream
            # instead of this stale trace-time stream.
            index = register_user_object(stream, source)
            assert index == CURRENT_STREAM_INDEX, (
                f"Current stream must be registered at index {CURRENT_STREAM_INDEX}, "
                f"got {index}"
            )
            stream_var = LazyVariableTracker.create(stream, source=source)
            # Set user_object_index as an instance attribute so accessing it
            # does NOT trigger LazyVariableTracker realization.
            stream_var.user_object_index = index  # type: ignore[union-attr]
            cur_stack = [stream_var]  # type: ignore[list-item]

        self.cur_stream_stack: collections.deque[StreamVariable] = collections.deque(
            cur_stack
        )

    def enter_stream(self, stream: "StreamVariable") -> None:
        self.cur_stream_stack.append(stream)

    def exit_stream(self) -> None:
        self.cur_stream_stack.pop()

    def cur_stream(self, device: torch.device | None = None) -> "StreamVariable":
        if device is not None:
            for stream in reversed(self.cur_stream_stack):
                if stream.device == device:
                    return stream

        return self.cur_stream_stack[-1]

    def in_stream_context(self) -> bool:
        return len(self.cur_stream_stack) > 0

    def cur_stream_id(self) -> int:
        """Get a Python object id for the current stream without realizing lazy variables."""
        stream = self.cur_stream_stack[-1]
        if isinstance(stream, LazyVariableTracker) and not stream.is_realized():
            return id(stream.peek_value())
        return id(stream.value)


class StreamContextVariable(FxTracebackAnnotateVariable):
    """This represents torch.cuda.StreamContext"""

    @staticmethod
    def create(
        tx: "InstructionTranslator",
        stream_to_enter: "StreamVariable",
        **kwargs: dict[str, Any],
    ) -> "StreamContextVariable":
        return StreamContextVariable(
            stream_to_enter,
            **kwargs,
        )

    def __init__(self, stream: Optional["StreamVariable"], **kwargs: Any) -> None:
        self.stream = stream
        super().__init__(
            target_values={"stream": self.get_stream().user_object_index},
            initial_values=None,
            **kwargs,
        )

    def enter(
        self, tx: "InstructionTranslator", *args: VariableTracker
    ) -> VariableTracker:
        # to stream, from stream is the order of the arguments
        # we are entering the target, and leaving the initial stream
        tx.symbolic_stream_state.enter_stream(self.get_stream())
        return super().enter(tx)

    def exit(
        self, tx: "InstructionTranslator", *args: VariableTracker
    ) -> VariableTracker:
        # to stream, from stream is the order of the arguments
        # we are leaving the target, and entering the initial stream
        tx.symbolic_stream_state.exit_stream()
        return super().exit(tx, *args)

    def python_type(self) -> type:
        return torch.cuda.StreamContext

    def supports_graph_breaks(self) -> bool:
        return True

    def get_stream(self) -> "StreamVariable":
        assert self.stream, "Stream context should have a separate stream"
        return self.stream


class StreamVariable(StreamContextVariable):
    """Represents the device-agnostic torch.Stream class"""

    _cpython_type = torch.Stream

    def __init__(
        self,
        proxy: Proxy,
        value: torch.Stream,
        user_object_index: int | None = None,
        **kwargs: Any,
    ) -> None:
        # Index into the user object table
        # used to pass arbitrary objects to the graph
        if proxy is not None and "example_value" in proxy.node.meta:
            assert proxy.node.meta["example_value"] == value

        self.proxy = proxy
        self.value = value
        self.device = value.device

        self.user_object_index = user_object_index
        super().__init__(None, **kwargs)

    def python_type(self) -> type:
        return torch.Stream

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

    def call_method(
        self,
        tx: "InstructionTranslator",
        name: str,
        args: list[VariableTracker],
        kwargs: dict[str, VariableTracker],
    ) -> VariableTracker:
        assert hasattr(self.value, name), f"no stream method found named {name}"

        from ..utils import cmp_name_to_op_mapping, proxy_args_kwargs
        from .builder import wrap_fx_proxy_cls

        if name == "wait_event":
            event_arg = args[0]
            assert isinstance(event_arg, EventVariable)
            tx.output.create_proxy(
                "call_function",
                torch.ops.streams.wait_event,
                (event_arg.user_object_index, self.user_object_index),
                {},
            )
            return ConstantVariable.create(None)
        elif name == "wait_stream":
            other_stream = args[0]
            assert isinstance(other_stream, StreamVariable)
            tx.output.create_proxy(
                "call_function",
                torch.ops.streams.wait_stream,
                (self.user_object_index, other_stream.user_object_index),
                {},
            )
            return ConstantVariable.create(None)
        elif name == "synchronize":
            tx.output.create_proxy(
                "call_function",
                torch.ops.streams.synchronize_stream,
                (self.user_object_index,),
                {},
            )
            return ConstantVariable.create(None)
        elif name == "query":
            return wrap_fx_proxy_cls(
                target_cls=ConstantVariable,
                tx=tx,
                proxy=tx.output.create_proxy(
                    "call_method", name, *proxy_args_kwargs([self] + args, kwargs)
                ),
            )
        elif name == "record_event":
            from .builder import wrap_fx_proxy

            tx.output.check_event_record_after_input_mutation(id(self.value))
            if args and isinstance(args[0], EventVariable):
                event_var = args[0]
                event = event_var.value
                event_index = event_var.user_object_index
            else:
                event = self.value.record_event()
                event_index = register_graph_created_object(
                    event,
                    EventVariable.make_construct_in_graph_event_fn(
                        TupleVariable([]), ConstDictVariable({})
                    ),
                )
            tx.output.create_proxy(
                "call_function",
                torch.ops.streams.record_event,
                (event_index, self.user_object_index),
                {},
            )
            return wrap_fx_proxy(
                tx=tx,
                proxy=tx.output.create_proxy(
                    "call_function",
                    get_external_object_by_index,
                    (event_index,),
                    {},
                ),
            )
        elif name in cmp_name_to_op_mapping and len(args) == 1 and not kwargs:
            from ..guards import GuardBuilder, install_guard

            if self.source:
                install_guard(self.source.make_guard(GuardBuilder.EQUALS_MATCH))

            # NB : Checking for mutation is necessary because we compare
            # constant values
            other = args[0]
            if not isinstance(other, StreamVariable):
                return VariableTracker.build(tx, NotImplemented)

            if other.source:
                assert self.source is not None
                install_guard(self.source.make_guard(GuardBuilder.EQUALS_MATCH))
            return VariableTracker.build(
                tx,
                cmp_name_to_op_mapping[name](self.value, other.value),  # type: ignore[arg-type]
            )

        return super().call_method(tx, name, args, kwargs)

    def as_proxy(self) -> Proxy:
        return self.proxy

    def module_name(self) -> str:
        return "torch._C"

    def fn_name(self) -> str:
        return "Stream"

    def reconstruct(self, codegen: "PyCodegen") -> None:
        # If we got here, this stream is fully subsumed by the graph - this means it is
        # not an input or global
        assert not self.source
        if self.user_object_index is not None:
            codegen.add_push_null(
                lambda: codegen.load_import_from(
                    torch._dynamo.graph_bytecode_inputs.__name__,
                    "get_external_object_by_index",
                )
            )
            codegen.append_output(codegen.create_load_const(self.user_object_index))
            codegen.extend_output(create_call_function(1, False))
        else:
            # This will support the legacy behavior
            prefix = f"_stream_{self.device}"
            name = codegen.tx.output.install_global_by_id(prefix, self.value)
            codegen.append_output(codegen.create_load_global(name, add=True))

    def get_stream(self) -> "StreamVariable":
        return self

    @staticmethod
    def make_construct_in_graph_stream_fn(
        args: TupleVariable, kwargs: ConstDictVariable
    ) -> Callable[[int, "PyCodegen"], None]:
        def fn(index: int, codegen: "PyCodegen") -> None:
            codegen.add_push_null(
                lambda: codegen.load_import_from(
                    torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
                    "stash_graph_created_object",
                )
            )
            codegen.add_push_null(
                lambda: codegen.load_import_from(
                    torch._dynamo.utils.__name__, "build_stream"
                )
            )
            codegen(args)
            codegen(kwargs)
            codegen.extend_output(create_call_function(2, False))
            codegen.extend_output(create_call_function(1, False))

        return fn


class CudaStreamVariable(StreamVariable):
    """Represents torch.cuda.Stream, preserving device-specific type and attributes."""

    _cpython_type = torch.cuda.Stream

    def python_type(self) -> type:
        return torch.cuda.Stream

    def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracker":
        from . import ConstantVariable

        if name == "cuda_stream":
            from ..guards import GuardBuilder, install_guard

            if self.source:
                install_guard(self.source.make_guard(GuardBuilder.EQUALS_MATCH))

            if hasattr(self.value, "cuda_stream"):
                return ConstantVariable.create(self.value.cuda_stream)

            if hasattr(self.value, "native_handle"):
                return ConstantVariable.create(self.value.native_handle)

        return super().var_getattr(tx, name)


class EventVariable(VariableTracker):
    def __init__(
        self,
        proxy: Proxy,
        value: torch.Event,
        user_object_index: int | None,
        **kwargs: Any,
    ) -> None:
        if proxy is not None and "example_value" in proxy.node.meta:
            assert proxy.node.meta["example_value"] == value
        super().__init__(**kwargs)
        self.proxy = proxy
        self.value = value
        self.user_object_index = user_object_index

    def python_type(self) -> type:
        return torch.Event

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

    def call_method(
        self,
        tx: "InstructionTranslator",
        name: str,
        args: list[VariableTracker],
        kwargs: dict[str, VariableTracker],
    ) -> VariableTracker:
        from ..utils import proxy_args_kwargs
        from .builder import wrap_fx_proxy_cls

        if name == "wait":
            _, stream_index = EventVariable._get_stream_arg(tx, args, kwargs)
            tx.output.create_proxy(
                "call_function",
                torch.ops.streams.wait_event,
                (
                    self.user_object_index,
                    stream_index,
                ),
                {},
            )
            return ConstantVariable.create(None)
        elif name == "record":
            stream_arg, stream_index = EventVariable._get_stream_arg(tx, args, kwargs)
            tx.output.check_event_record_after_input_mutation(id(stream_arg.value))
            tx.output.create_proxy(
                "call_function",
                torch.ops.streams.record_event,
                (
                    self.user_object_index,
                    stream_index,
                ),
                {},
            )
            return ConstantVariable.create(None)
        elif name == "synchronize":
            tx.output.create_proxy(
                "call_function",
                torch.ops.streams.synchronize_event,
                (self.user_object_index,),
                {},
            )
            return ConstantVariable.create(None)
        elif name == "query":
            return wrap_fx_proxy_cls(
                target_cls=ConstantVariable,
                tx=tx,
                proxy=tx.output.create_proxy(
                    "call_method", name, *proxy_args_kwargs([self] + args, kwargs)
                ),
            )
        else:
            method_name = (
                f"{type(self.value).__module__}.{type(self.value).__qualname__}.{name}"
            )
            unimplemented(
                gb_type="Unsupported event method",
                context=str(name),
                explanation=f"Dynamo doesn't support tracing the {method_name} method. "
                f"We currently support wait, record, synchronize, and query.",
                hints=[
                    *graph_break_hints.SUPPORTABLE,
                ],
            )

    def as_proxy(self) -> Proxy:
        return self.proxy

    @staticmethod
    def _get_stream_arg(
        tx: "InstructionTranslator",
        args: list[VariableTracker],
        kwargs: dict[str, VariableTracker],
    ) -> tuple["StreamVariable", int]:
        """Returns (stream_variable, stream_index_for_op).

        The ambient current stream is registered at index 0 in the external
        object registry.  The inductor wrapper updates index 0 at runtime so
        that cudagraph capture sees the capture stream, not the stale
        trace-time default stream.
        """
        stream_arg = None
        if args:
            stream_arg = args[0]
        elif kwargs:
            stream_arg = kwargs.get("stream")

        if not stream_arg:
            stream_var = tx.symbolic_stream_state.cur_stream()
            return stream_var, stream_var.user_object_index  # type: ignore[return-value]

        return stream_arg, stream_arg.user_object_index  # type: ignore[return-value]

    @staticmethod
    def make_construct_in_graph_event_fn(
        args: TupleVariable, kwargs: ConstDictVariable
    ) -> Callable[[int, "PyCodegen"], None]:
        def fn(index: int, codegen: "PyCodegen") -> None:
            codegen.add_push_null(
                lambda: codegen.load_import_from(
                    torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
                    "stash_graph_created_object",
                )
            )
            codegen.add_push_null(
                lambda: codegen.load_import_from(
                    torch._dynamo.utils.__name__, "build_event"
                )
            )
            codegen(args)
            codegen(kwargs)
            codegen.extend_output(create_call_function(2, False))
            codegen.extend_output(create_call_function(1, False))

        return fn

    def reconstruct(self, codegen: "PyCodegen") -> None:
        # If we got here, this event is fully subsumed by the graph - this means it is
        # not an input or global
        assert not self.source
        # Similar to stream handling, we lift the event into a global and then codegen bytecode to load it from there.
        prefix = "_event"
        name = codegen.tx.output.install_global_by_id(prefix, self.value)
        codegen.append_output(codegen.create_load_global(name, add=True))
