#  Copyright (c) Meta Platforms, Inc. and affiliates.

import atexit
import fnmatch
import gzip
import hashlib
import importlib
import inspect
import io
import json
import logging
import math
import os
import subprocess
import tempfile
from collections import defaultdict
from collections.abc import Mapping
from dataclasses import asdict, is_dataclass
from datetime import date, datetime
from enum import Enum
from functools import partial
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

from triton.knobs import JITHook, LaunchHook

from .shared_vars import DEFAULT_TRACE_FILE_PREFIX


log = logging.getLogger(__name__)

TEXT_FILE_EXTENSIONS = [".ttir", ".ttgir", ".llir", ".ptx", ".amdgcn", ".json"]
MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB limit for file content extraction

triton_trace_log = logging.getLogger("tritonparse_trace")
# The folder to store the triton trace log.
triton_trace_folder = os.environ.get("TRITON_TRACE", None)
# Enable debug logging for tritonparse itself
TRITONPARSE_DEBUG = os.getenv("TRITONPARSE_DEBUG", None) in ["1", "true", "True"]
# Kernel allowlist for filtering traced kernels. Use comma separated list of fnmatch patterns.
TRITONPARSE_KERNEL_ALLOWLIST = os.environ.get("TRITONPARSE_KERNEL_ALLOWLIST", None)
# Parsed kernel allowlist patterns (set during init)
_KERNEL_ALLOWLIST_PATTERNS: Optional[List[str]] = None
# Enable launch trace. WARNNING: it will overwrite launch_metadata function for each triton kernel.
TRITON_TRACE_LAUNCH = os.getenv("TRITON_TRACE_LAUNCH", None) in ["1", "true", "True"]
# Enable more tensor information collection in trace logs.
TRITONPARSE_MORE_TENSOR_INFORMATION = os.getenv(
    "TRITONPARSE_MORE_TENSOR_INFORMATION", None
) in ["1", "true", "True"]
# Enable full Python source file extraction instead of just the function definition
TRITON_FULL_PYTHON_SOURCE = os.getenv("TRITON_FULL_PYTHON_SOURCE", "0") in [
    "1",
    "true",
    "True",
]
# Compression algorithm for raw trace files
# When enabled, each JSON record is written as a separate gzip member,
# concatenated in sequence within a .bin.ndjson file.
# Supported values: "none", "gzip", "clp", "zstd" (future)
# - "gzip": Outputs .bin.ndjson (gzip member concatenation format)
# - "clp": Outputs .clp (Compressed Log Processor format)
# - "none": Outputs .ndjson (plain text)
# If TRITON_TRACE_COMPRESSION is explicitly set, respect that value;
# otherwise, default to "gzip" for better storage efficiency.
_compression_env = os.getenv("TRITON_TRACE_COMPRESSION")
TRITON_TRACE_COMPRESSION = _compression_env if _compression_env is not None else "gzip"
if TRITON_TRACE_COMPRESSION == "clp":
    from .clp import clp_open
# Maximum file size for full source extraction (default 10MB)
TRITON_MAX_SOURCE_SIZE = int(os.getenv("TRITON_MAX_SOURCE_SIZE", str(10 * 1024 * 1024)))
# Inductor compiled kernel's launch tracing needs this flag to be set.
# If TRITON_TRACE_LAUNCH is enabled, also enable TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK
TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK = (
    os.getenv("TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK", None) in ["1", "true", "True"]
    or TRITON_TRACE_LAUNCH
)
# Enable NVIDIA SASS dump. It requires the CUBIN file to be localable.
# WARNNING: it will slow down the compilation significantly.
TRITONPARSE_DUMP_SASS = os.getenv("TRITONPARSE_DUMP_SASS", None) in [
    "1",
    "true",
    "True",
]

# The flag to mark if launch is traced. It is used to avoid initilizing the launch hook twice.
_trace_launch_enabled = False
# Enable tensor blob storage
TRITONPARSE_SAVE_TENSOR_BLOBS = os.getenv("TRITONPARSE_SAVE_TENSOR_BLOBS", "0") in [
    "1",
    "true",
    "True",
]
# Tensor size limit in bytes (default 10GB)
TRITONPARSE_TENSOR_SIZE_LIMIT = int(
    os.getenv("TRITONPARSE_TENSOR_SIZE_LIMIT", str(10 * 1024 * 1024 * 1024))
)
# Tensor storage quota in bytes (default 100GB) - tracks compressed size for current run
TRITONPARSE_TENSOR_STORAGE_QUOTA = int(
    os.getenv("TRITONPARSE_TENSOR_STORAGE_QUOTA", str(100 * 1024 * 1024 * 1024))
)
# Compression threshold in bytes (default 1MB) - only compress blobs >= this size
TRITONPARSE_COMPRESSION_THRESHOLD = 1 * 1024 * 1024
# Compression level for gzip (0-9, higher = better compression but slower)
TRITONPARSE_COMPRESSION_LEVEL = 4
# Log statistics every N saved blobs
TRITONPARSE_STATS_LOG_FREQUENCY = 100

TRITON_TRACE_HANDLER = None
# Global tensor blob manager instance
TENSOR_BLOB_MANAGER = None

if importlib.util.find_spec("torch") is not None:
    TORCH_INSTALLED = True
    import torch
    from torch.utils._traceback import CapturedTraceback
else:
    TORCH_INSTALLED = False


