import enum
import functools
import json
import os
import random
import re
import socket
import threading
import time
from abc import ABC, abstractmethod
from collections.abc import Callable
from contextlib import contextmanager
from copy import deepcopy
from dataclasses import asdict, dataclass, field, fields, replace
from functools import cached_property
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from typing import Any, List, Literal, Optional, Union, get_args, get_type_hints

import torch
import torch.distributed as dist

# -------------------------------------- config base ------------------------------------------


@dataclass(frozen=True)
class _BaseConfig(ABC):
    def __post_init__(self) -> None:
        self._verify_types()

    def _verify_types(self) -> None:
        hints = get_type_hints(type(self))
        cls_name = type(self).__name__
        for f in fields(self):
            value = getattr(self, f.name)
            if value is None:
                continue
            expected = self._unwrap_type(hints[f.name])
            if not isinstance(value, expected):
                raise TypeError(
                    f"{cls_name}.{f.name}: expected {expected.__name__}, "
                    f"got {type(value).__name__}"
                )

    @classmethod
    @abstractmethod
    def _env_prefix(cls) -> str: ...

    @classmethod
    def _env_name(cls, field_name: str) -> str:
        return f"{cls._env_prefix()}{field_name.upper()}"

    @classmethod
    def from_env(cls) -> "_BaseConfig":
        return cls(
            **{
                f.name: cls._parse_env_field(cls._env_name(f.name), f.default)
                for f in fields(cls)
            }
        )

    def with_defaults(self, **kwargs) -> "_BaseConfig":
        cls = type(self)
        actual = {
            key: value
            for key, value in kwargs.items()
            if os.getenv(cls._env_name(key)) is None
        }
        return replace(self, **actual) if actual else self

    @staticmethod
    def _unwrap_type(hint) -> type:
        args = get_args(hint)
        if args:
            return next(a for a in args if a is not type(None))
        return hint

    @classmethod
    def _parse_env_field(cls, env_name: str, default):
        return cls._parse_env_value(os.getenv(env_name), default)

    @staticmethod
    def _parse_env_value(raw, default):
        if raw is None or not raw.strip():
            return default
        if isinstance(default, bool):
            return raw.lower() in ("true", "1")
        if isinstance(default, int):
            return int(raw)
        return raw

    @classmethod
    def from_kv_pairs(cls, pairs: Optional[List[str]]) -> "_BaseConfig":
        return cls(**cls._kv_pairs_to_dict(pairs))

    @classmethod
    def _kv_pairs_to_dict(cls, pairs: Optional[List[str]]) -> dict:
        if not pairs:
            return {}

        missing = object()
        defaults = {f.name: f.default for f in fields(cls)}
        result: dict = {}

        for pair in pairs:
            key, sep, value = pair.partition("=")
            if not sep:
                raise ValueError(f"Invalid config pair (missing '='): {pair!r}")
            default = defaults.get(key, missing)
            if default is missing:
                raise ValueError(
                    f"Unknown config key {key!r}. Valid keys: {sorted(defaults)}"
                )
            try:
                result[key] = cls._parse_env_value(value, default)
            except (ValueError, TypeError) as exc:
                field_type = type(default).__name__
                raise TypeError(f"{key}: expected {field_type}, got {value!r}") from exc

        return result


_DEFAULT_EXP_NAME_PREFIX = "dump_"


@dataclass(frozen=True)
class DumperConfig(_BaseConfig):
    enable: bool = False
    filter: Optional[str] = None
    dir: str = "/tmp/dumper"
    enable_output_file: bool = True
    enable_output_console: bool = True
    enable_value: bool = True
    enable_grad: bool = False
    enable_model_value: bool = False
    enable_model_grad: bool = False
    exp_name: Optional[str] = None
    cleanup_previous: bool = False
    collective_timeout: int = 60
    server_port: str = "-1"
    non_intrusive_mode: str = "core"
    source_patcher_config: Optional[str] = None

    @classmethod
    def _env_prefix(cls) -> str:
        # NOTE: should not be `SGLANG_DUMPER_`, otherwise it is weird when dumping Megatron in Miles
        return "DUMPER_"

    @property
    def server_port_parsed(self) -> Optional[Union[int, Literal["reuse"]]]:
        raw = self.server_port
        if raw == "reuse":
            return "reuse"
        port = int(raw)
        if port <= 0:
            return None
        return port


# -------------------------------------- dumper core ------------------------------------------


@dataclass
class _DumperState:
    dump_index: int = 0
    step: int = 0
    global_ctx: dict = field(default_factory=dict)
    captured_output_data: Optional[dict] = None
    cleanup_previous_handled: bool = False


