"""Kernel API crash debugging helpers for SGLang.

This module was developed with reference to FlashInfer's kernel API logging utility:
https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/api_logging.py
"""

from __future__ import annotations

import fnmatch
import functools
import inspect
import json
import logging
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, TypeVar, overload

import torch

_logger = logging.getLogger("sglang.kernel_api")

_T = TypeVar("_T")
_F = TypeVar("_F", bound=Callable[..., Any])


def _str_with_pid(path: str) -> str:
    if "%i" in path:
        return path.replace("%i", str(os.getpid()))
    return path


def _get_env(key: str, type: Callable[..., _T], default: _T) -> _T:
    value_str = os.environ.get(key, None)
    if value_str is None:
        return default
    try:
        return type(value_str)
    except Exception:
        _logger.warning(
            "Failed to parse environment variable %s=%r as %s, using default %r",
            key,
            value_str,
            type.__name__,
            default,
        )
        return default


def _parse_pattern(value: str) -> list[str]:
    return [p.strip() for p in value.split(",") if p.strip()]


_KERNEL_API_LOG_LEVEL = _get_env("SGLANG_KERNEL_API_LOGLEVEL", int, 0)
_KERNEL_API_LOG_DEST = _get_env("SGLANG_KERNEL_API_LOGDEST", _str_with_pid, "stdout")
_DUMP_DIR = Path(
    _get_env("SGLANG_KERNEL_API_DUMP_DIR", _str_with_pid, "sglang_kernel_api_dumps")
)
_DUMP_INCLUDE_PATTERNS = _get_env("SGLANG_KERNEL_API_DUMP_INCLUDE", _parse_pattern, [])
_DUMP_EXCLUDE_PATTERNS = _get_env("SGLANG_KERNEL_API_DUMP_EXCLUDE", _parse_pattern, [])

_dump_call_counter: dict[str, int] = {}


def _setup_logger() -> None:
    for handler in list(_logger.handlers):
        _logger.removeHandler(handler)
        try:
            handler.close()
        except Exception:
            pass

    if _KERNEL_API_LOG_LEVEL == 0:
        _logger.addHandler(logging.NullHandler())
        _logger.setLevel(logging.CRITICAL + 1)
        return

    _logger.setLevel(logging.DEBUG)

    if _KERNEL_API_LOG_DEST == "stdout":
        handler = logging.StreamHandler(sys.stdout)
    elif _KERNEL_API_LOG_DEST == "stderr":
        handler = logging.StreamHandler(sys.stderr)
    else:
        handler = logging.FileHandler(_KERNEL_API_LOG_DEST, mode="a")

    handler.setFormatter(logging.Formatter("%(message)s"))
    _logger.addHandler(handler)
    _logger.propagate = False


_setup_logger()


def _is_compiling() -> bool:
    try:
        if hasattr(torch, "compiler") and hasattr(torch.compiler, "is_compiling"):
            return bool(torch.compiler.is_compiling())
        if hasattr(torch, "_dynamo") and hasattr(torch._dynamo, "is_compiling"):
            return bool(torch._dynamo.is_compiling())
    except Exception:
        return False
    return False


def _timestamp() -> str:
    return datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")


def _is_cuda_graph_capture_active() -> bool:
    try:
        return torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()
    except Exception:
        return False


def _append_line(lines: list[str], indent: int, text: str) -> None:
    lines.append(" " * indent + text)


def _should_dump_function(func_name: str) -> bool:
    if _DUMP_INCLUDE_PATTERNS and not any(
        fnmatch.fnmatch(func_name, pattern) for pattern in _DUMP_INCLUDE_PATTERNS
    ):
        return False
    if _DUMP_EXCLUDE_PATTERNS and any(
        fnmatch.fnmatch(func_name, pattern) for pattern in _DUMP_EXCLUDE_PATTERNS
    ):
        return False
    return True