class TensorBlobManager:
    """
    Manager for storing tensor data as content-addressed blobs.

    Uses BLAKE2b hashing for content addressing and stores blobs in a two-level
    directory structure to avoid filesystem limitations with large numbers of files.
    """

    def __init__(
        self,
        root_dir: Optional[str] = None,
        storage_quota: Optional[int] = None,
    ):
        self.root_dir = None
        self.hash_to_path_cache = {}  # In-memory cache for hash -> path mapping
        self.compression_threshold = TRITONPARSE_COMPRESSION_THRESHOLD
        self.storage_quota = (
            storage_quota
            if storage_quota is not None
            else TRITONPARSE_TENSOR_STORAGE_QUOTA
        )

        # Resource statistics (tracks current run only)
        self.total_compressed_bytes = 0  # Total compressed size written in this run
        self.total_uncompressed_bytes = (
            0  # Total uncompressed size (for compression ratio)
        )
        self.blob_count = 0  # Total blob references (including dedup hits)
        self.blob_saved_count = 0  # Actual blobs saved (excluding dedup hits)
        self.storage_disabled = False  # Whether storage has been disabled due to quota
        self.storage_disabled_reason = None  # Reason for disabling storage

        if root_dir:
            self.set_root_dir(root_dir)

    def set_root_dir(self, root_dir: str):
        """Set the root directory for blob storage."""
        self.root_dir = Path(root_dir) / "saved_tensors"
        self.root_dir.mkdir(parents=True, exist_ok=True)
        log.debug(f"TensorBlobManager: using root directory {self.root_dir}")

    def _compute_hash(self, data: bytes) -> str:
        """Compute BLAKE2b hash of the data."""
        return hashlib.blake2b(data).hexdigest()

    def _get_blob_path(self, hash_hex: str, extension: str = ".bin.gz") -> Path:
        """Get the file path for a given hash using two-level directory structure."""
        if not self.root_dir:
            raise ValueError("Root directory not set")

        # Two-level directory: first 2 chars / full_hash{extension}
        subdir = hash_hex[:2]
        filename = f"{hash_hex}{extension}"
        return (self.root_dir / subdir / filename).resolve()

    def _get_tensor_size_bytes(self, tensor) -> int:
        """Get tensor size in bytes before serialization."""
        if hasattr(tensor, "numel") and hasattr(tensor, "element_size"):
            return tensor.numel() * tensor.element_size()
        return 0

    def _log_statistics(self, final: bool = False):
        """Print statistics about tensor blob storage.

        Args:
            final: If True, this is the final statistics message (e.g., when storage is disabled)
        """
        prefix = "📊 Final" if final else "📊"
        compression_ratio = (
            self.total_uncompressed_bytes / max(1, self.total_compressed_bytes)
            if self.total_compressed_bytes > 0
            else 0.0
        )
        dedup_count = self.blob_count - self.blob_saved_count

        log.info(
            f"{prefix} Tensor blob stats: "
            f"{self.blob_saved_count} saved ({self.blob_count} total, {dedup_count} dedup), "
            f"{self.total_compressed_bytes / 1024**3:.2f}GB compressed "
            f"({self.total_uncompressed_bytes / 1024**3:.2f}GB uncompressed), "
            f"compression ratio: {compression_ratio:.2f}x"
        )

    def _disable_storage(self, reason: str):
        """Disable blob storage and log warning with statistics.

        Args:
            reason: The reason why storage is being disabled
        """
        if not self.storage_disabled:  # Only disable once
            self.storage_disabled = True
            self.storage_disabled_reason = reason
            log.warning(f"⚠️  TENSOR BLOB STORAGE DISABLED: {reason}")
            self._log_statistics(final=True)

    def save_tensor_blob(self, tensor) -> Dict[str, Any]:
        """
        Save tensor as a blob and return metadata.

        Args:
            tensor: PyTorch tensor to save

        Returns:
            Dictionary with blob metadata or error information:
            - Success: {'tensor_hash': str, 'blob_path': str, 'blob_size': int,
                       'blob_size_uncompressed': int, 'compression': str,
                       'compression_ratio': float, 'serialization_method': str}
            - Dedup hit: Same as success but from cache (not counted in quota)
            - Error: {'error': str, 'tensor_hash': None}
        """
        # Early exit: Check if storage is disabled
        if self.storage_disabled:
            return {"error": self.storage_disabled_reason, "tensor_hash": None}

        # Early exit: Check if root directory is set
        if not self.root_dir:
            return {"error": "Blob storage not initialized", "tensor_hash": None}

        try:
            # Check tensor size before serialization
            tensor_size = self._get_tensor_size_bytes(tensor)
            if tensor_size > TRITONPARSE_TENSOR_SIZE_LIMIT:
                log.warning(
                    f"Tensor size {tensor_size} bytes exceeds limit {TRITONPARSE_TENSOR_SIZE_LIMIT} bytes, skipping blob storage"
                )
                return {
                    "error": f"Tensor size {tensor_size} bytes exceeds limit {TRITONPARSE_TENSOR_SIZE_LIMIT} bytes",
                    "tensor_hash": None,
                }

            # Serialize tensor using torch.save
            import io

            buffer = io.BytesIO()
            if TORCH_INSTALLED:
                torch.save(tensor.cpu(), buffer)
            else:
                return {
                    "error": "PyTorch not available for tensor serialization",
                    "tensor_hash": None,
                }

            blob_data = buffer.getvalue()
            uncompressed_size = len(blob_data)

            # Compute hash on uncompressed data for content addressing
            hash_hex = self._compute_hash(blob_data)

            # Check for deduplication (before compression to save work)
            if hash_hex in self.hash_to_path_cache:
                blob_path = self.hash_to_path_cache[hash_hex]
                try:
                    # Try to access the file - handles race condition where file might be deleted
                    disk_size = blob_path.stat().st_size
                    compression = (
                        "gzip" if str(blob_path).endswith(".bin.gz") else "none"
                    )
                    compression_ratio = uncompressed_size / max(1, disk_size)

                    # Deduplication hit - increment count but don't add to quota
                    self.blob_count += 1

                    return {
                        "tensor_hash": hash_hex,
                        "blob_path": str(blob_path),
                        "blob_size": disk_size,
                        "blob_size_uncompressed": uncompressed_size,
                        "compression": compression,
                        "compression_ratio": compression_ratio,
                        "serialization_method": "torch_save",
                        "deduplicated": True,
                    }
                except (FileNotFoundError, OSError):
                    # File was deleted or inaccessible - remove from cache and continue to save
                    log.debug(
                        f"Cached blob file no longer exists: {blob_path}, will re-save"
                    )
                    self.hash_to_path_cache.pop(hash_hex, None)

            # Decide whether to compress based on size threshold
            if uncompressed_size >= self.compression_threshold:
                # Compress the data
                data_to_write = gzip.compress(
                    blob_data, compresslevel=TRITONPARSE_COMPRESSION_LEVEL
                )
                file_extension = ".bin.gz"
                compression = "gzip"
            else:
                # Don't compress small files (overhead not worth it)
                data_to_write = blob_data
                file_extension = ".bin"
                compression = "none"

            disk_size = len(data_to_write)

            # Check quota BEFORE writing
            if self.total_compressed_bytes + disk_size > self.storage_quota:
                self._disable_storage(
                    f"Storage quota would be exceeded: "
                    f"{(self.total_compressed_bytes + disk_size) / 1024**3:.2f}GB > "
                    f"{self.storage_quota / 1024**3:.2f}GB limit"
                )
                return {"error": self.storage_disabled_reason, "tensor_hash": None}

            # Create blob file path with appropriate extension
            blob_path = self._get_blob_path(hash_hex, extension=file_extension)
            blob_path.parent.mkdir(parents=True, exist_ok=True)

            # Atomic write using temporary file + rename
            with tempfile.NamedTemporaryFile(
                mode="wb",
                dir=blob_path.parent,
                prefix=f".tmp_{hash_hex}_",
                delete=False,
            ) as tmp_file:
                tmp_file.write(data_to_write)
                tmp_path = Path(tmp_file.name)

            # Atomic rename
            tmp_path.rename(blob_path)

            # Update cache and statistics
            self.hash_to_path_cache[hash_hex] = blob_path
            self.total_compressed_bytes += disk_size
            self.total_uncompressed_bytes += uncompressed_size
            self.blob_count += 1
            self.blob_saved_count += 1

            # Log progress periodically
            if self.blob_saved_count % TRITONPARSE_STATS_LOG_FREQUENCY == 0:
                self._log_statistics()

            log.debug(
                f"Saved tensor blob: {hash_hex} -> {blob_path} ({disk_size} bytes, compression={compression})"
            )

            compression_ratio = uncompressed_size / max(1, disk_size)

            return {
                "tensor_hash": hash_hex,
                "blob_path": str(blob_path),
                "blob_size": disk_size,
                "blob_size_uncompressed": uncompressed_size,
                "compression": compression,
                "compression_ratio": compression_ratio,
                "serialization_method": "torch_save",
            }

        except OSError as e:
            # Disk full, permission errors, etc. - disable storage to avoid repeated failures
            error_msg = f"Failed to save tensor blob (I/O error): {str(e)}"
            log.error(error_msg)
            self._disable_storage(error_msg)
            return {"error": error_msg, "tensor_hash": None}
        except Exception as e:
            # Other unexpected errors - log but don't disable storage
            error_msg = f"Failed to save tensor blob: {str(e)}"
            log.error(error_msg)
            return {"error": error_msg, "tensor_hash": None}


class TritonLogRecord(logging.LogRecord):
    """
    Custom LogRecord class for structured logging of Triton operations.

    Extends the standard LogRecord with additional attributes for storing
    structured metadata and payload information.
    """

    def __init__(
        self,
        name,
        level,
        pathname,
        lineno,
        msg,
        args,
        exc_info,
        metadata=None,
        payload=None,
        **kwargs,
    ):
        super().__init__(name, level, pathname, lineno, msg, args, exc_info, **kwargs)
        self.metadata: Dict[str, Any] = metadata or {}
        self.payload: Optional[Union[str, Dict[str, Any], list]] = payload