class _Dumper:
    """Utility to dump tensors, which can be useful when comparison checking models.

    Example usage:
    dumper.dump("layer_start__hidden_states", hidden_states, layer_id=self.layer_id)
    dumper.step()

    Import from non-SGLang system:
    ```
    import sys
    sys.path.append("/YOUR_PATH/sglang/python/sglang/srt/debug_utils")
    from dumper import dumper
    ```

    Then run the program:
    `DUMPER_ENABLE=1 python ...`

    Auto-cleanup old dumps before first write:
    `DUMPER_CLEANUP_PREVIOUS=1 python ...`

    Alternatively, disable at startup and configure via HTTP:
    1. `python ...`
    2. sglang mode:  `curl -X POST http://localhost:30000/dumper/configure -d '{"enable": true}'`
       standalone:   `curl -X POST http://localhost:40000/dumper/configure -d '{"enable": true}'`
    3. `curl -X POST http://localhost:30000/dumper/configure -d '{"enable": true, "filter": "layer_id=[0-3]"}'`
    4. `curl -X POST http://localhost:30000/dumper/reset`

    Related: `sglang.srt.debug_utils.dump_comparator` for dump comparison
    """

    def __init__(self, *, config: DumperConfig):
        self._config = config
        self._state = _DumperState()
        self._non_intrusives: list["_NonIntrusiveDumper"] = []

    # ------------------------------- public :: core ---------------------------------

    @property
    def may_enable(self) -> bool:
        return self._config.enable or self._config.server_port_parsed is not None

    def step(self):
        """This should be called on all ranks at the end of each iteration."""

        self._http_manager  # noqa: B018

        if not self._config.enable:
            return

        # Users may want to `dump` only on some ranks, thus determine name here
        self._ensure_exp_name()

        self._state.step += 1
        print(f"[Dumper] [{time.time()}] step={self._state.step}")

    def dump(
        self,
        name: str,
        value,
        save: bool = True,
        dims: Optional[str] = None,
        dims_grad: Optional[str] = None,
        **kwargs,
    ) -> None:
        value_meta: dict = {}
        grad_meta: dict = {}
        if dims is not None:
            value_meta["dims"] = dims
            grad_meta["dims"] = dims
        if dims_grad is not None:
            value_meta["dims_grad"] = dims_grad
            grad_meta["dims"] = dims_grad

        self._dump_inner(
            name=name,
            value=value,
            extra_kwargs=kwargs,
            save=save,
            enable_value=self._config.enable_value,
            enable_curr_grad=False,
            enable_future_grad=self._config.enable_grad,
            value_tag="Dumper.Value",
            grad_tag="Dumper.Grad",
            value_meta_only_fields=value_meta,
            grad_meta_only_fields=grad_meta,
        )

    def dump_model(
        self,
        model: "torch.nn.Module",
        name_prefix: str = "param",
        save: bool = True,
        **kwargs,
    ) -> None:
        for param_name, param in model.named_parameters():
            self._dump_inner(
                name=f"{name_prefix}__{param_name}",
                value=param,
                extra_kwargs=kwargs,
                save=save,
                enable_value=self._config.enable_model_value,
                enable_curr_grad=self._config.enable_model_grad,
                enable_future_grad=False,
                value_tag="Dumper.ParamValue",
                grad_tag="Dumper.ParamGrad",
            )

    def dump_dict(self, name_prefix, data, save: bool = True, **kwargs):
        data = _obj_to_dict(data)
        for name, value in data.items():
            self.dump(f"{name_prefix}_{name}", value, save=save, **kwargs)

    def set_ctx(self, **kwargs):
        """
        Example:

        dumper.configure_default(filter='layer_id=[0-3]')
        dumper.set_ctx(layer_id=self.layer_id)
        ...
        dumper.set_ctx(layer_id=None)
        """
        self._state.global_ctx = {
            k: v for k, v in (self._state.global_ctx | kwargs).items() if v is not None
        }

    def ctx(
        self,
        _extractor: Optional[Callable[..., dict]] = None,
        **static_ctx: Any,
    ) -> Callable:
        """Decorator that sets context before calling the wrapped function and clears it after.

        Two forms:
            @dumper.ctx(lambda self: dict(layer_id=self.layer_id))
            def forward(self, x): ...

            @dumper.ctx(phase="decode")
            def decode_step(self, x): ...
        """
        if _extractor is not None and static_ctx:
            raise ValueError("cannot mix lambda extractor with static kwargs")
        if _extractor is None and not static_ctx:
            raise ValueError("must provide either a lambda or static kwargs")

        def decorator(fn: Callable) -> Callable:
            @functools.wraps(fn)
            def wrapper(*args: Any, **kwargs: Any) -> Any:
                ctx_dict: dict = _extractor(args[0]) if _extractor else static_ctx
                self.set_ctx(**ctx_dict)
                try:
                    return fn(*args, **kwargs)
                finally:
                    self.set_ctx(**{k: None for k in ctx_dict})

            return wrapper

        return decorator

    def apply_source_patches(self) -> None:
        """Apply source patches from DUMPER_SOURCE_PATCHER_CONFIG if set.

        Automatically injects ``from sglang.srt.debug_utils.dumper import dumper``
        into every replacement block so users don't need to write it in YAML.
        """
        config_path = self._config.source_patcher_config
        if not config_path:
            return

        from sglang.srt.debug_utils.source_patcher import apply_patches_from_config

        yaml_content: str = Path(config_path).read_text()
        print(f"[source_patcher] loading config from {config_path}")
        apply_patches_from_config(
            yaml_content,
            extra_imports=["from sglang.srt.debug_utils.dumper import dumper"],
        )

    def register_non_intrusive_dumper(
        self,
        model: "torch.nn.Module",
    ) -> Optional["_NonIntrusiveDumper"]:
        self._http_manager  # noqa: B018
        mode = self._config.non_intrusive_mode
        if mode == "off":
            return None
        non_intrusive = _NonIntrusiveDumper(dumper=self, model=model, mode=mode)
        self._non_intrusives.append(non_intrusive)
        return non_intrusive

    # ------------------------------- public :: secondary ---------------------------------

    def configure(self, **kwargs) -> None:
        self._config = replace(self._config, **kwargs)

    def configure_default(self, **kwargs) -> None:
        self._config = self._config.with_defaults(**kwargs)

    def reset(self) -> None:
        for non_intrusive in self._non_intrusives:
            non_intrusive.remove()
        self._non_intrusives.clear()
        self._state = _DumperState()

    @contextmanager
    def capture_output(self):
        assert self._state.captured_output_data is None
        self._state.captured_output_data = {}
        try:
            yield self._state.captured_output_data
        finally:
            self._state.captured_output_data = None

    def get_state(self) -> dict:
        return {
            "config": asdict(self._config),
            "dump_index": self._state.dump_index,
            "step": self._state.step,
        }

    @cached_property
    def _http_manager(self) -> Optional["_DumperHttpManager"]:
        if self._config.server_port_parsed is None:
            return None
        return _DumperHttpManager(self)

    # ------------------------- private :: related to dump -----------------------------

    def _dump_inner(
        self,
        *,
        name: str,
        value,
        extra_kwargs: dict,
        save: bool,
        enable_value: bool,
        enable_curr_grad: bool,
        enable_future_grad: bool,
        value_tag: str,
        grad_tag: str,
        value_meta_only_fields: Optional[dict] = None,
        grad_meta_only_fields: Optional[dict] = None,
    ) -> None:
        self._http_manager  # noqa: B018

        if not self._config.enable:
            return

        recompute_status = _detect_recompute_status()
        tags = dict(
            name=name,
            recompute_status=recompute_status.value,
            **extra_kwargs,
            **self._state.global_ctx,
        )

        if (f := self._config.filter) is not None and not _evaluate_filter(f, tags):
            return

        if not (enable_value or enable_curr_grad or enable_future_grad):
            return

        recompute_meta = recompute_status.to_pseudo_parallel_meta()
        value = _materialize_value(value)

        if enable_value:
            self._dump_single(
                tag=value_tag,
                tags=tags,
                value=value,
                save=save,
                meta_only_fields={**(value_meta_only_fields or {}), **recompute_meta},
            )

        if (
            enable_curr_grad
            and isinstance(value, torch.Tensor)
            and (g := value.grad) is not None
        ):
            self._dump_single(
                tag=grad_tag,
                tags={**tags, "name": f"grad__{name}"},
                value=g,
                save=save,
                meta_only_fields={**(grad_meta_only_fields or {}), **recompute_meta},
            )

        if enable_future_grad:
            self._register_dump_grad_hook(
                name=name,
                tensor=value,
                extra_kwargs=extra_kwargs,
                save=save,
                meta_only_fields=grad_meta_only_fields or {},
            )

    def _register_dump_grad_hook(
        self,
        *,
        name: str,
        tensor,
        extra_kwargs: dict,
        save: bool,
        meta_only_fields: Optional[dict] = None,
    ) -> None:
        if not isinstance(tensor, torch.Tensor):
            return
        if not tensor.requires_grad:
            return

        captured_step = self._state.step
        captured_tags = dict(
            name=f"grad__{name}",
            **deepcopy(extra_kwargs),
        )
        captured_meta_only = meta_only_fields or {}

        def grad_hook(grad: torch.Tensor) -> None:
            self._dump_single(
                tag="Dumper.Grad",
                tags=captured_tags,
                value=grad,
                save=save,
                step=captured_step,
                meta_only_fields=captured_meta_only,
            )

        tensor.register_hook(grad_hook)

    def _dump_single(
        self,
        *,
        tag: str,
        tags: dict,
        value,
        save: bool,
        step: Optional[int] = None,
        meta_only_fields: Optional[dict] = None,
    ) -> None:
        self._ensure_exp_name()
        self._state.dump_index += 1

        rank = _get_rank()
        full_kwargs = dict(
            step=(step if step is not None else self._state.step),
            rank=rank,
            dump_index=self._state.dump_index,
            **tags,
        )
        full_filename = _format_tags(full_kwargs) + ".pt"
        path = Path(self._config.dir) / self._config.exp_name / full_filename

        if self._config.enable_output_console:
            print(
                f"[{tag}] [{rank}, {time.time()}] {path} "
                f"type={type(value)} "
                f"shape={value.shape if isinstance(value, torch.Tensor) else None} "
                f"dtype={value.dtype if isinstance(value, torch.Tensor) else None} "
                f"device={value.device if isinstance(value, torch.Tensor) else None} "
                f"id={id(value)} "
                f"sample_value={get_truncated_value(value)}"
            )

        capturing = self._state.captured_output_data is not None
        if save and (self._config.enable_output_file or capturing):
            output_data = {
                "value": value,
                "meta": dict(
                    **full_kwargs,
                    **self._static_meta,
                    **(meta_only_fields or {}),
                ),
            }

            if capturing:
                output_data["value"] = _deepcopy_or_clone(output_data["value"])
                self._state.captured_output_data[tags["name"]] = output_data
            else:
                if (
                    not self._state.cleanup_previous_handled
                    and self._config.cleanup_previous
                ):
                    self._state.cleanup_previous_handled = True
                    _cleanup_old_dumps(
                        Path(self._config.dir), exp_name=self._config.exp_name
                    )

                path.parent.mkdir(parents=True, exist_ok=True)
                _torch_save(output_data, str(path))

    # ------------------------------- private :: misc ---------------------------------

    @cached_property
    def _static_meta(self) -> dict:
        return _compute_static_meta()

    def _ensure_exp_name(self):
        if self._config.exp_name is None:
            name = _get_default_exp_name(
                timeout_seconds=self._config.collective_timeout
            )
            self.configure(exp_name=name)
            print(f"[Dumper] Choose exp_name={name}")


