# 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.

from __future__ import annotations

import concurrent.futures
import faulthandler
import logging
import sys
import threading
import traceback
from collections.abc import Sequence
from typing import Any, cast

import spu.libspu as libspu

from mplang.v1.core import (
    ClusterSpec,
    CollectiveMixin,
    CommunicatorBase,
    InterpContext,
    InterpVar,
    IrReader,
    IrWriter,
    Mask,
    MPObject,
    MPType,
    PFunction,  # for spu.seed_env kernel seeding
    TensorLike,
)
from mplang.v1.core.expr.ast import Expr
from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
from mplang.v1.kernels.context import RuntimeContext
from mplang.v1.runtime.link_comm import LinkCommunicator
from mplang.v1.utils.spu_utils import parse_field, parse_protocol


class ThreadCommunicator(CommunicatorBase, CollectiveMixin):
    """Thread-based communicator for in-memory communication between threads"""

    def __init__(self, rank: int, world_size: int):
        super().__init__(rank, world_size)
        self.peers: list[ThreadCommunicator] = []
        logging.debug(
            f"ThreadCommunicator initialized with rank={self.rank}, world_size={self.world_size}"
        )

    def set_peers(self, peers: list[ThreadCommunicator]) -> None:
        assert self.world_size == len(peers)
        self.peers = peers

    def send(self, to: int, key: str, data: Any) -> None:
        assert 0 <= to < self.world_size
        # print(f"send {key}: {self.rank} -> {to_rank}")
        self.peers[to].onSent(self.rank, key, data)


class SimVar(InterpVar):
    """A variable that references a value in an interpreter.

    SimVar represents a value that has been computed and exists
    in the interpreter's variable store.
    """

    def __init__(self, ctx: Simulator, mptype: MPType, values: list[Any]):
        # Initialize the parent InterpVar with a generated name
        super().__init__(ctx, mptype)
        self._values = values

    @property
    def values(self) -> list[Any]:
        """Converted values across all ranks for user inspection."""
        return [v.to_numpy() if hasattr(v, "to_numpy") else v for v in self._values]

    def __repr__(self) -> str:
        return f"SimVar({self.mptype})"