def create_triton_log_record(
    name=None,
    level=logging.DEBUG,
    pathname=None,
    lineno=None,
    msg="",
    args=(),
    exc_info=None,
    metadata=None,
    payload=None,
    **kwargs,
):
    """
    Factory method to create TritonLogRecord instances with sensible defaults.

    Args:
        name (str, optional): Logger name. Defaults to triton_trace_log.name.
        level (int, optional): Log level. Defaults to DEBUG.
        pathname (str, optional): Path to the file where the log call was made. Defaults to current file.
        lineno (int, optional): Line number where the log call was made. Defaults to current line.
        msg (str, optional): Log message. Defaults to empty string.
        args (tuple, optional): Arguments to interpolate into the message. Defaults to empty tuple.
        exc_info (optional): Exception information. Defaults to None.
        metadata (Dict[str, Any], optional): Structured metadata for the log record. Defaults to empty dict.
        payload (optional): Payload data. Defaults to None.
        **kwargs: Additional keyword arguments for LogRecord

    Returns:
        TritonLogRecord: A custom log record with structured data
    """
    if pathname is None:
        pathname = __file__
    if lineno is None:
        lineno = inspect.currentframe().f_back.f_lineno
    if name is None:
        name = triton_trace_log.name

    record = TritonLogRecord(
        name,
        level,
        pathname,
        lineno,
        msg,
        args,
        exc_info,
        metadata=metadata,
        payload=payload,
        **kwargs,
    )
    return record


def convert(obj):
    """
    Recursively converts dataclasses, dictionaries, and lists to their serializable forms.

    Args:
        obj: The object to convert, which can be a dataclass instance, dictionary, list, or any other type

    Returns:
        A serializable version of the input object where dataclasses are converted to dictionaries
    """
    from triton.language.core import dtype

    # 1. primitives that JSON already supports  -------------------------------
    if obj is None or isinstance(obj, (bool, int, str)):
        return obj

    if isinstance(obj, float):
        # JSON spec forbids NaN/Infinity – keep precision but stay valid
        if math.isfinite(obj):
            return obj
        return str(obj)  # "NaN", "inf", "-inf"

    # 2. simple containers ----------------------------------------------------
    if isinstance(obj, (list, tuple)):
        # Handle namedtuple specially to preserve field names
        if hasattr(obj, "_asdict"):
            return convert(obj._asdict())
        return [convert(x) for x in obj]

    if isinstance(obj, (set, frozenset)):
        return [convert(x) for x in sorted(obj, key=str)]

    if isinstance(obj, Mapping):
        return {str(k): convert(v) for k, v in obj.items()}

    # 3. time, enum, path, bytes ---------------------------------------------
    if isinstance(obj, (datetime, date)):
        return obj.isoformat()

    if isinstance(obj, Enum):
        return convert(obj.value)

    if isinstance(obj, Path):
        return str(obj)

    if is_dataclass(obj):
        return convert(
            asdict(obj)
        )  # Convert dataclass to dict and then process that dict

    if _is_triton_kernels_layout(obj):
        layout_info = {"type": type(obj).__name__}
        if hasattr(obj, "initial_shape"):
            layout_info["initial_shape"] = convert(obj.initial_shape)
        if hasattr(obj, "name"):
            layout_info["name"] = convert(obj.name)
        return layout_info

    # 4. Common Triton constexpr objects
    if isinstance(obj, dtype):
        return f"triton.language.core.dtype('{str(obj)}')"

    if TORCH_INSTALLED and isinstance(obj, torch.dtype):
        return str(obj)

    log.warning(f"Unknown type: {type(obj)}")
    return str(obj)  # Return primitive types as-is


def _is_triton_kernels_layout(obj):
    """
    Check if an object is an instance of a Layout class from a
    triton_kernels module by checking its MRO.
    """
    t = type(obj)
    for base_class in t.__mro__:
        module_name = getattr(base_class, "__module__", "")
        type_name = getattr(base_class, "__name__", "")
        if type_name == "Layout" and module_name.startswith("triton_kernels"):
            return True
    return False


def _is_from_triton_kernels_module(obj):
    """
    Check if an object is an instance of Tensor or Storage from a
    triton_kernels module.
    """
    t = type(obj)
    module_name = getattr(t, "__module__", "")
    type_name = getattr(t, "__name__", "")
    return type_name in ("Tensor", "Storage") and module_name.startswith(
        "triton_kernels"
    )


def _log_torch_tensor_info(tensor_value):
    """
    Extracts metadata from a torch.Tensor object.

    Args:
        tensor_value (torch.Tensor): The tensor to extract information from.

    Returns:
        dict: A dictionary containing tensor metadata.
    """
    arg_info = {}
    arg_info["type"] = "tensor"
    arg_info["shape"] = list(tensor_value.shape)
    arg_info["dtype"] = str(tensor_value.dtype)
    arg_info["device"] = str(tensor_value.device)
    arg_info["stride"] = list(tensor_value.stride())
    arg_info["numel"] = tensor_value.numel()
    arg_info["is_contiguous"] = tensor_value.is_contiguous()
    arg_info["element_size"] = tensor_value.element_size()
    arg_info["storage_offset"] = tensor_value.storage_offset()
    # Memory usage in bytes
    arg_info["memory_usage"] = tensor_value.numel() * tensor_value.element_size()
    # Add data_ptr for memory tracking (optional)
    if hasattr(tensor_value, "data_ptr"):
        arg_info["data_ptr"] = hex(tensor_value.data_ptr())
    if TRITONPARSE_MORE_TENSOR_INFORMATION:
        try:
            # Convert to float for reliable statistics computation across all dtypes
            # This creates a new tensor without modifying the original
            float_tensor = tensor_value.float()
            arg_info["min"] = float_tensor.min().item()
            arg_info["max"] = float_tensor.max().item()
            arg_info["mean"] = float_tensor.mean().item()
            arg_info["std"] = float_tensor.std().item()
        except (RuntimeError, ValueError, TypeError) as e:
            log.error(f"Unable to compute tensor statistics: {e}")
            arg_info["tensor_capture_error"] = str(e)

    # Add tensor blob storage if enabled
    if TRITONPARSE_SAVE_TENSOR_BLOBS and TENSOR_BLOB_MANAGER is not None:
        blob_info = TENSOR_BLOB_MANAGER.save_tensor_blob(tensor_value)
        arg_info.update(blob_info)
    return arg_info


def maybe_enable_debug_logging():
    """
    This logging is for logging module itself, not for logging the triton compilation.
    """
    if TRITONPARSE_DEBUG:
        # Always set debug level if TRITONPARSE_DEBUG is set
        log.setLevel(logging.DEBUG)

        # Prevent propagation to root logger to avoid duplicate messages
        log.propagate = False

        # Check if we already have a debug handler
        has_debug_handler = any(
            isinstance(handler, logging.StreamHandler)
            and handler.level <= logging.DEBUG
            for handler in log.handlers
        )

        if not has_debug_handler:
            log_handler = logging.StreamHandler()
            log_handler.setLevel(logging.DEBUG)
            formatter = logging.Formatter("%(asctime)s[%(levelname)s] %(message)s")
            formatter.default_time_format = "%Y%m%d %H:%M:%S"
            formatter.default_msec_format = None
            log_handler.setFormatter(formatter)
            log.addHandler(log_handler)