# -------------------------------------- hook dumper ------------------------------------------


class _NonIntrusiveDumper:
    _NAME_PREFIX = "non_intrusive__"
    _LAYER_NAME_RE = re.compile(r"(?:.+\.)?layers\.(\d+)$")

    def __init__(
        self,
        dumper: _Dumper,
        model: "torch.nn.Module",
        mode: str,
    ):
        self._dumper = dumper
        self._mode = mode
        self._handles: list = []
        self._core_fields: frozenset[str] = frozenset().union(
            *(p.core_fields() for p in _plugins)
        )

        for module_name, module in model.named_modules():
            if ctx := self._detect_module_ctx(module_name, module):
                self._register_ctx_hooks(module, ctx=ctx)

            is_root = module_name == ""
            pre_hook = self._make_forward_pre_hook(
                module_name=module_name, is_root=is_root
            )
            hook = self._make_forward_hook(module_name=module_name, is_root=is_root)
            self._handles += _register_forward_hook_or_replace_fn(
                module,
                pre_hook=pre_hook,
                hook=hook,
                mode="replace_fn" if is_root else "hook",
            )

    def remove(self) -> None:
        for handle in self._handles:
            handle.remove()
        self._handles.clear()

    @classmethod
    def _detect_module_ctx(
        cls, module_name: str, module: "torch.nn.Module"
    ) -> Optional[dict]:
        match = cls._LAYER_NAME_RE.fullmatch(module_name)
        if match:
            for plugin in _plugins:
                layer_id = plugin.detect_layer_id(module)
                if layer_id is not None:
                    return {"layer_id": layer_id}
            return {"layer_id": int(match.group(1))}
        return None

    def _register_ctx_hooks(self, module: "torch.nn.Module", *, ctx: dict) -> None:
        clear_ctx = {k: None for k in ctx}
        self._handles.append(
            module.register_forward_pre_hook(
                lambda _mod, _input, _ctx=ctx: self._dumper.set_ctx(**_ctx)
            )
        )
        self._handles.append(
            module.register_forward_hook(
                lambda _mod, _input, _output, _clear=clear_ctx: self._dumper.set_ctx(
                    **_clear
                )
            )
        )

    def _make_forward_pre_hook(self, *, module_name: str, is_root: bool):
        def _hook(_module, args, kwargs):
            for i, item in enumerate(args):
                self._dump_value(
                    module_name, item, sub_name=f"inputs.{i}", is_root=is_root
                )
            for name, value in kwargs.items():
                self._dump_value(
                    module_name,
                    value,
                    sub_name=f"inputs.{name}",
                    is_root=is_root,
                )

        return _hook

    def _make_forward_hook(self, *, module_name: str, is_root: bool):
        def _hook(_module, input, output):
            if output is not None:
                self._dump_value(module_name, output, sub_name="output", is_root=False)

        return _hook

    def _dump_value(
        self, module_name: str, value: Any, sub_name: str, *, is_root: bool
    ) -> None:
        for key, item in self._convert_value(
            value, skip_forward_batch=(not is_root)
        ).items():
            effective_key = key or sub_name.rsplit(".", 1)[-1]
            if effective_key in self._core_fields:
                self._dumper.dump(effective_key, item)
            elif self._mode == "all":
                parts = [p for p in (module_name, sub_name, key) if p]
                self._dumper.dump(self._NAME_PREFIX + ".".join(parts), item)

    @staticmethod
    def _convert_value(value, *, skip_forward_batch: bool = False) -> dict[str, Any]:
        if isinstance(value, torch.Tensor):
            return {"": value}

        if isinstance(value, (tuple, list)):
            tensors = [t for t in value if isinstance(t, torch.Tensor)]
            if len(tensors) == 1:
                return {"": tensors[0]}
            return {str(i): t for i, t in enumerate(tensors)}

        for plugin in _plugins:
            result = plugin.convert_value(value, skip_forward_batch=skip_forward_batch)
            if result is not None:
                return result

        return {}


