# Copyright 2025 Ant Group Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Interpreter: Execute Graph IR and Eager Operations.

Interpreter is a Context that executes operations immediately.
It can execute both:
1. Graph IR (via GraphInterpreter)
2. Eager operations on InterpObject (via backend executors)
"""

from __future__ import annotations

import collections
import concurrent.futures
import contextlib
import hashlib
import json
import os
import pathlib
import queue
import threading
import time
from collections.abc import Callable, Iterator
from typing import TYPE_CHECKING, Any, cast

from mplang.edsl.context import AbstractInterpreter
from mplang.edsl.graph import Graph, Value
from mplang.edsl.object import Object
from mplang.edsl.registry import get_impl
from mplang.edsl.typing import BaseType
from mplang.runtime.dialect_state import DialectState
from mplang.runtime.object_store import ObjectStore
from mplang.utils.logging import get_logger

if TYPE_CHECKING:
    from mplang.edsl.primitive import Primitive

logger = get_logger(__name__)


class ExecutionTracer:
    """Tracer for DAG execution events (Chrome Tracing format).

    This tracer captures:
    1. Op-level execution events (schedule, start, end)
    2. Communication statistics (via registered CommStats)
    3. Custom events (e.g., comm.send, comm.recv)

    The output is compatible with Chrome's chrome://tracing viewer.

    Integration with CommStats:
        >>> tracer = ExecutionTracer(enabled=True, trace_dir="./traces")
        >>> comm = HttpCommunicator(..., tracer=tracer)
        >>> tracer.register_comm_stats("rank_0", comm.stats)
        >>> # ... run computation ...
        >>> tracer.stop()  # Includes comm stats in output
    """

    def __init__(self, enabled: bool = False, *, trace_dir: str | pathlib.Path):
        self.enabled = enabled
        self.start_time = 0.0
        self.end_time = 0.0
        self.active_tasks_samples: list[tuple[float, int]] = []
        self.queue_size_samples: list[tuple[float, int]] = []
        self.completed_ops = 0
        self.total_ops = 0
        self.trace_dir = pathlib.Path(trace_dir)

        # Tracing
        self.trace_events: list[dict[str, Any]] = []
        self.op_schedule_times: dict[
            tuple[int, Any], float
        ] = {}  # (id(op), namespace) -> ts (us)
        self.pid = os.getpid()

        # Communication statistics registry: name -> CommStats-like object
        # Objects must have a snapshot() method returning dict[str, Any]
        self._comm_stats: dict[str, Any] = {}

    def register_comm_stats(self, name: str, stats: Any) -> None:
        """Register a CommStats object for inclusion in trace output.

        The stats object must have a snapshot() method that returns a dict.

        Args:
            name: Identifier for this stats source (e.g., "rank_0", "worker_1")
            stats: CommStats-like object with snapshot() method
        """
        self._comm_stats[name] = stats

    def start(self) -> None:
        self.start_time = time.time()

    def stop(self, filename_prefix: str = "dag_trace") -> None:
        self.end_time = time.time()
        self.save_trace(filename_prefix)

    def sample(self, active_tasks: int, queue_size: int) -> None:
        now = time.time() - self.start_time
        self.active_tasks_samples.append((now, active_tasks))
        self.queue_size_samples.append((now, queue_size))

    def log_schedule(self, op: Any, namespace: Any = None) -> None:
        if not self.enabled:
            return
        key = (id(op), namespace)
        self.op_schedule_times[key] = time.time() * 1e6

    def log_start(
        self, op: Any, pid: int | None = None, namespace: Any = None
    ) -> float:
        if not self.enabled:
            return 0.0
        start_ts = time.time() * 1e6
        if pid is None:
            pid = self.pid

        # Record scheduling latency (Queue Time)
        key = (id(op), namespace)
        if key in self.op_schedule_times:
            sched_ts = self.op_schedule_times.pop(key)
            self.trace_events.append({
                "name": f"Queue: {op.opcode}",
                "cat": "scheduler",
                "ph": "X",
                "ts": sched_ts,
                "dur": start_ts - sched_ts,
                "pid": pid,
                "tid": "SchedulerQueue",
            })
        return start_ts

    def log_end(self, op: Any, start_ts: float, pid: int | None = None) -> None:
        if not self.enabled:
            return
        end_ts = time.time() * 1e6
        tid = threading.get_ident()
        if pid is None:
            pid = self.pid

        self.trace_events.append({
            "name": op.opcode,
            "cat": "op",
            "ph": "X",
            "ts": start_ts,
            "dur": end_ts - start_ts,
            "pid": pid,
            "tid": tid,
            "args": {
                "opcode": op.opcode,
            },
        })

    def log_custom_event(
        self,
        name: str,
        start_ts: float,
        end_ts: float,
        cat: str = "custom",
        args: dict[str, Any] | None = None,
    ) -> None:
        """Log a custom event with explicit start/end timestamps (in seconds)."""
        if not self.enabled:
            return
        tid = threading.get_ident()

        # Convert to microseconds
        ts_us = start_ts * 1e6
        dur_us = (end_ts - start_ts) * 1e6

        self.trace_events.append({
            "name": name,
            "cat": cat,
            "ph": "X",
            "ts": ts_us,
            "dur": dur_us,
            "pid": self.pid,
            "tid": tid,
            "args": args or {},
        })

    def save_trace(
        self,
        filename_prefix: str = "dag_trace",
        job_id: str | None = None,
        rank: int | None = None,
    ) -> None:
        if not self.enabled or not self.trace_events:
            return
        try:
            if len(self.trace_events) < 100:
                return  # Skip small graphs

            # Use unique filename to avoid overwriting
            if job_id:
                # Format: trace_<job_id>_rank_<rank>.json
                rank_str = f"_rank_{rank}" if rank is not None else ""
                filename = f"trace_{job_id}{rank_str}.json"
            else:
                timestamp = int(time.time() * 1000)
                tid = threading.get_ident()
                filename = f"{filename_prefix}_{timestamp}_{tid}.json"

            # Save trace to trace_dir
            self.trace_dir.mkdir(parents=True, exist_ok=True)
            filepath = self.trace_dir / filename

            # Build trace output with optional metadata
            trace_output: dict[str, Any] = {"traceEvents": self.trace_events}

            # Include communication statistics as metadata
            if self._comm_stats:
                comm_snapshots = {}
                for name, stats in self._comm_stats.items():
                    if hasattr(stats, "snapshot") and callable(stats.snapshot):
                        comm_snapshots[name] = stats.snapshot()
                if comm_snapshots:
                    trace_output["metadata"] = {
                        "comm_stats": comm_snapshots,
                        "duration_s": self.end_time - self.start_time,
                        "total_ops": self.total_ops,
                    }

            with open(filepath, "w") as f:
                json.dump(trace_output, f)
            print(f"\n[Tracer] Trace saved to {filepath.absolute()}")
        except Exception as e:
            print(f"[Tracer] Failed to save trace: {e}")

    def print_summary(self) -> None:
        duration = self.end_time - self.start_time
        if duration <= 0:
            return

        avg_active = (
            sum(c for _, c in self.active_tasks_samples)
            / len(self.active_tasks_samples)
            if self.active_tasks_samples
            else 0
        )
        max_active = (
            max(c for _, c in self.active_tasks_samples)
            if self.active_tasks_samples
            else 0
        )
        avg_queue = (
            sum(c for _, c in self.queue_size_samples) / len(self.queue_size_samples)
            if self.queue_size_samples
            else 0
        )

        print("\n" + "=" * 80)
        print("DAG EXECUTION PROFILER")
        print("=" * 80)
        print(f"Total Duration: {duration:.3f}s")
        print(f"Total Ops:      {self.total_ops}")
        print(f"Throughput:     {self.total_ops / duration:.1f} ops/s")
        print("-" * 80)
        print(f"Active Tasks:   Avg={avg_active:.1f}, Max={max_active}")
        print(f"Ready Queue:    Avg={avg_queue:.1f}")

        # Print communication statistics if available
        if self._comm_stats:
            print("-" * 80)
            print("COMMUNICATION STATS:")
            for name, stats in self._comm_stats.items():
                if hasattr(stats, "snapshot") and callable(stats.snapshot):
                    snap = stats.snapshot()
                    print(f"  [{name}]")
                    print(
                        f"    Messages: sent={snap.get('messages_sent', 0)}, "
                        f"recv={snap.get('messages_received', 0)}"
                    )
                    bytes_sent = snap.get("bytes_sent", 0)
                    bytes_recv = snap.get("bytes_received", 0)
                    print(
                        f"    Bytes:    sent={self._format_bytes(bytes_sent)}, "
                        f"recv={self._format_bytes(bytes_recv)}"
                    )
                    print(
                        f"    Latency:  send_avg={snap.get('avg_send_latency_ms', 0):.1f}ms, "
                        f"recv_wait_avg={snap.get('avg_recv_wait_time_ms', 0):.1f}ms"
                    )
                    errors = snap.get("send_errors", 0) + snap.get("recv_timeouts", 0)
                    if errors > 0:
                        print(
                            f"    Errors:   send_errors={snap.get('send_errors', 0)}, "
                            f"recv_timeouts={snap.get('recv_timeouts', 0)}"
                        )

        print("=" * 80 + "\n")

    @staticmethod
    def _format_bytes(num_bytes: int) -> str:
        """Format bytes as human-readable string."""
        if num_bytes < 1024:
            return f"{num_bytes}B"
        elif num_bytes < 1024 * 1024:
            return f"{num_bytes / 1024:.1f}KB"
        elif num_bytes < 1024 * 1024 * 1024:
            return f"{num_bytes / (1024 * 1024):.1f}MB"
        else:
            return f"{num_bytes / (1024 * 1024 * 1024):.1f}GB"


class _NullTracer:
    """No-op tracer stub for when tracing is disabled."""

    enabled = False
    total_ops = 0

    def log_schedule(self, op: Any, namespace: Any = None) -> None:
        pass

    def log_start(
        self, op: Any, pid: int | None = None, namespace: Any = None
    ) -> float:
        return 0.0

    def log_end(self, op: Any, start_ts: float, pid: int | None = None) -> None:
        pass

    def stop(self) -> None:
        pass

    def save_trace(self, **kwargs: Any) -> None:
        pass


class InterpObject(Object):
    """Interp-time object (during eager execution).

    Holds a runtime object (the actual data/handle owned by the backend executor)
    and a reference to the Interpreter (Context).
    Operations delegate to primitives which execute immediately.

    The runtime object can be:
    - FHE backend: Local TenSEAL/SEAL ciphertext
    - JAX backend: Local jax.Array
    - MP backend: Backend handle (pointer to party-side data)
    - SQL backend: DatabaseHandle
    - etc.

    Example:
        >>> # FHE backend (local execution)
        >>> x = fhe.encrypt([1, 2, 3])  # InterpObject with local ciphertext
        >>> y = fhe.encrypt([4, 5, 6])
        >>> z = x + y  # InterpObject.__add__ → add_p.bind(x, y)

        >>> # MP backend (distributed execution)
        >>> x = mp.random.uniform(shape=(10,))  # InterpObject with backend handle
        >>> y = mp.random.uniform(shape=(10,))
        >>> z = x + y  # InterpObject.__add__ → add_p.bind(x, y)
    """

    def __init__(
        self,
        runtime_obj: Any,
        obj_type: BaseType,
        interpreter: Interpreter | None = None,
    ):
        """Initialize InterpObject.

        Args:
            runtime_obj: Backend-specific runtime object (ciphertext, array, handle, etc.)
            obj_type: Type of the object (BaseType from edsl.typing)
            interpreter: Interpreter context (if None, uses default interpreter)
        """
        self._runtime_obj = runtime_obj
        self._type = obj_type
        self._context = interpreter  # InterpObject holds its Interpreter (Context)

    @property
    def type(self) -> BaseType:
        return self._type

    @property
    def runtime_obj(self) -> Any:
        """Get the underlying runtime object (backend-specific)."""
        return self._runtime_obj

    def __repr__(self) -> str:
        runtime_repr = repr(self._runtime_obj)
        # Truncate long representations
        if len(runtime_repr) > 50:
            runtime_repr = runtime_repr[:47] + "..."
        return f"InterpObject({runtime_repr}, type={self.type})"


class Interpreter(AbstractInterpreter):
    """Execution context for eager execution.

    Inherits from Context and implements bind_primitive() by executing immediately.

    Responsibilities:
    1. Execute primitives on InterpObject immediately
    2. Delegate to backend-specific executors
    3. Execute Graph IR (via GraphInterpreter)

    Example:
        >>> interp = Interpreter()
        >>> x = InterpObject(np.array([1, 2, 3]), Tensor[f32, (3,)])
        >>> y = InterpObject(np.array([4, 5, 6]), Tensor[f32, (3,)])
        >>> z = x + y  # InterpObject.__add__ → add_p.bind(x, y)
    """

    def __init__(
        self,
        executor: concurrent.futures.Executor | None = None,
        name: str = "Interpreter",
        tracer: ExecutionTracer | None = None,
        trace_pid: int | None = None,
        store: ObjectStore | None = None,
        root_dir: str | pathlib.Path | None = None,
        handlers: dict[str, Callable[..., Any]] | None = None,
    ) -> None:
        # Persistence Root
        self.root_dir = (
            pathlib.Path(root_dir)
            if root_dir
            else pathlib.Path(os.environ.get("MPLANG_DATA_ROOT", ".mpl"))
        )

        # Initialize Context base class (for state management)
        super().__init__()

        # Instance-level handler registry (overrides global registry)
        self.handlers: dict[str, Callable] = handlers or {}
        self.tracer = tracer

        logger.debug(
            "Initialized Interpreter '%s' with store=%s, tracer=%s",
            name,
            store is not None,
            tracer is not None,
        )

        # GraphValue -> InterpObject cache
        # Maps a GraphValue (IR node) to its computed InterpObject (Runtime result).
        # This serves two purposes:
        # 1. Caching: Avoid re-evaluating the same graph node multiple times.
        # 2. MIMO Optimization: When one output of a multi-output op is computed,
        #    all sibling outputs are cached here to avoid re-execution.
        self._execution_cache: dict[Any, InterpObject] = {}

        # -----------------------------------------------------------------
        # Graph-local op execution ids (for deterministic communication tags)
        # -----------------------------------------------------------------
        # We assign a monotonically increasing exec_id to each op execution
        # within a graph namespace, and keep it deterministic across parties.
        #
        # IMPORTANT:
        # - We intentionally make exec_id grow across repeated executions of the
        #   same region graph (e.g., while_loop iterations) to avoid tag/key reuse.
        #
        # Implementation:
        # - Each evaluate_graph(graph, ...) reserves a contiguous exec_id range
        #   [base, base + len(graph.operations)).
        # - Op exec_id = base + op_index_in_graph.
        # - Reservation is persisted per graph_exec_key (structural hash).
        # - We forbid concurrent execution of the same graph_hash to avoid
        #   message tag confusion when a backend uses only per-op tags.
        self._exec_id_lock = threading.Lock()
        self._graph_next_exec_base: dict[str, int] = {}
        self._active_graph_exec_keys: set[str] = set()
        self._tls = threading.local()
        self.executor = executor
        self.async_ops: set[str] = set()
        self.name = name
        self.trace_pid = trace_pid
        self.store: ObjectStore | None = store

    @contextlib.contextmanager
    def _tls_exec_context(
        self,
        *,
        graph_exec_key: str | None = None,
        op_exec_id: int | None = None,
    ) -> Iterator[None]:
        """Temporarily set execution context in thread-local storage."""

        prev_graph_key = getattr(self._tls, "current_graph_exec_key", None)
        prev_exec_id = getattr(self._tls, "current_op_exec_id", None)

        if graph_exec_key is not None:
            self._tls.current_graph_exec_key = graph_exec_key
        if op_exec_id is not None:
            self._tls.current_op_exec_id = op_exec_id

        try:
            yield
        finally:
            if graph_exec_key is not None:
                if prev_graph_key is None:
                    delattr(self._tls, "current_graph_exec_key")
                else:
                    self._tls.current_graph_exec_key = prev_graph_key

            if op_exec_id is not None:
                if prev_exec_id is None:
                    delattr(self._tls, "current_op_exec_id")
                else:
                    self._tls.current_op_exec_id = prev_exec_id

    def _graph_exec_key(self, graph: Graph) -> str:
        """Return a deterministic, structural hash for a graph.

        Used for:
        - Namespacing per-graph exec_id counters
        - Communication tag disambiguation (worker ops may include this key)

        Note: we cache on the Graph object assuming graphs are immutable during
        execution (finalized graphs / regions).
        """

        cached = getattr(graph, "_exec_key", None)
        if cached is not None:
            return cast(str, cached)

        # NOTE: We intentionally do NOT use graph.to_json() here.
        # graph.to_json() requires all attrs to be JSON-serializable via serde,
        # but graphs may legitimately contain runtime-only objects (e.g. JAX
        # PyTreeDef used by func.func). For communication tag namespaces we use
        # a simple structural fingerprint that is deterministic across parties.

        def _stable_attr_value(obj: Any) -> Any | None:
            """Return a JSON-compatible stable value or None if unsupported.

            We include only values that are likely deterministic across parties.
            Unknown runtime objects are skipped (e.g. PyTreeDef, callables, etc.).
            """

            if obj is None or isinstance(obj, (bool, int, float, str)):
                return obj

            if isinstance(obj, (bytes, bytearray, memoryview)):
                b = bytes(obj)
                return {
                    "_kind": "bytes",
                    "len": len(b),
                    "sha256": hashlib.sha256(b).hexdigest(),
                }

            try:
                import numpy as np  # type: ignore

                if isinstance(obj, np.ndarray):
                    b = obj.tobytes(order="C")
                    return {
                        "_kind": "ndarray",
                        "dtype": str(obj.dtype),
                        "shape": list(obj.shape),
                        "sha256": hashlib.sha256(b).hexdigest(),
                    }
                if isinstance(obj, (np.integer, np.floating)):
                    return obj.item()
            except Exception:
                pass

            if isinstance(obj, (list, tuple)):
                items: list[Any] = []
                for x in obj:
                    sx = _stable_attr_value(x)
                    if sx is None:
                        return None
                    items.append(sx)
                return items

            if isinstance(obj, dict):
                stable_items: list[tuple[Any, Any]] = []
                for k, v in obj.items():
                    sk = _stable_attr_value(k)
                    sv = _stable_attr_value(v)
                    if sk is None or sv is None:
                        return None
                    stable_items.append((sk, sv))
                stable_items.sort(
                    key=lambda kv: json.dumps(
                        kv[0], sort_keys=True, separators=(",", ":"), ensure_ascii=False
                    )
                )
                return {"_kind": "dict", "items": stable_items}

            return None

        def _graph_fingerprint(g: Graph) -> Any:
            # Map SSA Values to stable indices independent of their textual names.
            value_to_index: dict[Value, int] = {}

            def _index(v: Value) -> int:
                if v in value_to_index:
                    return value_to_index[v]
                value_to_index[v] = len(value_to_index)
                return value_to_index[v]

            for v in g.inputs:
                _index(v)
            for op in g.operations:
                for out in op.outputs:
                    _index(out)

            ops_fp: list[dict[str, Any]] = []
            for op in g.operations:
                attr_keys = sorted(op.attrs.keys())
                stable_attr_items: list[tuple[str, Any]] = []
                for k in attr_keys:
                    attr_val = op.attrs.get(k)
                    sv = _stable_attr_value(attr_val)
                    if sv is not None:
                        stable_attr_items.append((k, sv))

                ops_fp.append({
                    "opcode": op.opcode,
                    "inputs": [_index(v) for v in op.inputs],
                    "outputs": [str(v.type) for v in op.outputs],
                    "attrs": {"keys": attr_keys, "stable": stable_attr_items},
                    "regions": [_graph_fingerprint(r) for r in op.regions],
                })

            return {
                "inputs": [str(v.type) for v in g.inputs],
                "ops": ops_fp,
                "outputs": [_index(v) for v in g.outputs],
            }

        fingerprint = _graph_fingerprint(graph)

        payload = json.dumps(
            fingerprint,
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=False,
        ).encode("utf-8")
        key = hashlib.sha256(payload).hexdigest()

        # Store on graph to avoid id(graph) reuse pitfalls.
        graph._exec_key = key  # type: ignore[attr-defined]
        return key

    def shutdown(self) -> None:
        """Shutdown the interpreter and release resources.

        This method is idempotent and safe to call multiple times.
        It performs the following cleanup:
        1. Shuts down the internal executor (if any).
        2. Stops the execution tracer (if any).
        3. Shuts down any attached dialect states (e.g., stopping drivers).
        """
        logger.info("Shutting down Interpreter '%s'", self.name)

        # 1. Shutdown Executor
        if self.executor:
            logger.debug("Shutting down executor")
            self.executor.shutdown(wait=True)
            self.executor = None

        # 2. Stop Tracer
        if self.tracer:
            logger.debug("Stopping execution tracer")
            self.tracer.stop()
            # Don't clear self.tracer, as we might want to read stats later

        # 3. Shutdown Dialect States
        # Iterate over all attached states (e.g., drivers, cluster managers)
        # and shut them down if they support it.
        logger.debug("Shutting down dialect states")
        for state in self._states.values():
            if hasattr(state, "shutdown") and callable(state.shutdown):
                state.shutdown()

    # =========================================================================
    # Dialect State Management
    # =========================================================================
    def get_dialect_state(self, dialect: str) -> DialectState | None:
        """Get the state object for a specific dialect.

        This is a convenience wrapper around get_state("dialect.{dialect}").

        Args:
            dialect: Name of the dialect (e.g., "simp", "bfv", "spu")

        Returns:
            The dialect state object, or None if not set.

        Example:
            simp_state = interpreter.get_dialect_state("simp")
            if simp_state is not None:
                simp_state.submit(rank, graph, inputs)
        """
        state = self.get_state(f"dialect.{dialect}")
        # Type assertion: dialect states are always DialectState or None
        return cast(DialectState | None, state)

    def set_dialect_state(self, dialect: str, state: DialectState) -> None:
        """Set the state object for a specific dialect.

        This is a convenience wrapper around set_state("dialect.{dialect}", state).

        Args:
            dialect: Name of the dialect (e.g., "simp", "bfv", "spu")
            state: The dialect state object (should implement DialectState protocol)

        Example:
            interpreter.set_dialect_state("simp", cluster.connect())
        """
        self.set_state(f"dialect.{dialect}", state)

    def bind_primitive(
        self, primitive: Primitive, args: tuple[Any, ...], kwargs: dict[str, Any]
    ) -> InterpObject | list[InterpObject] | Any:
        """Execute primitive by tracing and interpreting.

        Implements the unified trace → interpret flow:
        1. All InterpObject arguments already registered via lift()
        2. Create a Tracer and push it as context
        3. Call primitive.bind() to build Graph IR (uses obj id in value names)
        4. Execute the graph via evaluate_graph() (resolves inputs via registry)

        Args:
            primitive: The primitive to execute
            args: Positional arguments (already lifted by Primitive.bind)
            kwargs: Keyword arguments (already lifted by Primitive.bind)

        Returns:
            Execution result (InterpObject or list of InterpObject or mixed with immediates)
        """
        from mplang.edsl.tracer import Tracer
        from mplang.utils import var_demorph, var_morph

        # Create tracer and build graph
        # Note: primitive.bind() internally calls Tracer.lift() with is_param=False,
        # so all args become captures (not params). This is correct because we're
        # tracing a primitive execution, not a user function with explicit parameters.
        with Tracer() as ctx:
            # Finalize graph by setting outputs
            result_traced = primitive.bind(*args, **kwargs)

            # Separate outputs into variables (Objects) and immediates (constants)
            out_vars, out_imms, morph_struct = var_morph(
                result_traced, lambda x: isinstance(x, Object)
            )

            if out_vars:
                graph = ctx.finalize(out_vars)
            else:
                # All outputs are immediates, no graph outputs
                graph = ctx.graph
                graph.outputs = []

        # Build inputs list for interpret
        # _captured_vars contains all inputs (no params in this context)
        inputs_list = [
            obj.runtime_obj if isinstance(obj, InterpObject) else obj
            for obj, _ in ctx._captured_vars.values()
        ]

        # Execute graph (may have 0 outputs if all were immediates)
        if graph.outputs:
            result_runtime_list = self.evaluate_graph(graph, inputs_list)
        else:
            result_runtime_list = []

        # Wrap runtime results as InterpObjects
        interp_results = [
            InterpObject(rt_val, tr_obj.type, self)
            for rt_val, tr_obj in zip(result_runtime_list, out_vars, strict=True)
        ]

        # Reconstruct the output tree: merge InterpObjects and immediates
        return var_demorph(interp_results, out_imms, morph_struct)

    def lift(self, obj: Any) -> InterpObject | Any:
        """Lift an object to the Interpreter's native representation.

        This is THE central method that manages the boundary between
        InterpObject and TraceObject:

        1. **InterpObject → TraceObject** (during nested tracing):
           - Register the InterpObject in self._objects for later resolution
           - The InterpObject must belong to this Interpreter
           - When the object flows into Tracer.lift() during bind_primitive,
             it will be captured as input with a clean SSA name like "%arg0"

        2. **TraceObject → InterpObject** (evaluate traced computation):
           - Extract the graph from the TraceObject's context (Tracer)
           - Execute the graph via evaluate_graph() to get runtime result
           - Wrap result as InterpObject and register it

        3. **Constants**: Pass through unchanged

        Args:
            obj: Object to lift (InterpObject, TraceObject, or constant)

        Returns:
            InterpObject (if Object input) or constant (pass-through)

        Example:
            >>> # InterpObject case
            >>> x = InterpObject(np.array([1, 2]), Tensor[f32, (2,)])
            >>> x_lifted = interp.lift(x)  # registers in _objects, returns x
            >>>
            >>> # TraceObject case
            >>> tracer = Tracer()
            >>> push_context(tracer)
            >>> z_trace = some_primitive.bind(x, y)  # TraceObject
            >>> pop_context()
            >>> interp = Interpreter()
            >>> z_interp = interp.lift(z_trace)  # evaluate graph → InterpObject
        """
        from mplang.edsl.tracer import TraceObject

        if isinstance(obj, InterpObject):
            # InterpObject must belong to this interpreter
            if obj._context is not None and obj._context is not self:
                raise ValueError(
                    f"InterpObject belongs to a different Interpreter. "
                    f"Object context: {obj._context}, Current interpreter: {self}"
                )
            return obj

        elif isinstance(obj, TraceObject):
            # Check execution cache
            # If this value was computed as part of a previous execution (e.g. sibling output)
            # we can return it immediately without re-execution.
            graph_value = obj._graph_value
            if graph_value in self._execution_cache:
                return self._execution_cache[graph_value]

            # First time seeing this Value.
            # We need to execute the graph to compute it.
            # MIMO Optimization:
            # Instead of just asking for this single value, we ask for ALL outputs
            # of the operation that produced this value. This ensures that if we
            # later ask for a sibling output, it will be in the cache.

            tracer = obj._context
            graph = tracer.graph
            defining_op = graph_value.defining_op

            if defining_op is None:
                # Value is likely a constant or input (no defining op in graph)
                # Just execute graph for this single value
                target_outputs = [graph_value]
            else:
                # Fetch all outputs of the defining op
                target_outputs = defining_op.outputs

            # Temporarily set graph outputs to the target outputs
            # We must save/restore original outputs to avoid side effects
            original_outputs = graph.outputs
            graph.outputs = target_outputs

            try:
                # Resolve inputs from Tracer's captured vars
                # _captured_vars preserves insertion order which matches graph.inputs order
                inputs_list = []
                for captured_obj, _ in tracer._captured_vars.values():
                    # Recursively lift captured objects to ensure they are ready
                    lifted = self.lift(captured_obj)
                    if isinstance(lifted, InterpObject):
                        inputs_list.append(lifted.runtime_obj)
                    else:
                        inputs_list.append(lifted)

                # Execute graph
                results_runtime = self.evaluate_graph(graph, inputs_list)

                # Cache all results
                for val, res in zip(target_outputs, results_runtime, strict=True):
                    # Wrap as InterpObject and cache
                    # Note: We use obj.type for the requested value, but for siblings
                    # we should ideally use their types. However, we don't have TraceObjects
                    # for siblings here, only GraphValues.
                    # InterpObject needs a type. GraphValue has a type.
                    self._execution_cache[val] = InterpObject(res, val.type, self)

            finally:
                # Restore original outputs
                graph.outputs = original_outputs

            # Now the result for our requested object should be in the cache
            if graph_value not in self._execution_cache:
                logger.error(
                    "Failed to compute value for %s after graph execution", obj
                )
                raise RuntimeError(
                    f"Failed to compute value for {obj} even after graph execution"
                )

            return self._execution_cache[graph_value]

        else:
            # Constants: pass through unchanged
            return obj

    def evaluate_graph(
        self, graph: Graph, inputs: list[Any], job_id: str | None = None
    ) -> list[Any]:
        """Execute a Graph IR with runtime data.

        Can be overridden by subclasses to implement remote execution or compilation.

        Args:
            graph: Finalized Graph IR to execute
            inputs: Runtime objects corresponding to graph.inputs (positional)
            job_id: Optional unique ID for this execution job (for profiling/tracing).

        Returns:
            List of runtime execution results corresponding to graph.outputs.
        """
        graph_exec_key = self._graph_exec_key(graph)

        # Prevent concurrent execution of the same graph hash.
        with self._exec_id_lock:
            if graph_exec_key in self._active_graph_exec_keys:
                raise RuntimeError(
                    "Concurrent execution of the same graph is not allowed. "
                    f"graph_exec_key={graph_exec_key}"
                )
            self._active_graph_exec_keys.add(graph_exec_key)

        try:
            with self._tls_exec_context(graph_exec_key=graph_exec_key):
                logger.debug(
                    "Evaluating graph: %d inputs, %d ops, %d outputs (job_id=%s, async=%s, graph_key=%s)",
                    len(inputs),
                    len(graph.operations),
                    len(graph.outputs),
                    job_id,
                    self.executor is not None,
                    graph_exec_key,
                )
                if self.executor:
                    return self._evaluate_graph_async(graph, inputs, job_id)
                else:
                    return self._evaluate_graph_sync(graph, inputs, job_id)
        finally:
            with self._exec_id_lock:
                self._active_graph_exec_keys.discard(graph_exec_key)

    def _reserve_op_exec_base(self, graph: Graph) -> int:
        """Reserve a contiguous exec_id range for a single evaluate_graph call.

        Counter is namespaced by the current graph_exec_key.
        """
        key = self.current_graph_exec_key()
        with self._exec_id_lock:
            base = self._graph_next_exec_base.get(key, 0)
            self._graph_next_exec_base[key] = base + len(graph.operations)
        return base

    def current_graph_exec_key(self) -> str:
        """Return current graph execution key during evaluate_graph execution."""

        key = getattr(self._tls, "current_graph_exec_key", None)
        if key is None:
            raise RuntimeError(
                "current_graph_exec_key() called outside of evaluate_graph execution"
            )
        return cast(str, key)

    def current_op_exec_id(self) -> int:
        """Return current op exec_id during graph execution.

        Worker-side implementations can use this to build deterministic,
        unique communication tags without coupling to any specific op.
        """

        exec_id = getattr(self._tls, "current_op_exec_id", None)
        if exec_id is None:
            raise RuntimeError(
                "current_op_exec_id() called outside of evaluate_graph execution"
            )
        return cast(int, exec_id)

    def _evaluate_graph_sync(
        self, graph: Graph, inputs: list[Any], job_id: str | None = None
    ) -> list[Any]:
        """Synchronous execution (Baseline)."""
        # Local environment: Value -> Runtime Object
        env = dict(zip(graph.inputs, inputs, strict=True))

        op_exec_base = self._reserve_op_exec_base(graph)

        for op_index, op in enumerate(graph.operations):
            exec_id = op_exec_base + op_index
            # Resolve inputs
            try:
                args = [env[val] for val in op.inputs]
            except KeyError as e:
                missing_keys = [str(k) for k in op.inputs if k not in env]
                # Limit available keys output to avoid flooding logs if env is huge
                available_keys = [str(k) for k in list(env.keys())[:20]]
                if len(env) > 20:
                    available_keys.append("...")

                raise RuntimeError(
                    f"Failed to resolve inputs for op '{op.opcode}'.\n"
                    f"Missing values: {missing_keys}\n"
                    f"Available values (partial): {available_keys}"
                ) from e

            # Dispatch
            # 1. Check instance-level handlers
            handler = self.handlers.get(op.opcode)
            # 2. Check global registry
            if not handler:
                handler = get_impl(op.opcode)

            with self._tls_exec_context(op_exec_id=exec_id):
                if handler:
                    # Pass interpreter to support recursive execution (HOFs)
                    # Pass op to access attributes and regions
                    # Pass args as runtime values
                    results = handler(self, op, *args)
                else:
                    raise NotImplementedError(
                        f"No implementation registered for opcode: {op.opcode}"
                    )

            # Update environment with outputs
            # Handler should return a single value or a tuple/list of values
            if len(op.outputs) == 0:
                pass  # Void operation
            elif len(op.outputs) == 1:
                env[op.outputs[0]] = results
            else:
                if len(results) != len(op.outputs):
                    raise RuntimeError(
                        f"Op {op.opcode} returned {len(results)} values, expected {len(op.outputs)}"
                    )
                for out_val, res in zip(op.outputs, results, strict=True):
                    env[out_val] = res

        # Return outputs
        if self.tracer and job_id:
            self.tracer.save_trace(job_id=job_id, rank=self.trace_pid)

        return [env[out] for out in graph.outputs]

    def _evaluate_graph_async(
        self, graph: Graph, inputs: list[Any], job_id: str | None = None
    ) -> list[Any]:
        """Asynchronous execution with non-blocking DAG scheduling."""
        graph_exec_key = self.current_graph_exec_key()
        op_exec_base = self._reserve_op_exec_base(graph)
        op_to_index = {op: i for i, op in enumerate(graph.operations)}
        # Tracer setup (if not provided, use a disabled stub)
        tracer: ExecutionTracer | _NullTracer
        if self.tracer:
            tracer = self.tracer
            tracer.total_ops += len(graph.operations)
        else:
            # No tracer provided - use minimal stub (no trace_dir needed)
            tracer = _NullTracer()

        active_tasks = 0

        # 1. Setup State
        # Value -> Runtime Object (initially inputs)
        env = dict(zip(graph.inputs, inputs, strict=True))

        # Op -> Pending Input Count
        pending_counts = {}
        # Value -> list[Op] (Consumers)
        value_to_consumers: dict[Any, list[Any]] = collections.defaultdict(list)
        # Value -> Remaining Consumers Count (for GC)
        remaining_consumers: dict[Any, int] = collections.defaultdict(int)

        # 2. Build Dependency Graph
        for op in graph.operations:
            count = 0
            for val in op.inputs:
                if val not in env:  # If not already resolved (input or constant)
                    value_to_consumers[val].append(op)
                    remaining_consumers[val] += 1
                    count += 1
            pending_counts[op] = count

        # Mark graph outputs as having an extra consumer (the user)
        # so they are not GC'd before return
        for out in graph.outputs:
            remaining_consumers[out] += 1

        # 3. Synchronization
        lock = threading.Lock()
        ready_queue: queue.Queue[Any] = queue.Queue()
        remaining_ops = len(graph.operations)

        # Error propagation
        error_occurred = False

        # 4. Execution Helper
        def on_op_done(op: Any, result: Any, error: Exception | None = None) -> None:
            nonlocal remaining_ops, error_occurred, active_tasks

            if error:
                with lock:
                    if not error_occurred:
                        error_occurred = True
                        ready_queue.put(error)
                return

            with lock:
                if op.opcode in self.async_ops and self.executor:
                    active_tasks -= 1
                    # profiler.sample(active_tasks, ready_queue.qsize())

                if error_occurred:
                    return

                # Store results
                if len(op.outputs) == 1:
                    env[op.outputs[0]] = result
                else:
                    for out_val, res in zip(op.outputs, result, strict=True):
                        env[out_val] = res

                # Trigger consumers
                for out_val in op.outputs:
                    if out_val in value_to_consumers:
                        for consumer_op in value_to_consumers[out_val]:
                            pending_counts[consumer_op] -= 1
                            if pending_counts[consumer_op] == 0:
                                tracer.log_schedule(
                                    consumer_op, namespace=self.trace_pid
                                )
                                ready_queue.put(consumer_op)

                # GC Inputs
                for val in op.inputs:
                    if val in remaining_consumers:
                        remaining_consumers[val] -= 1
                        if remaining_consumers[val] == 0:
                            env.pop(val, None)

                remaining_ops -= 1
                if remaining_ops == 0:
                    ready_queue.put(None)  # Sentinel

        def execute_op(op: Any) -> None:
            nonlocal active_tasks
            # Extract args from env (must be ready)
            args = [env[val] for val in op.inputs]

            exec_id = op_exec_base + op_to_index[op]

            handler = self.handlers.get(op.opcode)
            if not handler:
                handler = get_impl(op.opcode)

            if not handler:
                raise NotImplementedError(
                    f"No implementation registered for opcode: {op.opcode}"
                )

            if op.opcode in self.async_ops and self.executor:
                with lock:
                    active_tasks += 1
                    # profiler.sample(active_tasks, ready_queue.qsize())

                # Submit to executor
                def task() -> Any:
                    with self._tls_exec_context(
                        graph_exec_key=graph_exec_key, op_exec_id=exec_id
                    ):
                        start_ts = tracer.log_start(
                            op, pid=self.trace_pid, namespace=self.trace_pid
                        )
                        res = handler(self, op, *args)
                        tracer.log_end(op, start_ts, pid=self.trace_pid)
                        return res

                def callback(fut: Any) -> None:
                    try:
                        res = fut.result()
                        on_op_done(op, res)
                    except Exception as e:
                        on_op_done(op, None, error=e)

                fut = self.executor.submit(task)
                fut.add_done_callback(callback)
            else:
                # Sync execution (run immediately)
                try:
                    with self._tls_exec_context(
                        graph_exec_key=graph_exec_key, op_exec_id=exec_id
                    ):
                        start_ts = tracer.log_start(
                            op, pid=self.trace_pid, namespace=self.trace_pid
                        )
                        res = handler(self, op, *args)
                        tracer.log_end(op, start_ts, pid=self.trace_pid)
                        on_op_done(op, res)
                except Exception as e:
                    on_op_done(op, None, error=e)

        # 5. Initial Submission
        # Submit all ops with 0 pending inputs
        initial_ops = [op for op, count in pending_counts.items() if count == 0]
        if not initial_ops and remaining_ops > 0:
            # Cycle detected or empty graph?
            pass

        for op in initial_ops:
            tracer.log_schedule(op, namespace=self.trace_pid)
            ready_queue.put(op)

        # Handle empty graph case
        if remaining_ops == 0:
            ready_queue.put(None)

        # 6. Main Loop
        while True:
            item = ready_queue.get()
            if item is None:
                break
            if isinstance(item, Exception):
                raise item

            # It's an op
            execute_op(item)

        # 7. Return outputs
        if not self.tracer:
            tracer.stop()

        if self.tracer and job_id:
            self.tracer.save_trace(job_id=job_id, rank=self.trace_pid)

        return [env[out] for out in graph.outputs]