def get_stack_trace(skip=1):
    """
    Get call stack trace for the current execution context.

    Extracts stack trace information using torch's CapturedTraceback utility,
    providing detailed information about each frame in the call stack.

    Args:
        skip (int): Number of frames to skip from the start of the stack

    Returns:
        List[Dict]: List of frame information dictionaries containing line numbers,
                   function names, filenames, and code snippets
    """
    if not TORCH_INSTALLED:
        return []
    frames = []
    for frame in CapturedTraceback.extract(skip=skip).summary():
        frames.append(
            {
                "line": frame.lineno,
                "name": frame.name,
                "filename": frame.filename,
                "loc": frame.line,
            }
        )
    return frames


def parse_kernel_allowlist() -> Optional[List[str]]:
    """
    Parse the kernel allowlist from environment variable.

    Returns:
        List[str] or None: List of kernel name patterns to trace, or None if all kernels should be traced
    """
    if not TRITONPARSE_KERNEL_ALLOWLIST:
        return None

    # Split by comma and strip whitespace
    patterns = [pattern.strip() for pattern in TRITONPARSE_KERNEL_ALLOWLIST.split(",")]
    # Filter out empty patterns
    patterns = [pattern for pattern in patterns if pattern]

    if not patterns:
        return None

    log.debug(f"Kernel allowlist patterns: {patterns}")
    return patterns


def extract_kernel_name(src) -> Optional[str]:
    """
    Extract kernel name from the source object.

    Args:
        src (Union[ASTSource, IRSource]): Source object containing kernel information

    Returns:
        str or None: Kernel name if extractable, None otherwise
    """
    from triton.compiler import IRSource

    try:
        if isinstance(src, IRSource):
            return src.getattr("name", None)
        else:
            # For ASTSource, get the function name
            if (
                hasattr(src, "fn")
                and hasattr(src.fn, "fn")
                and hasattr(src.fn.fn, "__name__")
            ):
                return src.fn.fn.__name__
            return None
    except Exception as e:
        log.warn(f"Error extracting kernel name: {e}")
        return None


def should_trace_kernel(
    kernel_name: Optional[str], allowlist_patterns: Optional[List[str]]
) -> bool:
    """
    Check if a kernel should be traced based on the allowlist.

    Args:
        kernel_name (str or None): Name of the kernel
        allowlist_patterns (List[str] or None): List of patterns to match against

    Returns:
        bool: True if the kernel should be traced, False otherwise
    """
    # If no allowlist is set, trace all kernels
    if allowlist_patterns is None:
        return True

    # If we can't extract kernel name, don't trace (conservative approach)
    if kernel_name is None:
        log.debug("Cannot extract kernel name, skipping trace")
        return False

    # Check if kernel name matches any pattern in the allowlist
    for pattern in allowlist_patterns:
        if fnmatch.fnmatch(kernel_name, pattern):
            log.debug(f"Kernel '{kernel_name}' matches pattern '{pattern}', will trace")
            return True

    log.debug(
        f"Kernel '{kernel_name}' does not match any allowlist pattern, skipping trace"
    )
    return False


def extract_python_source_info(trace_data: Dict[str, Any], source):
    """
    Extract Python source code information from the source object and add it to trace_data.

    This function uses Python's inspect module to extract source code information
    from the provided source object (typically an ASTSource or IRSource instance).
    It adds file path, line numbers, and the actual source code to the trace_data.

    By default, only the function definition is extracted. Set TRITON_FULL_PYTHON_SOURCE=1
    to extract the entire Python source file.
    @TODO: we should enable it by default in next diff and track the compilation time regression

    Environment Variables:
        TRITON_FULL_PYTHON_SOURCE: If set to "1", extract the full Python file
                                   instead of just the function definition.
        TRITON_MAX_SOURCE_SIZE: Maximum file size in bytes for full source extraction
                               (default: 10MB). Files larger than this will fall back
                               to function-only mode.

    Args:
        trace_data (Dict[str, Any]): Dictionary to store extracted information
        source (Union[ASTSource, IRSource]): Source object containing kernel function information
    """
    # @TODO: add support for IRSource
    from triton.compiler import IRSource
    from triton.runtime.jit import JITFunction

    if isinstance(source, IRSource):
        return

    # Get the function reference
    if isinstance(fn := source.fn, JITFunction):
        fn_ref = fn.fn
    else:
        fn_ref = source.fn

    python_source_file = inspect.getfile(fn_ref)

    # Get function range information
    if (
        isinstance(fn := source.fn, JITFunction)
        and hasattr(fn, "starting_line_number")
        and hasattr(fn, "raw_src")
    ):
        function_start_line = fn.starting_line_number
        source_lines = fn.raw_src
    else:
        source_lines, function_start_line = inspect.getsourcelines(fn_ref)

    function_end_line = function_start_line + len(source_lines) - 1

    if TRITON_FULL_PYTHON_SOURCE:
        # Full file mode: read the entire Python file
        try:
            # Check file size before reading
            file_size = os.path.getsize(python_source_file)
        except OSError as e:
            log.warning(
                f"Failed to check file size for {python_source_file}: {e}. "
                f"Falling back to function-only mode."
            )
            use_full_source = False
        else:
            if file_size > TRITON_MAX_SOURCE_SIZE:
                log.warning(
                    f"Source file {python_source_file} is too large ({file_size} bytes, "
                    f"limit: {TRITON_MAX_SOURCE_SIZE} bytes). Falling back to function-only mode."
                )
                use_full_source = False
            else:
                use_full_source = True

        if use_full_source:
            try:
                with open(python_source_file, "r", encoding="utf-8") as f:
                    file_content = f.read()

                # Calculate total lines
                total_lines = len(file_content.split("\n"))

                trace_data["python_source"] = {
                    "file_path": python_source_file,
                    "start_line": 1,
                    "end_line": total_lines,
                    "code": file_content,
                    # Add function range for frontend highlighting and scrolling
                    "function_start_line": function_start_line,
                    "function_end_line": function_end_line,
                }
                return
            except (OSError, UnicodeDecodeError) as e:
                log.warning(
                    f"Failed to read full source file {python_source_file}: {e}. "
                    f"Falling back to function-only mode."
                )

    # Default behavior: only extract function definition
    trace_data["python_source"] = {
        "file_path": python_source_file,
        "start_line": function_start_line,
        "end_line": function_end_line,
        "code": "".join(source_lines),
    }


def extract_file_content(trace_data: Dict[str, Any], metadata_group: Dict[str, str]):
    """
    Extract file content from metadata_group and add it to trace_data.

    Args:
        trace_data (Dict): Dictionary to store extracted information
        metadata_group (Dict): Dictionary mapping filenames to file paths
    """
    for ir_filename, file_path in metadata_group.items():
        # Add file path to trace data
        trace_data["file_path"][ir_filename] = file_path

        # Check if this is a text file we can read
        if any(ir_filename.endswith(ext) for ext in TEXT_FILE_EXTENSIONS):
            try:
                # Check file size before reading to avoid memory issues
                file_size = os.path.getsize(file_path)
                if file_size > MAX_FILE_SIZE:
                    message = f"<file too large: {file_size} bytes>"
                    trace_data["file_content"][ir_filename] = message
                    continue

                with open(file_path, "r") as f:
                    trace_data["file_content"][ir_filename] = f.read()
            except (UnicodeDecodeError, OSError) as e:
                # add more specific error type
                message = f"<error reading file: {str(e)}>"
                trace_data["file_content"][ir_filename] = message
                log.debug(f"Error reading file {file_path}: {e}")
    cubin_keys = [key for key in metadata_group.keys() if key.endswith(".cubin")]
    cubin_path = metadata_group[cubin_keys[0]] if cubin_keys else None

    if TRITONPARSE_DUMP_SASS and cubin_path:
        filename_no_ext = os.path.splitext(os.path.basename(cubin_path))[0]
        sass_filename = f"{filename_no_ext}.sass"
        try:
            import tritonparse.tools.disasm

            sass_content = tritonparse.tools.disasm.extract(cubin_path)
            trace_data["file_content"][sass_filename] = sass_content
        except subprocess.CalledProcessError as e:
            message = f"<nvdisasm failed: {str(e)}>"
            trace_data["file_content"][sass_filename] = message
        except OSError as e:
            message = f"<error reading cubin file: {str(e)}>"
            trace_data["file_content"][sass_filename] = message
        except Exception as e:
            message = f"<error dumping SASS: {str(e)}>"
            trace_data["file_content"][sass_filename] = message