def _register_forward_hook_or_replace_fn(
    module: "torch.nn.Module",
    *,
    pre_hook,
    hook,
    mode: str,
) -> list:
    """Attach pre/post forward hooks to *module*.

    mode="hook"       — standard ``register_forward_pre_hook`` / ``register_forward_hook``
                        (fires only via ``__call__``).
    mode="replace_fn" — monkey-patch ``module.forward`` so hooks fire even when
                        callers invoke ``.forward()`` directly (as sglang does for the
                        root model).

    Returns a list of handle objects with a ``.remove()`` method that undoes
    the registration.
    """
    if mode == "hook":
        return [
            module.register_forward_pre_hook(pre_hook, with_kwargs=True),
            module.register_forward_hook(hook),
        ]
    elif mode == "replace_fn":
        original_forward = module.forward

        @functools.wraps(original_forward)
        def _wrapped(*args, **kwargs):
            pre_hook(module, args, kwargs)
            output = original_forward(*args, **kwargs)
            hook(module, args, output)
            return output

        module.forward = _wrapped

        class _Handle:
            def remove(self) -> None:
                assert module.forward is _wrapped
                module.forward = original_forward

        return [_Handle()]
    else:
        raise ValueError(f"Unknown mode {mode!r}")


# -------------------------------------- util fn ------------------------------------------


def _torch_save(value, path: str):
    value = _clone_if_view(value)
    try:
        try:
            return torch.save(value, path)
        except RuntimeError as e:
            if "not pickleable" in str(e):
                stripped = _strip_parameter(value)
                if stripped is not value:
                    print(f"[Dumper] Observe error={e} and try pickling .data")
                    return _torch_save(stripped, path)
            raise
    except Exception as e:
        print(f"[Dumper] Observe error={e} when saving data, skip the tensor")