def _serialize_tensor(tensor: torch.Tensor) -> list[str]:
    lines = ["Tensor("]
    _append_line(lines, 2, f"shape={tuple(tensor.shape)}")
    _append_line(lines, 2, f"dtype={tensor.dtype}")
    _append_line(lines, 2, f"device={tensor.device}")
    _append_line(lines, 2, f"requires_grad={tensor.requires_grad}")
    _append_line(lines, 2, f"is_contiguous={tensor.is_contiguous()}")

    if _KERNEL_API_LOG_LEVEL >= 5:
        if tensor.numel() == 0:
            _append_line(lines, 2, "statistics=[empty tensor]")
        elif tensor.device.type == "cuda" and _is_cuda_graph_capture_active():
            _append_line(
                lines, 2, "statistics=[skipped: CUDA graph capture in progress]"
            )
        else:
            try:
                detached = tensor.detach()
                if detached.is_complex():
                    stats_source = detached.abs().float()
                    nan_count = int(torch.isnan(detached).sum().item())
                    inf_count = int(torch.isinf(detached).sum().item())
                else:
                    stats_source = detached.float()
                    if detached.is_floating_point():
                        nan_count = int(torch.isnan(detached).sum().item())
                        inf_count = int(torch.isinf(detached).sum().item())
                    else:
                        nan_count = 0
                        inf_count = 0

                _append_line(lines, 2, f"min={stats_source.min().item():.6f}")
                _append_line(lines, 2, f"max={stats_source.max().item():.6f}")
                _append_line(lines, 2, f"mean={stats_source.mean().item():.6f}")
                _append_line(lines, 2, f"nan_count={nan_count}")
                _append_line(lines, 2, f"inf_count={inf_count}")
            except Exception as exc:
                _append_line(
                    lines, 2, f"statistics=[unavailable: {type(exc).__name__}]"
                )

    lines.append(")")
    return lines


def _serialize_value(value: Any, depth: int = 0) -> list[str]:
    if depth >= 2:
        return [f"{type(value).__name__}(...)"]

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

    if isinstance(value, (str, int, float, bool, type(None))):
        return [repr(value)]

    if isinstance(value, (list, tuple)):
        opener = "[" if isinstance(value, list) else "("
        closer = "]" if isinstance(value, list) else ")"
        lines = [opener]
        for idx, item in enumerate(value[:4]):
            item_lines = _serialize_value(item, depth + 1)
            lines.append(f"  [{idx}] {item_lines[0]}")
            for extra in item_lines[1:]:
                lines.append(f"      {extra}")
        if len(value) > 4:
            lines.append(f"  ... ({len(value) - 4} more items)")
        lines.append(closer)
        return lines

    if isinstance(value, dict):
        lines = ["{"]
        items = list(value.items())
        for key, item in items[:8]:
            item_lines = _serialize_value(item, depth + 1)
            lines.append(f"  {key!r}: {item_lines[0]}")
            for extra in item_lines[1:]:
                lines.append(f"      {extra}")
        if len(items) > 8:
            lines.append(f"  ... ({len(items) - 8} more items)")
        lines.append("}")
        return lines

    summary = [f"{type(value).__name__}("]
    for attr in ("shape", "dtype", "device"):
        if hasattr(value, attr):
            try:
                _append_line(summary, 2, f"{attr}={getattr(value, attr)}")
            except Exception:
                pass
    if len(summary) == 1:
        _append_line(summary, 2, f"repr={repr(value)[:200]}")
    summary.append(")")
    return summary


def _serialize_json_value(value: Any) -> Any:
    if isinstance(value, torch.dtype):
        return {"type": "torch.dtype", "value": str(value)}
    if isinstance(value, (str, int, float, bool, type(None))):
        return value
    if isinstance(value, (list, tuple)):
        return [_serialize_json_value(item) for item in value[:16]]
    if isinstance(value, dict):
        return {
            str(key): _serialize_json_value(item)
            for key, item in list(value.items())[:32]
        }
    return {"type": type(value).__name__, "repr": repr(value)[:200]}


def _collect_dump_entries(
    prefix: str,
    value: Any,
    tensor_entries: dict[str, torch.Tensor],
    metadata_entries: dict[str, Any],
) -> None:
    if isinstance(value, torch.Tensor):
        tensor_entries[prefix] = value.detach().cpu()
        return

    if isinstance(value, (list, tuple)):
        for idx, item in enumerate(value):
            _collect_dump_entries(
                f"{prefix}_{idx}", item, tensor_entries, metadata_entries
            )
        metadata_entries[f"{prefix}__container"] = {
            "type": type(value).__name__,
            "length": len(value),
        }
        return

    if isinstance(value, dict):
        for key, item in value.items():
            _collect_dump_entries(
                f"{prefix}_{str(key)}", item, tensor_entries, metadata_entries
            )
        metadata_entries[f"{prefix}__container"] = {
            "type": "dict",
            "keys": [str(k) for k in value.keys()],
        }
        return

    metadata_entries[prefix] = _serialize_json_value(value)