def extract_metadata_from_src(trace_data, src):
    from triton._C.libtriton import get_cache_invalidating_env_vars

    env_vars = get_cache_invalidating_env_vars()
    # extra_options = src.parse_options()
    # options = backend.parse_options(dict(options or dict(), **extra_options))

    # trace_data["extra_options"] = extra_options
    trace_data["metadata"].update(
        {
            "env": env_vars,
            "src_attrs": src.attrs if hasattr(src, "attrs") else {},
            "src_cache_key": src.fn.cache_key if hasattr(src, "fn") else "",
            "src_constants": src.constants if hasattr(src, "constants") else {},
        }
    )


class TritonJsonFormatter(logging.Formatter):
    """
    Format log records as JSON for Triton compilation tracing.

    This formatter converts log records with metadata and payload into NDJSON format,
    suitable for structured logging and later analysis. It handles special attributes
    added by the tritonparse, such as metadata dictionaries and payload data.
    """

    def format(self, record: logging.LogRecord):
        log_entry = record.metadata
        payload = record.payload

        log_entry["timestamp"] = self.formatTime(record, "%Y-%m-%dT%H:%M:%S.%fZ")
        if payload is not None:
            log_entry["payload"] = json.loads(payload)
        clean_log_entry = convert(log_entry)
        # NDJSON format requires a newline at the end of each line
        json_str = json.dumps(clean_log_entry, separators=(",", ":"))
        return json_str + "\n"


class TritonTraceHandler(logging.StreamHandler):
    """
    A handler for Triton compilation tracing that outputs NDJSON files.

    This handler creates and manages log files for Triton kernel compilation traces.
    It supports creating new files for different compilation events and handles
    proper cleanup of file resources. When running in a distributed environment,
    it automatically adds rank information to filenames.

    The handler automatically detects when torch.distributed becomes initialized
    and switches to a rank-specific log file to avoid multiple ranks writing to
    the same file.
    """

    # Sentinel value to distinguish "uninitialized" from "no rank available (None)"
    _UNINITIALIZED = object()

    def __init__(
        self, root_dir: Optional[str] = None, prefix=DEFAULT_TRACE_FILE_PREFIX
    ):
        logging.Handler.__init__(self)
        self.root_dir = root_dir
        self.prefix = prefix
        self.stream = None
        self.first_record = True
        # Track the rank used when creating the current log file
        # _UNINITIALIZED means we haven't created a file yet
        self._last_rank = self._UNINITIALIZED
        # If the program is unexpected terminated, atexit can ensure  file resources are properly closed and released.
        # it is because we use `self.stream` to keep the opened file stream, if the program is interrupted by some errors, the stream may not be closed.
        atexit.register(self._cleanup)

    def _get_current_rank(self) -> Optional[int]:
        """
        Get the current distributed rank if available.

        Returns:
            The rank number if torch.distributed is initialized, None otherwise.
        """
        if TORCH_INSTALLED:
            import torch.distributed as dist

            if dist.is_available() and dist.is_initialized():
                return dist.get_rank()
        return None

    def get_root_dir(self):
        # For meta internal runs, we use the /logs directory by default
        # reference implementation
        # https://github.com/pytorch/pytorch/blob/5fe58ab5bd9e14cce3107150a9956a2ed40d2f79/torch/_logging/_internal.py#L1071
        if self.root_dir:
            return self.root_dir
        TRACE_LOG_DIR = "/logs"
        should_set_root_dir = True
        if TORCH_INSTALLED:
            import torch.version as torch_version

            if (
                hasattr(torch_version, "git_version")
                and os.getenv("MAST_HPC_JOB_NAME") is None
            ):
                log.info(
                    "TritonTraceHandler: disabled because not fbcode or conda on mast"
                )
                should_set_root_dir = False
            # TODO: change to tritonparse knob
            # The following check is necessary because the possible version mismatch between torch and tritonparse
            elif (
                hasattr(torch, "_utils_internal")
                and hasattr(torch._utils_internal, "justknobs_check")
                and not torch._utils_internal.justknobs_check("pytorch/trace:enable")
            ):
                log.info(
                    "TritonTraceHandler: disabled because justknobs_check('pytorch/trace:enable') returned False"
                )
                should_set_root_dir = False
        if should_set_root_dir:
            if not os.path.exists(TRACE_LOG_DIR):
                log.info(
                    "TritonTraceHandler: disabled because %s does not exist",
                    TRACE_LOG_DIR,
                )
            elif not os.access(TRACE_LOG_DIR, os.W_OK):
                log.info(
                    "TritonTraceHandler: disabled because %s is not writable",
                    TRACE_LOG_DIR,
                )
            else:
                self.root_dir = TRACE_LOG_DIR
        return self.root_dir

    def emit(self, record):
        # reference implementation
        # https://github.com/pytorch/pytorch/blob/5fe58ab5bd9e14cce3107150a9956a2ed40d2f79/torch/_logging/_internal.py#L1071
        try:
            # Check if rank status has changed since we created the current log file
            # This handles the case where torch.distributed is initialized after
            # some kernel compilations have already been logged
            current_rank = self._get_current_rank()
            if (
                self.stream is not None
                and self._last_rank is not self._UNINITIALIZED
                and current_rank != self._last_rank
            ):
                log.debug(
                    "TritonTraceHandler: rank changed from %s to %s, switching log file",
                    self._last_rank,
                    current_rank,
                )
                self._ensure_stream_closed()
                self.stream = None

            if self.stream is None:
                root_dir = self.get_root_dir()
                if root_dir is not None:
                    os.makedirs(root_dir, exist_ok=True)
                    ranksuffix = ""
                    if current_rank is not None:
                        ranksuffix = f"rank_{current_rank}_"
                    filename = f"{self.prefix}{ranksuffix}"
                    self._ensure_stream_closed()
                    # Choose file extension and mode based on compression setting
                    if TRITON_TRACE_COMPRESSION == "gzip":
                        file_extension = ".bin.ndjson"
                        file_mode = "ab+"  # Binary mode for gzip member concatenation
                    elif TRITON_TRACE_COMPRESSION == "clp":
                        file_extension = ".clp"
                        file_mode = "w"  # CLP write mode
                    else:
                        file_extension = ".ndjson"
                        file_mode = "a+"
                    log_file_name = os.path.abspath(
                        os.path.join(root_dir, f"{filename}{file_extension}")
                    )
                    if TRITON_TRACE_COMPRESSION == "clp":
                        self.stream = clp_open(log_file_name, file_mode)
                    else:
                        self.stream = open(
                            log_file_name,
                            mode=file_mode,
                        )
                    # Update _last_rank after successfully creating the file
                    self._last_rank = current_rank
                    log.debug("TritonTraceHandler: logging to %s", log_file_name)
                else:
                    triton_trace_log.removeHandler(self)
                    return

            if self.stream:
                formatted = self.format(record)
                if TRITON_TRACE_COMPRESSION == "gzip":
                    # Create a separate gzip member for each record
                    # This allows standard gzip readers to handle member concatenation automatically
                    buffer = io.BytesIO()
                    with gzip.GzipFile(fileobj=buffer, mode="wb") as gz:
                        gz.write(formatted.encode("utf-8"))
                    # Write the complete gzip member to the file
                    compressed_data = buffer.getvalue()
                    self.stream.write(compressed_data)
                elif TRITON_TRACE_COMPRESSION == "clp":
                    self.stream.add(io.StringIO(formatted))
                else:
                    self.stream.write(formatted)
                self.flush()
        except Exception as e:
            # record exception and ensure resources are released
            log.error(f"Error in TritonTraceHandler.emit: {e}")
            self._ensure_stream_closed()
            self.handleError(record)  # call Handler's standard error handling

    def close(self):
        """Close the current file."""
        self.acquire()
        try:
            try:
                if self.stream:
                    try:
                        self.flush()
                    finally:
                        if TRITON_TRACE_COMPRESSION == "clp":
                            self.stream._compress()
                        self.stream.close()
                        self.stream = None
            finally:
                # Solution adopted from PyTorch PR #120289
                logging.StreamHandler.close(self)
        finally:
            self.release()

    def _cleanup(self):
        """Ensure proper cleanup on program exit"""
        if self.stream is not None:
            self.close()

    def _ensure_stream_closed(self):
        """ensure stream is closed"""
        if self.stream is not None:
            try:
                self.flush()
            finally:
                self.stream.close()
                self.stream = None