def _map_tensor(value, fn: Callable[[torch.Tensor], torch.Tensor]):
    if isinstance(value, dict):
        return {k: _map_tensor(v, fn) for k, v in value.items()}
    if isinstance(value, torch.Tensor):
        return fn(value)
    return value


def _clone_if_view(value):
    def _fn(t: torch.Tensor) -> torch.Tensor:
        if t.untyped_storage().nbytes() > t.nelement() * t.element_size():
            return t.clone()
        return t

    return _map_tensor(value, _fn)


def _strip_parameter(value):
    def _fn(t: torch.Tensor) -> torch.Tensor:
        if isinstance(t, torch.nn.Parameter):
            return t.data
        return t

    return _map_tensor(value, _fn)


def _collective_with_timeout(fn, operation_name: str, timeout_seconds: int = 60):
    completed = threading.Event()

    def watchdog():
        if not completed.wait(timeout=timeout_seconds):
            print(
                f"\n[Dumper] WARNING: '{operation_name}' has not completed after "
                f"{timeout_seconds}s. This usually means not all ranks are "
                f"participating in this collective operation.\n",
                flush=True,
            )

    thread = threading.Thread(target=watchdog, daemon=True)
    thread.start()
    try:
        return fn()
    finally:
        completed.set()


def _get_default_exp_name(timeout_seconds: int = 60):
    rank = _get_rank()
    now = time.time()
    ms = int((now % 1) * 1000)
    rand_suffix = random.randint(0, 999)
    object_list = [
        (
            (
                f"{_DEFAULT_EXP_NAME_PREFIX}"
                f"{time.strftime('%Y%m%d_%H%M%S', time.gmtime(now))}"
                f"_{ms:03d}{rand_suffix:03d}"
            )
            if rank == 0
            else None
        )
    ]

    if dist.is_initialized():
        _collective_with_timeout(
            lambda: dist.broadcast_object_list(object_list, device="cuda"),
            operation_name="broadcast_object_list in _get_default_exp_name",
            timeout_seconds=timeout_seconds,
        )

    return object_list[0]


def _cleanup_old_dumps(base_dir: Path, exp_name: Optional[str] = None) -> None:
    import shutil

    if _get_rank() == 0:
        targets = {entry for entry in base_dir.glob(f"{_DEFAULT_EXP_NAME_PREFIX}*")}
        if exp_name:
            targets.add(base_dir / exp_name)
        targets = {d for d in targets if d.is_dir()}

        for entry in targets:
            shutil.rmtree(entry)
            print(f"[Dumper] Cleaned up {entry}")

    if dist.is_initialized():
        _collective_with_timeout(
            dist.barrier,
            operation_name="barrier in _cleanup_old_dumps",
        )


def _get_rank():
    if dist.is_initialized():
        return dist.get_rank()
    else:
        return 0


def _get_world_size():
    if dist.is_initialized():
        return dist.get_world_size()
    else:
        return 1


def _obj_to_dict(obj):
    if isinstance(obj, dict):
        return obj
    ret = {}
    for k in dir(obj):
        if k.startswith("__") and k.endswith("__"):
            continue
        try:
            v = getattr(obj, k)
            if not callable(v):
                ret[k] = v
        except Exception:
            # Skip attributes that raise an exception on access
            continue
    return ret


def _materialize_value(value):
    if callable(value):
        value = value()
    return value


def _format_tags(kwargs: dict) -> str:
    return "___".join(f"{k}={v}" for k, v in kwargs.items())


class _DefaultNoneDict(dict):
    """dict subclass that returns None for missing keys, for filter expression eval."""

    def __missing__(self, key: str):
        return None


_FILTER_BUILTINS: dict[str, Any] = {"search": re.search, "match": re.match}


def _evaluate_filter(filter_expr: str, tags: dict[str, Any]) -> bool:
    """Evaluate a Python filter expression against the tags dict.

    Unknown tag keys resolve to None, so `layer_id is None` works when layer_id is absent.
    `re.search` and `re.match` are available as `search()` and `match()`.
    """
    namespace = _DefaultNoneDict(tags)
    namespace.update(_FILTER_BUILTINS)
    return bool(eval(filter_expr, {"__builtins__": {}}, namespace))


def _deepcopy_or_clone(x):
    if isinstance(x, torch.Tensor):
        return x.clone()
    return deepcopy(x)


# -------------------------------------- static meta ------------------------------------------


def _compute_static_meta():
    result = {
        "world_rank": _get_rank(),
        "world_size": _get_world_size(),
    }

    for plugin in _plugins:
        if info := plugin.collect_parallel_info():
            result[f"{plugin.name}_parallel_info"] = info

    for plugin in _plugins:
        tokenizer_path: Optional[str] = plugin.get_tokenizer_path()
        if tokenizer_path is not None:
            result["tokenizer_path"] = tokenizer_path
            break

    return result


# -------------------------------------- http manager ------------------------------------------


class _DumperHttpManager:
    def __init__(self, dumper: "_Dumper"):
        self._dumper = dumper
        http_port = self._dumper._config.server_port_parsed

        rpc_broadcast = _create_zmq_rpc_broadcast(
            self,
            timeout_seconds=self._dumper._config.collective_timeout,
        )

        if _get_rank() == 0:
            assert rpc_broadcast is not None
            self._rpc_broadcast = rpc_broadcast

            if http_port == "reuse":
                print(
                    "[Dumper] Standalone HTTP server disabled, reusing existing ports"
                )
            else:
                _start_http_server(prefix="/dumper/", target=self, http_port=http_port)
                print(f"[Dumper] HTTP server started on port {http_port}")

    # ------------------------------- public ---------------------------------

    def handle_request(self, *, method: str, body: dict[str, Any]) -> list[dict]:
        return self._rpc_broadcast._handle_request_inner(method=method, body=body)

    # ------------------------------- private ---------------------------------

    def _handle_request_inner(self, *, method: str, body: dict[str, Any]) -> dict:
        if method == "get_state":
            return self._dumper.get_state()
        elif method == "configure":
            self._dumper.configure(**body)
            return {}
        elif method == "reset":
            self._dumper.reset()
            return {}
        else:
            raise ValueError(f"Unknown dumper control method: {method!r}")