class Simulator(InterpContext):
    def __init__(
        self,
        cluster_spec: ClusterSpec,
        *,
        trace_ranks: list[int] | None = None,
    ) -> None:
        """Initialize a simulator with the given cluster specification.

        Args:
            cluster_spec: The cluster specification defining the simulation environment.
            trace_ranks: List of ranks to trace execution for debugging.
                Per-node op binding overrides should now be provided via
                each node's `runtime_info.op_bindings` in the supplied
                `cluster_spec`.
        """
        super().__init__(cluster_spec)
        self._trace_ranks = trace_ranks or []

        spu_devices = cluster_spec.get_devices_by_kind("SPU")
        if not spu_devices:
            raise ValueError("No SPU device found in the cluster specification")
        if len(spu_devices) > 1:
            raise ValueError("Multiple SPU devices found in the cluster specification")
        spu_device = spu_devices[0]

        # compute spu_mask from spu_device members
        spu_mask = Mask.from_ranks([member.rank for member in spu_device.members])

        # Convert protocol and field from config using utility functions
        spu_protocol = parse_protocol(spu_device.config["protocol"])
        spu_field = parse_field(spu_device.config["field"])

        world_size = self.world_size()

        # Setup communicators
        self._comms = [
            ThreadCommunicator(rank, world_size) for rank in range(world_size)
        ]
        for comm in self._comms:
            comm.set_peers(self._comms)

        # Prepare link contexts for SPU parties (store for evaluator-time initialization)
        # Use Channels mode to reuse ThreadCommunicator instead of separate mem_link
        self._spu_link_ctxs: list[LinkCommunicator | None] = [None] * world_size

        # Create LinkCommunicators in parallel to avoid deadlock
        # (create_with_channels does handshake via TestSend/TestRecv)
        exceptions: dict[int, Exception] = {}

        def create_link(g_rank: int) -> None:
            try:
                self._spu_link_ctxs[g_rank] = LinkCommunicator(
                    rank=g_rank,
                    comm=self._comms[g_rank],
                    spu_mask=spu_mask,
                )
            except Exception as e:
                exceptions[g_rank] = e

        threads = [
            threading.Thread(target=create_link, args=(g_rank,)) for g_rank in spu_mask
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        # Check for exceptions during link creation
        if exceptions:
            first_exc = next(iter(exceptions.values()))
            raise RuntimeError(
                f"Failed to create SPU link contexts for ranks {list(exceptions.keys())}"
            ) from first_exc

        self._spu_runtime_cfg = libspu.RuntimeConfig(
            protocol=spu_protocol, field=spu_field
        )
        self._spu_world = spu_mask.num_parties()
        self._spu_mask = spu_mask

        # Persistent per-rank RuntimeContext instances (reused across evaluates).
        # We no longer pre-create evaluators since each evaluate has different env bindings.
        # Build per-rank runtime contexts.
        self._runtimes: list[RuntimeContext] = []
        for rank in range(self.world_size()):
            node = self.cluster_spec.get_node_by_rank(rank)
            rt = RuntimeContext(
                rank=rank,
                world_size=self.world_size(),
                initial_bindings=node.runtime_info.op_bindings,
            )
            self._runtimes.append(rt)

    @classmethod
    def simple(
        cls,
        world_size: int,
        op_bindings: dict[str, str] | None = None,
        **kwargs: Any,
    ) -> Simulator:
        """Create a simple simulator with the given number of parties.

        This is a convenience method that creates a ClusterSpec.simple()
        configuration for quick testing and prototyping.

        Args:
            world_size: Number of simulated parties.
            **kwargs: Additional arguments passed to the Simulator constructor.

        Returns:
            A Simulator instance with a simple cluster configuration.
        """
        cluster_spec = ClusterSpec.simple(world_size)
        if op_bindings:
            # Apply the same op_bindings to every node's runtime_info for convenience
            for node in cluster_spec.nodes.values():
                node.runtime_info.op_bindings.update(op_bindings)
        return cls(cluster_spec, **kwargs)

    def _do_evaluate(self, expr: Expr, evaluator_engine: IEvaluator) -> Any:
        """
        Helper function to simulate real-world MPIR serialization/deserialization
        process instead of direct expr.accept execution.

        This exposes potential MPIR serialization bugs by forcing expressions
        to go through the full serialize->deserialize cycle.
        """
        writer = IrWriter()
        graph_proto = writer.dumps(expr)

        reader = IrReader()
        deserialized_expr = reader.loads(graph_proto)

        if deserialized_expr is None:
            raise ValueError("Failed to deserialize expression")

        return evaluator_engine.evaluate(deserialized_expr)

    # override
    def fetch(self, obj: MPObject) -> list[TensorLike]:
        if not isinstance(obj, SimVar):
            raise ValueError(f"Expected SimVar, got {type(obj)}")
        return [v.to_numpy() if hasattr(v, "to_numpy") else v for v in obj._values]

    # override
    def evaluate(self, expr: Expr, bindings: dict[str, MPObject]) -> Sequence[MPObject]:
        # sanity check for bindings.
        for name, var in bindings.items():
            if var.ctx is not self:
                raise ValueError(f"Variable {name} not in this context, got {var.ctx}.")

        pts_env = [
            {name: cast(SimVar, var)._values[rank] for name, var in bindings.items()}
            for rank in range(self.world_size())
        ]

        # Build per-rank evaluators with the per-party environment (runtime reused)
        pts_evaluators: list[IEvaluator] = []
        for rank in range(self.world_size()):
            runtime = self._runtimes[rank]
            ev = create_evaluator(
                rank,
                pts_env[rank],
                self._comms[rank],
                runtime,
                None,
            )
            # Seed SPU once per runtime (idempotent logical requirement)
            # Use setdefault to both retrieve and create metadata dict in one step.
            spu_meta = runtime.state.setdefault("_spu", {})
            if not spu_meta.get("inited", False):
                link_ctx = self._spu_link_ctxs[rank]
                seed_fn = PFunction(
                    fn_type="spu.seed_env",
                    ins_info=(),
                    outs_info=(),
                    config=self._spu_runtime_cfg,
                    world=self._spu_world,
                    link=link_ctx,
                )
                ev.runtime.run_kernel(seed_fn, [])  # type: ignore[arg-type]
                spu_meta["inited"] = True
            pts_evaluators.append(ev)

        # Collect evaluation results from all parties
        pts_results: list[Any] = []

        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = [
                executor.submit(self._do_evaluate, expr, evaluator)
                for evaluator in pts_evaluators
            ]

            # Collect results with proper exception handling
            for i, future in enumerate(futures):
                try:
                    result = future.result(100)  # 100 second timeout
                    pts_results.append(result)
                except concurrent.futures.TimeoutError:
                    faulthandler.dump_traceback(file=sys.stderr, all_threads=True)
                    raise
                except Exception as e:
                    print(
                        f"Exception in party {i}: {type(e).__name__}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc(file=sys.stderr)
                    executor.shutdown(wait=False, cancel_futures=True)
                    raise

        # Convert results to SimVar objects
        # pts_results is a list of party results, where each party result is a list of values
        # We need to transpose this to get (n_outputs, n_parties) structure
        assert len(pts_results) == self.world_size()

        # Ensure all parties returned the same number of outputs (matrix validation)
        if pts_results and not all(
            len(row) == len(pts_results[0]) for row in pts_results
        ):
            raise ValueError("Inconsistent number of outputs across parties")

        # Transpose: (n_parties, n_outputs) -> (n_outputs, n_parties)
        output_values = list(zip(*pts_results, strict=False))

        # Get the output types from the expression
        output_types = expr.mptypes

        # Create SimVar objects for each output
        sim_vars = []
        for values, mptype in zip(output_values, output_types, strict=False):
            sim_var = SimVar(self, mptype, list(values))
            sim_vars.append(sim_var)

        return sim_vars