def _dump_metadata_path(dump_dir: Path) -> Path:
    return dump_dir / "metadata.json"


def _write_dump_metadata(dump_dir: Path, metadata: dict[str, Any]) -> None:
    _dump_metadata_path(dump_dir).write_text(json.dumps(metadata, indent=2))


def _read_dump_metadata(dump_dir: Path) -> dict[str, Any]:
    return json.loads(_dump_metadata_path(dump_dir).read_text())


def _dump_function_inputs(
    func_name: str, args: tuple[Any, ...], kwargs: dict[str, Any]
) -> Path | None:
    if not _should_dump_function(func_name):
        return None

    _DUMP_DIR.mkdir(parents=True, exist_ok=True)
    call_index = _dump_call_counter.get(func_name, 0) + 1
    _dump_call_counter[func_name] = call_index
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
    safe_func_name = func_name.replace("/", "_").replace("<", "_").replace(">", "_")
    dump_dir = (
        _DUMP_DIR
        / f"{timestamp}_pid{os.getpid()}_{safe_func_name}_call{call_index:04d}"
    )
    dump_dir.mkdir(parents=True, exist_ok=True)

    tensor_entries: dict[str, torch.Tensor] = {}
    metadata_entries: dict[str, Any] = {}
    for idx, arg in enumerate(args):
        _collect_dump_entries(f"arg_{idx}", arg, tensor_entries, metadata_entries)
    for key, value in kwargs.items():
        _collect_dump_entries(f"kwarg_{key}", value, tensor_entries, metadata_entries)

    if tensor_entries:
        torch.save(tensor_entries, dump_dir / "inputs.pt")

    metadata = {
        "function_name": func_name,
        "timestamp": timestamp,
        "process_id": os.getpid(),
        "execution_status": "inputs_saved",
        "input_metadata": metadata_entries,
        "input_tensor_keys": list(tensor_entries.keys()),
        "output_metadata": {},
        "output_tensor_keys": [],
    }
    _write_dump_metadata(dump_dir, metadata)
    _logger.debug("Dumped inputs to: %s", dump_dir)
    return dump_dir


def _dump_function_outputs(dump_dir: Path, result: Any) -> None:
    tensor_entries: dict[str, torch.Tensor] = {}
    metadata_entries: dict[str, Any] = {}
    _collect_dump_entries("result", result, tensor_entries, metadata_entries)
    if tensor_entries:
        torch.save(tensor_entries, dump_dir / "outputs.pt")

    metadata = _read_dump_metadata(dump_dir)
    metadata["execution_status"] = "completed"
    metadata["output_metadata"] = metadata_entries
    metadata["output_tensor_keys"] = list(tensor_entries.keys())
    _write_dump_metadata(dump_dir, metadata)
    _logger.debug("Dumped outputs to: %s", dump_dir)


def _mark_dump_exception(dump_dir: Path, exc: Exception) -> None:
    metadata = _read_dump_metadata(dump_dir)
    metadata["execution_status"] = "exception"
    metadata["exception"] = {
        "type": type(exc).__name__,
        "message": str(exc),
    }
    _write_dump_metadata(dump_dir, metadata)


def _log_section(title: str, data: dict[str, Any]) -> None:
    _logger.debug(title)
    for key, value in data.items():
        lines = _serialize_value(value)
        _logger.debug("  %s=%s", key, lines[0])
        for line in lines[1:]:
            _logger.debug("    %s", line)


def _infer_func_name(func: Callable) -> str:
    qualname = getattr(func, "__qualname__", getattr(func, "__name__", "unknown"))
    qualname = qualname.replace(".<locals>.", ".").replace("<locals>.", "")

    module = getattr(func, "__module__", "")
    for prefix in ("sglang.", "sgl_kernel."):
        if module.startswith(prefix):
            module = module[len(prefix) :]
            break

    if module and module not in {"__main__", "builtins"}:
        return f"{module}.{qualname}"

    source_path = inspect.getsourcefile(func)
    if source_path is not None:
        return f"{Path(source_path).stem}.{qualname}"

    return qualname