# -------------------------------------- http control server ------------------------------------------


def _start_http_server(*, prefix: str, target: object, http_port: int):
    handler_class = _make_http_handler(prefix=prefix, target=target)
    server = HTTPServer(("0.0.0.0", http_port), handler_class)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()


def _make_http_handler(*, prefix: str, target):
    class _HTTPHandler(BaseHTTPRequestHandler):
        def do_POST(self):
            if not self.path.startswith(prefix):
                self.send_error(404)
                return
            method = self.path[len(prefix) :]
            try:
                req_body = self._get_request_body()
                print(f"[Dumper#{_get_rank()}] HTTP {self.path} {req_body=}")
                result = target.handle_request(method=method, body=req_body)
                resp_body = json.dumps(result).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(resp_body)))
                self.end_headers()
                self.wfile.write(resp_body)
            except Exception as e:
                self.send_error(400, str(e))

        def _get_request_body(self) -> dict:
            content_length = int(self.headers.get("Content-Length", 0))
            if content_length == 0:
                return {}
            return json.loads(self.rfile.read(content_length))

    return _HTTPHandler


# -------------------------------------- zmq rpc ------------------------------------------


def _create_zmq_rpc_broadcast(
    handler, timeout_seconds: int = 60
) -> Optional["_ZmqRpcBroadcast"]:
    """A general-purpose minimal RPC to support broadcasting executions to multi processes"""
    import zmq

    rank = _get_rank()
    world_size = dist.get_world_size() if dist.is_initialized() else 1

    ctx = zmq.Context()
    sock = ctx.socket(zmq.REP)
    sock.bind("tcp://*:0")
    bound_port = int(sock.getsockopt_string(zmq.LAST_ENDPOINT).rsplit(":", 1)[1])
    local_addr = f"tcp://{_get_local_ip_by_remote()}:{bound_port}"

    def serve_loop():
        while True:
            try:
                req = sock.recv_pyobj()
                result = getattr(handler, req["method"])(*req["args"], **req["kwargs"])
                resp = {"result": result, "error": None}
            except Exception as e:
                print(f"[Dumper.ZmqRpc] error inside handler: {e}")
                resp = {"result": None, "error": str(e)}
            sock.send_pyobj(resp)

    thread = threading.Thread(target=serve_loop, daemon=True)
    thread.start()
    print(f"[Dumper.ZmqRpc] rank={rank} server started at {local_addr}")

    if dist.is_initialized():
        all_addresses = [None] * world_size
        _collective_with_timeout(
            lambda: dist.all_gather_object(all_addresses, local_addr),
            operation_name="all_gather_object in _create_zmq_rpc_broadcast",
            timeout_seconds=timeout_seconds,
        )
    else:
        all_addresses = [local_addr]
    print(f"[Dumper.ZmqRpc] rank={rank} all_addresses={all_addresses}")

    if rank == 0:
        handles = []
        for i, addr in enumerate(all_addresses):
            req_socket = ctx.socket(zmq.REQ)
            req_socket.connect(addr)
            handles.append(_ZmqRpcHandle(req_socket, debug_name=f"rank-{i}"))
        return _ZmqRpcBroadcast(handles)
    else:
        return None


class _ZmqRpcHandle:
    """Proxy object to call remote handler methods via ZMQ."""

    def __init__(self, socket, debug_name: str):
        self._socket = socket
        self._debug_name = debug_name

    def __getattr__(self, method_name: str):
        def call(*args, **kwargs):
            self._socket.send_pyobj(
                {
                    "method": method_name,
                    "args": args,
                    "kwargs": kwargs,
                }
            )
            response = self._socket.recv_pyobj()
            if response["error"]:
                raise RuntimeError(
                    f"RPC error on {self._debug_name}: {response['error']}"
                )
            return response["result"]

        return call


class _RpcBroadcastBase:
    """Base for broadcasting method calls to dumper instance(s)."""

    def __getattr__(self, method_name: str):
        raise NotImplementedError

    def __init__(self, handles: List[_ZmqRpcHandle]):
        self._handles = handles


class _ZmqRpcBroadcast(_RpcBroadcastBase):
    """Broadcasts method calls to all ZMQ RPC handles.

    Returns a list of results, one per rank (ordered by rank).
    """

    def __init__(self, handles: List[_ZmqRpcHandle]):
        self._handles = handles

    def __getattr__(self, method_name: str):
        def call(*args, **kwargs):
            return [
                getattr(handle, method_name)(*args, **kwargs)
                for handle in self._handles
            ]

        return call


# --------------------------------- copied code (avoid dependency) --------------------------------------


def _get_local_ip_by_remote() -> Optional[str]:
    # try ipv4
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))  # Doesn't need to be reachable
        return s.getsockname()[0]
    except Exception:
        pass

    try:
        hostname = socket.gethostname()
        ip = socket.gethostbyname(hostname)
        if ip and ip != "127.0.0.1" and ip != "0.0.0.0":
            return ip
    except Exception:
        pass

    # try ipv6
    try:
        s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
        # Google's public DNS server, see
        # https://developers.google.com/speed/public-dns/docs/using#addresses
        s.connect(("2001:4860:4860::8888", 80))  # Doesn't need to be reachable
        return s.getsockname()[0]
    except Exception:
        print("Can not get local ip by remote")
    return None