def init_logs():
    """
    Initialise tritonparse's logging system.

    Requirements handled:
    1. First call may or may not pass `trace_folder`.
    2. A later call *can* pass `trace_folder` and must activate an
       existing handler whose `root_dir` was None.
    3. When tracing is disabled (no writable directory), prevent the
       empty                                     →
           DEBUG:tritonparse_trace:
       lines by blocking propagation to the root logger.
    """
    global TRITON_TRACE_HANDLER, triton_trace_folder, TENSOR_BLOB_MANAGER

    # Basic logger settings (safe to run on every call)
    triton_trace_log.setLevel(logging.DEBUG)
    triton_trace_log.propagate = False  # stops bubbling to root logger. see 3)
    # 1) Create the handler on first use (root_dir may be None)
    if TRITON_TRACE_HANDLER is None:
        TRITON_TRACE_HANDLER = TritonTraceHandler(triton_trace_folder)
    if triton_trace_folder is not None:
        TRITON_TRACE_HANDLER.root_dir = triton_trace_folder
    # 2) Re-evaluate whether we have a writable directory
    #    (`get_root_dir()` also checks /logs logic, permissions, etc.)
    root_dir = TRITON_TRACE_HANDLER.get_root_dir()
    if root_dir is None:
        # Tracing still disabled: ensure the handler is NOT attached
        if TRITON_TRACE_HANDLER in triton_trace_log.handlers:
            triton_trace_log.removeHandler(TRITON_TRACE_HANDLER)
        return  # quiet exit, no blank lines
    # 3) Tracing is enabled: attach the handler (if not already
    #    attached) and set the JSON formatter.
    if TRITON_TRACE_HANDLER not in triton_trace_log.handlers:
        TRITON_TRACE_HANDLER.setFormatter(TritonJsonFormatter())
        triton_trace_log.addHandler(TRITON_TRACE_HANDLER)

    # Initialize tensor blob manager if enabled
    if TRITONPARSE_SAVE_TENSOR_BLOBS and root_dir:
        if TENSOR_BLOB_MANAGER is None:
            TENSOR_BLOB_MANAGER = TensorBlobManager(
                root_dir=root_dir, storage_quota=TRITONPARSE_TENSOR_STORAGE_QUOTA
            )
        elif TENSOR_BLOB_MANAGER.root_dir is None:
            # Update root_dir if it wasn't set during initialization
            TENSOR_BLOB_MANAGER.set_root_dir(root_dir)


def trace_structured_triton(
    name: str,
    metadata_fn: Optional[Callable[[], Dict[str, Any]]] = None,
    *,
    payload_fn: Optional[Callable[[], Optional[Union[str, object]]]] = None,
):
    """
    Record structured trace information for Triton kernel compilation.

    This function is the main entry point for logging structured trace events
    in the Triton system. It handles initialization of the logging system if needed,
    creates new log files, and formats the trace data with metadata
    and payload information.

    Args:
        name (str): Name of the trace event (e.g., "compilation", "execution")
        metadata_fn (Callable): Function that returns a dictionary of metadata to include
                               in the trace record
        payload_fn (Callable): Function that returns the payload data (can be a string,
                              dictionary, or other serializable object)
    """

    if metadata_fn is None:

        def metadata_fn():
            return {}

    if payload_fn is None:

        def payload_fn():
            return None

    metadata_dict: Dict[str, Any] = {"event_type": name}
    metadata_dict["pid"] = os.getpid()
    custom_metadata = metadata_fn()
    if custom_metadata:
        metadata_dict.update(custom_metadata)

    metadata_dict["stack"] = get_stack_trace()

    # Log the record using our custom LogRecord
    payload = payload_fn()
    # Use a custom factory to create the record with simplified parameters
    record = create_triton_log_record(metadata=metadata_dict, payload=payload)
    # Log the custom record
    triton_trace_log.handle(record)


def maybe_trace_triton(
    src,
    metadata: Dict[str, Any],
    metadata_group: Dict[str, str],
    times: Any,
    event_type: str = "compilation",
    cache_hit: bool = False,
):
    """
    Collect and trace Triton kernel compilation information for debugging and profiling.

    This function gathers metadata, IR files, and source code information about a Triton
    kernel compilation, then logs it through the tracing system if tracing is enabled.
    It collects information from multiple sources:
    1. JSON metadata file (if provided)
    2. PyTorch compilation context (if available)
    3. IR and other compilation artifact files
    4. Python source code of the kernel function

    This function is designed to be used as a CompilationListener in triton.knobs.compilation.listener,
    which now accepts a list of listeners.

    Args:
        src (Union[ASTSource, IRSource]): Source object containing kernel information
        metadata (Dict[str, Any]): Dictionary containing metadata for the compilation
        metadata_group (Dict[str, Any]): Dictionary mapping filenames to file paths for all compilation artifacts
        times (CompileTimes): Object containing timing information for the compilation
        event_type (str): Type of event being traced (default: "compilation")
        cache_hit (bool): Whether the compilation was a cache hit (default: False)

    Returns:
        Dict[str, Any]: Dictionary containing all collected trace data, even if tracing is disabled
    """
    # Check kernel allowlist early to avoid unnecessary work
    if _KERNEL_ALLOWLIST_PATTERNS is not None:
        kernel_name = extract_kernel_name(src)
        if not should_trace_kernel(kernel_name, _KERNEL_ALLOWLIST_PATTERNS):
            # Return empty dict to indicate no tracing was done
            return {}

    # Initialize a dictionary with defaultdict to avoid key errors
    trace_data = defaultdict(dict)
    # Add cache_hit to metadata
    trace_data["metadata"]["cache_hit"] = cache_hit
    if not metadata:
        metadata_path = next(
            (Path(p) for c, p in metadata_group.items() if c.endswith(".json"))
        )
        with open(metadata_path, "r") as f:
            metadata = json.load(f)
            trace_data["metadata"].update(metadata)
    else:
        trace_data["metadata"].update(metadata)
    # Handle torch._guards which might not be recognized by type checker
    if TORCH_INSTALLED:
        trace_id = torch._guards.CompileContext.current_trace_id()  # type: ignore
    else:
        trace_id = None
    cid = trace_id.compile_id if trace_id else None
    if cid is not None:
        for attr_name in ["compiled_autograd_id", "frame_id", "frame_compile_id"]:
            attr_value = getattr(cid, attr_name, None)
            if attr_value is not None:
                trace_data["pt_info"][attr_name] = attr_value
    if trace_id:
        trace_data["pt_info"]["attempt"] = trace_id.attempt
    # Extract content from all IR and other files in the metadata group
    extract_file_content(trace_data, metadata_group)
    # Extract Python source code information if available
    extract_python_source_info(trace_data, src)
    extract_metadata_from_src(trace_data, src)

    # Add timing information if available
    if times:
        trace_data["metadata"]["times"] = times
    # Log the collected information through the tracing system
    trace_structured_triton(
        event_type,
        payload_fn=lambda: json.dumps(convert(trace_data)),
    )

    return trace_data