@overload
def debug_kernel_api(
    func: _F,
    *,
    op_name: str | None = None,
) -> _F: ...


@overload
def debug_kernel_api(
    *,
    op_name: str | None = None,
) -> Callable[[_F], _F]: ...


def debug_kernel_api(
    func: Callable | None = None,
    *,
    op_name: str | None = None,
) -> Callable:
    # NOTE: avoid any overhead in the hot path when logging is disabled
    if _KERNEL_API_LOG_LEVEL == 0:
        if func is None:
            return lambda f: f
        return func

    def decorator(f: Callable) -> Callable:
        if hasattr(f, "_debug_kernel_wrapped"):
            return f

        @functools.wraps(f)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            if _is_compiling():
                return f(*args, **kwargs)

            func_name = op_name or _infer_func_name(f)
            dump_dir: Path | None = None
            positional_args = args
            try:
                parameters = tuple(inspect.signature(f).parameters.values())
            except (TypeError, ValueError):
                parameters = ()
            if args and parameters and parameters[0].name in {"self", "cls"}:
                positional_args = args[1:]
            _logger.debug("=" * 80)
            _logger.debug("%s SGLang Kernel API Call: %s", _timestamp(), func_name)

            if _KERNEL_API_LOG_LEVEL >= 3:
                if positional_args:
                    _log_section(
                        "Positional input arguments:",
                        {f"arg[{idx}]": arg for idx, arg in enumerate(positional_args)},
                    )
                if kwargs:
                    _log_section("Keyword input arguments:", kwargs)

            if _KERNEL_API_LOG_LEVEL >= 10:
                if _is_cuda_graph_capture_active():
                    _logger.debug("Tensor dump skipped: CUDA graph capture in progress")
                else:
                    dump_dir = _dump_function_inputs(func_name, positional_args, kwargs)

            try:
                result = f(*args, **kwargs)
            except Exception as exc:
                if dump_dir is not None:
                    _mark_dump_exception(dump_dir, exc)
                _logger.debug(
                    "%s SGLang Kernel API Exception: %s (%s: %s)",
                    _timestamp(),
                    func_name,
                    type(exc).__name__,
                    exc,
                )
                raise

            if dump_dir is not None:
                _dump_function_outputs(dump_dir, result)
            if _KERNEL_API_LOG_LEVEL >= 3:
                _log_section("Output:", {"return": result})
            return result

        setattr(wrapper, "_debug_kernel_wrapped", True)
        return wrapper

    return decorator if func is None else decorator(func)


def debug_torch_op(
    op_func: Callable,
    op_name: str,
    *,
    namespace: str = "sglang",
) -> Callable:
    """NOTE: For internal use. Prefer `debug_kernel_api` for general use cases."""
    # NOTE: avoid any overhead in the hot path when logging is disabled
    impl = getattr(getattr(torch.ops, namespace), op_name)
    if _KERNEL_API_LOG_LEVEL == 0:
        return impl
    # NOTE: propagate the marker to avoid double-wrapping
    if hasattr(op_func, "_debug_kernel_wrapped"):
        setattr(impl, "_debug_kernel_wrapped", True)
        return impl
    # NOTE: redirect the function name
    return debug_kernel_api(impl, op_name=_infer_func_name(op_func))


def wrap_method_with_debug_kernel_once(
    obj: Any,
    method_name: str,
    *,
    op_name: str,
    marker_attr: str | None = None,
) -> Any:
    # NOTE: avoid any overhead in the hot path when logging is disabled
    if _KERNEL_API_LOG_LEVEL == 0:
        return obj

    if marker_attr is None:
        marker_attr = f"_debug_kernel_{method_name}_wrapped"

    if getattr(obj, marker_attr, False):
        return obj

    setattr(
        obj,
        method_name,
        debug_kernel_api(getattr(obj, method_name), op_name=op_name),
    )
    setattr(obj, marker_attr, True)
    return obj