# -------------------------------------- framework plugins ------------------------------------------


class _RecomputeStatus(enum.Enum):
    DISABLED = "disabled"
    ORIGINAL = "original"  # inside checkpoint, original forward
    RECOMPUTE = "recompute"  # inside checkpoint, recompute forward

    def to_pseudo_parallel_meta(self) -> dict[str, Any]:
        if self == _RecomputeStatus.DISABLED:
            return {}
        return {
            "recompute_pseudo_rank": 1 if self == _RecomputeStatus.RECOMPUTE else 0,
            "recompute_pseudo_size": 2,
        }


class _FrameworkPlugin(ABC):
    @property
    @abstractmethod
    def name(self) -> str: ...

    @abstractmethod
    def collect_parallel_info(self) -> dict: ...

    @abstractmethod
    def convert_value(
        self, value: Any, *, skip_forward_batch: bool
    ) -> Optional[dict[str, Any]]:
        """Return converted dict, or None if this plugin doesn't handle the value."""
        ...

    @abstractmethod
    def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]:
        """Return 0-indexed layer_id, or None if not detectable."""
        ...

    def core_fields(self) -> frozenset[str]:
        return frozenset()

    def get_tokenizer_path(self) -> Optional[str]:
        return None

    def detect_recompute_status(self) -> _RecomputeStatus:
        return _RecomputeStatus.DISABLED


class _SGLangPlugin(_FrameworkPlugin):
    _available = True
    try:
        from sglang.srt import distributed as _dist
        from sglang.srt.layers import dp_attention as _dp_attn
        from sglang.srt.layers.logits_processor import LogitsProcessorOutput
        from sglang.srt.model_executor.forward_batch_info import (
            ForwardBatch,
            PPProxyTensors,
        )
    except ImportError:
        _available = False

    @property
    def name(self) -> str:
        return "sglang"

    def collect_parallel_info(self) -> dict:
        if not self._available:
            return {}

        info = {}

        try:
            info["tp_rank"] = self._dist.get_tensor_model_parallel_rank()
            info["tp_size"] = self._dist.get_tensor_model_parallel_world_size()
            info["pp_rank"] = self._dist.get_pipeline_model_parallel_rank()
            info["pp_size"] = self._dist.get_pipeline_model_parallel_world_size()
            info["moe_ep_rank"] = self._dist.get_moe_expert_parallel_rank()
            info["moe_ep_size"] = self._dist.get_moe_expert_parallel_world_size()
            info["moe_tp_rank"] = self._dist.get_moe_tensor_parallel_rank()
            info["moe_tp_size"] = self._dist.get_moe_tensor_parallel_world_size()
            info["moe_dp_rank"] = self._dist.get_moe_data_parallel_rank()
            info["moe_dp_size"] = self._dist.get_moe_data_parallel_world_size()
        except (AttributeError, AssertionError):
            info["distributed_error"] = True

        try:
            info["enable_dp_attention"] = self._dp_attn.is_dp_attention_enabled()
            info["attn_tp_rank"] = self._dp_attn.get_attention_tp_rank()
            info["attn_tp_size"] = self._dp_attn.get_attention_tp_size()
            info["attn_dp_rank"] = self._dp_attn.get_attention_dp_rank()
            info["attn_dp_size"] = self._dp_attn.get_attention_dp_size()
            info["local_attn_dp_rank"] = self._dp_attn.get_local_attention_dp_rank()
            info["local_attn_dp_size"] = self._dp_attn.get_local_attention_dp_size()
            info["attn_cp_rank"] = self._dp_attn.get_attention_cp_rank()
            info["attn_cp_size"] = self._dp_attn.get_attention_cp_size()
        except (AttributeError, AssertionError):
            info["dp_attention_error"] = True

        return info

    def convert_value(
        self, value: Any, *, skip_forward_batch: bool
    ) -> Optional[dict[str, Any]]:
        if not self._available:
            return None

        if isinstance(value, self.LogitsProcessorOutput):
            return {"next_token_logits": value.next_token_logits}
        if isinstance(value, self.ForwardBatch):
            if skip_forward_batch:
                return {}
            result = {
                "input_ids": value.input_ids,
                "seq_lens": value.seq_lens,
                "positions": value.positions,
                "req_pool_indices": value.req_pool_indices,
            }
            if value.rids is not None:
                result["rids"] = value.rids
            return result
        if isinstance(value, self.PPProxyTensors):
            return {k: v for k, v in value.tensors.items()}

        return None

    def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]:
        if hasattr(module, "layer_id"):
            return module.layer_id
        return None

    def core_fields(self) -> frozenset[str]:
        return frozenset(
            {"input_ids", "positions", "seq_lens", "req_pool_indices", "rids"}
        )

    def get_tokenizer_path(self) -> Optional[str]:
        if not self._available:
            return None

        try:
            from sglang.srt.server_args import get_global_server_args

            args = get_global_server_args()
            if args is None:
                return None

            return args.tokenizer_path
        except Exception:
            return None