def extract_arg_info(arg_dict):
    """
    Extract detailed information from kernel arguments, especially for PyTorch
    tensors.

    Args:
        arg_dict: Dictionary of kernel arguments

    Returns:
        Dictionary with extracted argument information including tensor properties
    """
    extracted_args = {}

    for arg_name, arg_value in arg_dict.items():
        arg_info = {}

        # Check if it's a PyTorch tensor
        if TORCH_INSTALLED and isinstance(arg_value, torch.Tensor):
            arg_info["type"] = "tensor"
            arg_info.update(_log_torch_tensor_info(arg_value))
        # Handle custom Tensor/Storage types from triton_kernels
        elif _is_from_triton_kernels_module(arg_value):
            type_name = type(arg_value).__name__
            arg_info["type"] = f"triton_kernels.tensor.{type_name}"

            if type_name == "Tensor":
                # Dump all attributes needed to reconstruct the Tensor wrapper
                if hasattr(arg_value, "shape"):
                    arg_info["shape"] = convert(arg_value.shape)
                if hasattr(arg_value, "shape_max"):
                    arg_info["shape_max"] = convert(arg_value.shape_max)
                if hasattr(arg_value, "dtype"):
                    arg_info["dtype"] = convert(arg_value.dtype)
                if hasattr(arg_value, "storage"):
                    # Recursively process the storage, which can be another
                    # custom type or a torch.Tensor
                    storage_arg = {"storage": arg_value.storage}
                    arg_info["storage"] = extract_arg_info(storage_arg)["storage"]

            elif type_name == "Storage":
                # Dump all attributes needed to reconstruct the Storage object
                if (
                    hasattr(arg_value, "data")
                    and TORCH_INSTALLED
                    and isinstance(arg_value.data, torch.Tensor)
                ):
                    # The 'data' is a torch.Tensor, log its metadata fully
                    arg_info["data"] = _log_torch_tensor_info(arg_value.data)
                if hasattr(arg_value, "layout"):
                    arg_info["layout"] = convert(arg_value.layout)
            else:
                log.warning(f"Unknown type: {type(arg_value)}")

        # Handle scalar values
        elif isinstance(arg_value, (int, float, bool)):
            arg_info["type"] = type(arg_value).__name__
            arg_info["value"] = arg_value
        # Handle strings
        elif isinstance(arg_value, str):
            arg_info["type"] = "str"
            arg_info["value"] = arg_value
            arg_info["length"] = len(arg_value)
        # Handle other types
        else:
            arg_info["type"] = type(arg_value).__name__
            # Try to convert to string for logging
            arg_info["repr"] = str(arg_value)
            if len(arg_info["repr"]) > 200:  # Truncate very long representations
                arg_info["repr"] = arg_info["repr"][:200] + "..."

        extracted_args[arg_name] = arg_info

    return extracted_args


def add_launch_metadata(grid, metadata, arg_dict, inductor_args=None):
    # Check if we're in CUDA graph capture mode - if so, skip detailed argument extraction
    # to avoid CUDA errors (cudaErrorStreamCaptureUnsupported)
    is_capturing = False
    if TORCH_INSTALLED:
        try:
            is_capturing = torch.cuda.is_current_stream_capturing()
        except (AttributeError, RuntimeError):
            pass

    if is_capturing:
        # During CUDA graph capture, return minimal metadata without argument extraction
        return {
            "launch_metadata_tritonparse": (
                grid,
                metadata._asdict(),
                {"_note": "argument extraction skipped during CUDA graph capture"},
                {},
            )
        }

    # Extract detailed argument information (only when NOT capturing)
    extracted_args = extract_arg_info(arg_dict)
    extracted_inductor_args = extract_arg_info(inductor_args) if inductor_args else {}
    return {
        "launch_metadata_tritonparse": (
            grid,
            metadata._asdict(),
            extracted_args,
            extracted_inductor_args,
        )
    }


class JITHookImpl(JITHook):
    """
    JIT Hook implementation that overrides or sets the launch_metadata function for Triton kernels.

    This hook is essential for capturing detailed kernel launch information beyond the basic
    metadata (like kernel name) that Triton provides by default. Without setting a custom
    launch_metadata function, only minimal launch information is available as shown in:
    https://github.com/triton-lang/triton/blob/7ce287dc24b43476cdeb30529089ac361564505d/python/triton/compiler/compiler.py#L504

    By intercepting the JIT compilation process and setting a custom launch_metadata function,
    we can capture comprehensive runtime information including grid parameters, kernel metadata,
    and argument dictionaries for detailed analysis and logging.
    """

    def __call__(
        self,
        *,
        key: str,
        repr: str,
        fn,
        compile,
        is_manual_warmup: bool,
        already_compiled: bool,
        inductor_args: Optional[Dict[str, Any]] = None,
    ) -> Optional[bool]:
        """
        Override or set the launch_metadata function for the JIT-compiled kernel.

        This method is called during the JIT compilation process and allows us to
        inject our custom launch_metadata function that will be used to collect
        detailed kernel launch information.

        Args:
            key: Unique identifier for the kernel
            repr: String representation of the kernel
            fn: The JIT function object
            compile: Compilation function
            is_manual_warmup: Whether this is a manual warmup call
            already_compiled: Whether the kernel is already compiled

        Returns:
            True to continue with compilation, None/False to skip
        """
        # Check kernel allowlist early to avoid unnecessary work
        if _KERNEL_ALLOWLIST_PATTERNS is not None:
            kernel_name = fn.name
            if not should_trace_kernel(kernel_name, _KERNEL_ALLOWLIST_PATTERNS):
                # Skip overriding launch_metadata if kernel is not in allowlist
                return True

        # Get the current launch_metadata function if it exists
        function = getattr(fn, "jit_function", fn)

        current_launch_metadata = getattr(function, "launch_metadata", None)
        if current_launch_metadata is not None:
            log.warning(
                f"fn {fn} launch_metadata is not None: {current_launch_metadata}. It will be overridden by tritonparse."
            )
        function.launch_metadata = partial(
            add_launch_metadata, inductor_args=inductor_args
        )
        return True