class _MegatronPlugin(_FrameworkPlugin):
    _available = True
    try:
        from megatron.core import parallel_state as _mpu
        from megatron.core.packed_seq_params import PackedSeqParams
    except ImportError:
        _available = False

    @property
    def name(self) -> str:
        return "megatron"

    def collect_parallel_info(self) -> dict:
        if not self._available:
            return {}

        info = {}
        try:
            info["tp_rank"] = self._mpu.get_tensor_model_parallel_rank()
            info["tp_size"] = self._mpu.get_tensor_model_parallel_world_size()
            info["pp_rank"] = self._mpu.get_pipeline_model_parallel_rank()
            info["pp_size"] = self._mpu.get_pipeline_model_parallel_world_size()
            info["dp_rank"] = self._mpu.get_data_parallel_rank()
            info["dp_size"] = self._mpu.get_data_parallel_world_size()
            info["cp_rank"] = self._mpu.get_context_parallel_rank()
            info["cp_size"] = self._mpu.get_context_parallel_world_size()
            info["vpp_rank"] = self._mpu.get_virtual_pipeline_model_parallel_rank()
            info["vpp_size"] = (
                self._mpu.get_virtual_pipeline_model_parallel_world_size()
            )
            info["ep_rank"] = self._mpu.get_expert_model_parallel_rank()
            info["ep_size"] = self._mpu.get_expert_model_parallel_world_size()
            info["etp_rank"] = self._mpu.get_expert_tensor_parallel_rank()
            info["etp_size"] = self._mpu.get_expert_tensor_parallel_world_size()
            info["edp_rank"] = self._mpu.get_expert_data_parallel_rank()
            info["edp_size"] = self._mpu.get_expert_data_parallel_world_size()
            info["tcp_rank"] = self._mpu.get_tensor_and_context_parallel_rank()
            info["tcp_size"] = self._mpu.get_tensor_and_context_parallel_world_size()
            info["etmp_rank"] = self._mpu.get_expert_tensor_and_model_parallel_rank()
            info["etmp_size"] = (
                self._mpu.get_expert_tensor_and_model_parallel_world_size()
            )
            info["tp_src_rank"] = self._mpu.get_tensor_model_parallel_src_rank()
            info["mp_src_rank"] = self._mpu.get_model_parallel_src_rank()
            info["dp_src_rank"] = self._mpu.get_data_parallel_src_rank()
        except (AttributeError, AssertionError):
            info["megatron_error"] = True

        # Megatron sequence parallel reuses the TP group (no dedicated parallel state API).
        # When sequence_parallel=True, inject sp_rank/sp_size for the comparator unsharder.
        try:
            from megatron.training.global_vars import get_args

            args = get_args()
            if getattr(args, "sequence_parallel", False) and "tp_rank" in info:
                info["sp_rank"] = info["tp_rank"]
                info["sp_size"] = info["tp_size"]
        except (ImportError, AssertionError, AttributeError):
            pass

        return info

    def convert_value(
        self, value: Any, *, skip_forward_batch: bool
    ) -> Optional[dict[str, Any]]:
        if not self._available:
            return None
        if isinstance(value, self.PackedSeqParams):
            return {
                "cu_seqlens_q": value.cu_seqlens_q,
                "cu_seqlens_kv": value.cu_seqlens_kv,
                "qkv_format": value.qkv_format,
            }
        return None

    def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]:
        if hasattr(module, "layer_number"):
            return module.layer_number - 1
        return None

    def core_fields(self) -> frozenset[str]:
        return frozenset(
            {"input_ids", "position_ids", "cu_seqlens_q", "cu_seqlens_kv", "qkv_format"}
        )

    def detect_recompute_status(self) -> _RecomputeStatus:
        if not self._available:
            return _RecomputeStatus.DISABLED
        try:
            from megatron.core.tensor_parallel.random import is_checkpointing

            if not is_checkpointing():
                return _RecomputeStatus.DISABLED
            if torch.is_grad_enabled():
                return _RecomputeStatus.RECOMPUTE
            return _RecomputeStatus.ORIGINAL
        except (ImportError, AttributeError):
            return _RecomputeStatus.DISABLED


_plugins: list[_FrameworkPlugin] = [_SGLangPlugin(), _MegatronPlugin()]


def _detect_recompute_status() -> _RecomputeStatus:
    for plugin in _plugins:
        info = plugin.detect_recompute_status()
        if info != _RecomputeStatus.DISABLED:
            return info
    return _RecomputeStatus.DISABLED


# -------------------------------------- singleton ------------------------------------------


dumper = _Dumper(config=DumperConfig.from_env())


# -------------------------------------- other utility functions ------------------------------------------


def get_truncated_value(value):
    if value is None:
        return None

    if isinstance(value, tuple):
        return [get_truncated_value(x) for x in value]

    if not isinstance(value, torch.Tensor):
        return value

    if value.numel() < 200:
        return value

    slices = [slice(0, 5) if dim_size > 50 else slice(None) for dim_size in value.shape]
    return value[tuple(slices)]


def get_tensor_info(x):
    """
    from sglang.srt.debug_utils.dumper import get_tensor_info
    """
    if not isinstance(x, torch.Tensor):
        return f"type={type(x)} value={x}"
    min = x.float().min() if x.numel() > 0 else None
    max = x.float().max() if x.numel() > 0 else None
    mean = x.float().mean() if x.numel() > 0 else None
    torch.set_printoptions(precision=10)
    x_sample_head = str(x.flatten()[:5])
    x_sample_tail = str(x.flatten()[-5:])
    torch.set_printoptions(precision=4)
    return (
        f"type={type(x)} "
        f"shape={x.shape} "
        f"dtype={x.dtype} "
        f"device={x.device} "
        f"stride={x.stride()} "
        f"req_grad={x.requires_grad} "
        f"min={min} "
        f"max={max} "
        f"mean={mean} "
        f"x_sample_head={x_sample_head} "
        f"x_sample_tail={x_sample_tail}"
    )