class LaunchHookImpl(LaunchHook):
    """
    Launch Hook implementation for capturing and logging kernel launch metadata.

    This hook is responsible for intercepting kernel launches and extracting the detailed
    metadata that was set up by the JITHookImpl. It provides entry point for
    kernel execution, allowing comprehensive logging and analysis of kernel launches
    including timing, parameters, and execution context.

    The metadata captured includes:
    - Kernel name and function details
    - Grid dimensions and launch parameters
    - Kernel arguments and their values
    - Stream information
    - Custom metadata added by the launch_metadata function
    """

    def __call__(self, metadata):
        """
        Handle kernel launch entry point.

        This method is called when a kernel is about to be launched, providing
        access to all the launch metadata for logging, profiling, or analysis.
        metadata format:

                Args:
            metadata: LazyDict containing comprehensive launch information including
                     kernel name, function, stream, grid parameters, and custom data
                     format: {'name': 'add_kernel', 'function': None, 'stream': 0,
                              'launch_metadata_tritonparse': (grid, self.metadata, extracted_args)}
                     where extracted_args contains detailed info for each argument:
                     - For tensors: shape, dtype, device, stride, memory_usage, etc.
                     - For scalars: type and value
                     - For other types: type and string representation
                 defined here:
                 https://github.com/triton-lang/triton/blob/7ce287dc24b43476cdeb30529089ac361564505d/
                 python/triton/compiler/compiler.py#L512.
        """
        metadata_dict = metadata.get()
        # Check kernel allowlist early to avoid unnecessary work
        if _KERNEL_ALLOWLIST_PATTERNS is not None:
            kernel_name = metadata_dict.get("name")

            if not should_trace_kernel(kernel_name, _KERNEL_ALLOWLIST_PATTERNS):
                # Skip tracing if kernel is not in allowlist
                return

        trace_data = defaultdict(dict)
        trace_data["name"] = metadata_dict["name"]
        trace_data["function"] = metadata_dict["function"]
        trace_data["stream"] = metadata_dict["stream"]
        launch_metadata_tritonparse = metadata_dict.get(
            "launch_metadata_tritonparse", None
        )
        if launch_metadata_tritonparse is not None:
            trace_data["grid"] = launch_metadata_tritonparse[0]
            trace_data["compilation_metadata"] = launch_metadata_tritonparse[1]
            trace_data["extracted_args"] = launch_metadata_tritonparse[
                2
            ]  # Now contains detailed arg info
            trace_data["extracted_inductor_args"] = launch_metadata_tritonparse[3]
        trace_structured_triton("launch", metadata_fn=lambda: convert(trace_data))


def maybe_enable_trace_launch():
    global _trace_launch_enabled
    if TRITON_TRACE_LAUNCH and not _trace_launch_enabled:
        from triton import knobs

        launch_hook = LaunchHookImpl()
        jit_hook = JITHookImpl()
        knobs.runtime.jit_post_compile_hook = jit_hook
        knobs.runtime.launch_enter_hook = launch_hook

        _trace_launch_enabled = True


def init_basic(trace_folder: Optional[str] = None):
    """
    Initialize the basic logging system for Triton compilation.

    This function sets up the basic logging system for Triton kernel compilation.

    Args:
        trace_folder (Optional[str]): The folder to store the trace files.
    """
    global triton_trace_folder, _KERNEL_ALLOWLIST_PATTERNS
    maybe_enable_debug_logging()
    if triton_trace_folder is not None and trace_folder is not None:
        log.info(
            "Conflict settings: triton_trace_folder is already set to %s, we will use provided trace_folder(%s) instead.",
            triton_trace_folder,
            trace_folder,
        )
    if trace_folder is not None:
        triton_trace_folder = trace_folder

    # Parse and store kernel allowlist configuration
    _KERNEL_ALLOWLIST_PATTERNS = parse_kernel_allowlist()
    if _KERNEL_ALLOWLIST_PATTERNS:
        log.debug(
            f"Kernel allowlist enabled with patterns: {_KERNEL_ALLOWLIST_PATTERNS}"
        )
    else:
        log.debug("Kernel allowlist not set, tracing all kernels")

    init_logs()
    maybe_enable_trace_launch()


def init(
    trace_folder: Optional[str] = None,
    enable_trace_launch: bool = False,
    enable_more_tensor_information: bool = False,
    enable_sass_dump: Optional[bool] = False,
    enable_tensor_blob_storage: bool = False,
    tensor_storage_quota: Optional[int] = None,
    compression: Optional[str] = None,
):
    """
    This function is a wrapper around init_basic() that also sets up the compilation listener. Its arguments have higher priority than the environment variables for same settings.

    Args:
        trace_folder (Optional[str]): The folder to store the trace files.
        enable_trace_launch (bool): Whether to enable the trace launch hook.
        enable_more_tensor_information (bool): Whether to enable more tensor information logging.
            It only works when enable_trace_launch/TRITON_TRACE_LAUNCH is True.
        enable_sass_dump (Optional[bool]): Whether to enable SASS dumping.
        enable_tensor_blob_storage (bool): Whether to enable tensor blob storage.
        tensor_storage_quota (Optional[int]): Storage quota in bytes for tensor blobs (default: 100GB).
        compression (Optional[str]): Compression format for trace files ("none", "gzip", or "clp").
            If not specified, respects TRITON_TRACE_COMPRESSION env var, or defaults to "gzip".
    """
    global TRITON_TRACE_LAUNCH, TRITONPARSE_MORE_TENSOR_INFORMATION
    global TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK, TRITONPARSE_DUMP_SASS
    global TRITONPARSE_SAVE_TENSOR_BLOBS, TRITONPARSE_TENSOR_STORAGE_QUOTA
    global TRITON_TRACE_COMPRESSION

    # Set global flags BEFORE calling init_basic, so init_logs() can see them
    if enable_trace_launch:
        TRITON_TRACE_LAUNCH = True
        TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK = True
    if enable_more_tensor_information:
        TRITONPARSE_MORE_TENSOR_INFORMATION = True
    if enable_sass_dump:
        TRITONPARSE_DUMP_SASS = True
    if enable_tensor_blob_storage:
        TRITONPARSE_SAVE_TENSOR_BLOBS = True

    # Set the quota in global var for TensorBlobManager creation in init_logs()
    if tensor_storage_quota is not None:
        TRITONPARSE_TENSOR_STORAGE_QUOTA = tensor_storage_quota

    # Handle compression setting (follows enable_trace_launch pattern)
    # Priority: env var > Python API > module default ("gzip")
    if compression is not None:
        if os.getenv("TRITON_TRACE_COMPRESSION") is None:
            TRITON_TRACE_COMPRESSION = compression

    init_basic(trace_folder)
    from triton import knobs

    knobs.compilation.listener = maybe_trace_triton


def init_with_env():
    """
    This function is used to initialize TritonParse with the environment variable TRITON_TRACE_FOLDER and TRITON_TRACE_LAUNCH specifically.
    It is only supposed to be used in OSS triton's source code.
    """
    if triton_trace_folder:
        init(triton_trace_folder, enable_trace_launch=TRITON_TRACE_LAUNCH)


def clear_logging_config():
    """
    Clear all configurations made by init() and init_basic().

    This function resets the logging handlers, global state variables,
    and Triton knobs to their default states, effectively disabling
    the custom tracing.

    WARNING: This function is not supposed to be called unless you are sure
    you want to clear the logging config.
    """
    global TRITON_TRACE_HANDLER, triton_trace_folder, _KERNEL_ALLOWLIST_PATTERNS
    global _trace_launch_enabled
    global TENSOR_BLOB_MANAGER
    # 1. Clean up the log handler
    if TRITON_TRACE_HANDLER is not None:
        if TRITON_TRACE_HANDLER in triton_trace_log.handlers:
            triton_trace_log.removeHandler(TRITON_TRACE_HANDLER)
        TRITON_TRACE_HANDLER.close()
        TRITON_TRACE_HANDLER = None

    # 2. Reset global state variables
    triton_trace_folder = None
    _KERNEL_ALLOWLIST_PATTERNS = None
    _trace_launch_enabled = False

    # 3. Reset tensor blob manager and related flags
    TENSOR_BLOB_MANAGER = None

    # 4. Reset Triton knobs
    # Check if triton was actually imported and used
    from triton import knobs

    knobs.compilation.listener = None
    knobs.runtime.jit_post_compile_hook = None
    knobs.runtime.launch_enter_hook = None
