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

import json
from abc import ABC
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Protocol

from tritonparse.reproducer.function_extractor import (
    extract_autotune_config_params,
    extract_function_with_decorators,
    extract_utility_functions,
    is_constexpr_param,
)
from tritonparse.reproducer.ingestion.ndjson import ContextBundle
from tritonparse.reproducer.templates.utils import (
    _disable_triton_autotune,
    get_function_source,
)
from tritonparse.reproducer.types import KernelImportMode
from tritonparse.reproducer.utils import (
    _generate_import_statements,
    _generate_invocation_snippet,
    _get_compile_params_for_invocation,
    _parse_kernel_signature,
)
from tritonparse.tp_logger import logger

# Size threshold for embedded JSON warning (50KB)
EMBEDDED_JSON_SIZE_THRESHOLD = 50 * 1024

# Imports already provided by the hardcoded import block in _replace_kernel_import.
# Used to deduplicate imports discovered by MultiFileCallGraphAnalyzer.
_SKIP_IMPORTS = {
    "import triton",
    "import torch",
    "import numpy",
    "import numpy as np",
}

# Bare module names that are submodule references (e.g. triton.language) and
# not valid standalone imports.  Discovered imports like "import language" are
# filtered out using this set.
_SKIP_BARE_MODULES = {
    "triton",
    "language",
    "tl",
    "tlx",
    "torch",
    "numpy",
    "np",
    "extra",
    "math",
    "typing",
    "functools",
    "types",
}

# Base import lines injected into every reproducer.  Shared between
# _replace_kernel_import (which emits them) and the skip-sets above
# (which prevent duplicates from the analyzer).
_BASE_IMPORT_LINES = [
    "import torch",
    "import numpy as np",
    "import triton",
    "import triton.language as tl",
    "from typing import List, Tuple",
]

# Known additional imports that may be needed by kernel source code.
# Maps module reference patterns found in kernel source to their import statement.
_EXTRA_IMPORT_PATTERNS = [
    # Triton Language Extensions (warp specialization, TMA, etc.)
    ("tlx.", "import triton.language.extra.tlx as tlx"),
    ("tlx ", "import triton.language.extra.tlx as tlx"),
]


def _detect_extra_imports(source_code: str) -> list[str]:
    """Scan kernel source code for usage of modules that need extra imports.

    Detects patterns like ``tlx.async_task(...)`` and returns the corresponding
    import statement (e.g. ``import triton.language.extra.tlx as tlx``).
    Deduplicates so each import appears at most once.

    Args:
        source_code: Combined kernel + dependent function source code.

    Returns:
        List of unique import statement strings.
    """
    seen: set[str] = set()
    imports: list[str] = []
    for pattern, import_stmt in _EXTRA_IMPORT_PATTERNS:
        if pattern in source_code and import_stmt not in seen:
            seen.add(import_stmt)
            imports.append(import_stmt)
    return imports


class HandlerProtocol(Protocol):
    def __call__(
        self, code: str, context_bundle: ContextBundle, **kwargs: Any
    ) -> str: ...


class PlaceholderReplacer(ABC):
    """
    Abstract base class for template placeholder replacement.

    Subclasses should register replacement handlers in their __init__ method
    by calling self.register(placeholder, handler_function).

    Each handler function should have the signature:
        handler(code: str, context_bundle: ContextBundle, **kwargs) -> str
    """

    def __init__(self):
        # Dictionary mapping placeholder strings to handler functions
        self.handlers: Dict[str, HandlerProtocol] = {}

    def register(self, placeholder: str, handler: HandlerProtocol):
        """
        Register a handler function for a specific placeholder.

        Args:
            placeholder: The placeholder string to replace (e.g., "{{JSON_FILE_NAME_PLACEHOLDER}}")
            handler: A callable that takes (code, context_bundle, **kwargs) and returns modified code
        """
        self.handlers[placeholder] = handler

    def replace(
        self, template_code: str, context_bundle: ContextBundle, **kwargs: Any
    ) -> str:
        """
        Replace all registered placeholders in the template code.

        Args:
            template_code: The template code containing placeholders
            context_bundle: Context information about the kernel
            **kwargs: Additional keyword arguments passed to handler functions

        Returns:
            The code with all placeholders replaced
        """
        code = template_code
        for handler in self.handlers.values():
            code = handler(code, context_bundle, **kwargs)
        return code


class DefaultPlaceholderReplacer(PlaceholderReplacer):
    """
    Default implementation of PlaceholderReplacer.

    Handles the following placeholders:
    - {{JSON_FILE_NAME_PLACEHOLDER}}: Replaced with the JSON file name
    - # {{KERNEL_SYSPATH_PLACEHOLDER}}: Replaced with sys.path setup code
    - # {{KERNEL_IMPORT_PLACEHOLDER}}: Replaced with kernel import statement
    - # {{KERNEL_INVOCATION_PLACEHOLDER}}: Replaced with kernel invocation code

    Args:
        preserve_autotune: If True, preserves @triton.autotune decorator behavior
            instead of disabling it. When enabled:
            1. Does NOT inject _disable_triton_autotune() function
            2. Re-extracts kernel WITH its @triton.autotune decorator from source
            3. Filters autotune-provided params from kernel invocation
            4. Places dependent functions before kernel (for config generators)
            This allows reproducers to re-autotune fresh, potentially finding
            better configurations for the given inputs. Default: False.
    """

    KERNEL_NAME_PLACEHOLDER = "{{KERNEL_NAME_PLACEHOLDER}}"
    JSON_FILE_NAME_PLACEHOLDER = "{{JSON_FILE_NAME_PLACEHOLDER}}"
    IR_OVERRIDE_SETUP_PLACEHOLDER = "# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}"
    KERNEL_SYSPATH_PLACEHOLDER = "# {{KERNEL_SYSPATH_PLACEHOLDER}}"
    KERNEL_IMPORT_PLACEHOLDER = "# {{KERNEL_IMPORT_PLACEHOLDER}}"
    UTILITY_FUNCTIONS_PLACEHOLDER = "# {{UTILITY_FUNCTIONS_PLACEHOLDER}}"
    KERNEL_INVOCATION_PLACEHOLDER = "# {{KERNEL_INVOCATION_PLACEHOLDER}}"
    # New placeholders for embed-context mode
    CONTEXT_JSON_PLACEHOLDER = "# {{CONTEXT_JSON_PLACEHOLDER}}"
    COMPILATION_JSON_PLACEHOLDER = "# {{COMPILATION_JSON_PLACEHOLDER}}"
    LAUNCH_KERNEL_BODY_PLACEHOLDER = "# {{LAUNCH_KERNEL_BODY_PLACEHOLDER}}"
    # Placeholder for reproducer metadata in docstring
    REPRODUCER_METADATA_PLACEHOLDER = "{{REPRODUCER_METADATA_PLACEHOLDER}}"

    def __init__(self, preserve_autotune: bool = False):
        super().__init__()
        self.preserve_autotune = preserve_autotune
        # Register all default handlers
        self.register(self.JSON_FILE_NAME_PLACEHOLDER, self._replace_json_filename)
        self.register(
            self.IR_OVERRIDE_SETUP_PLACEHOLDER, self._replace_ir_override_setup
        )
        self.register(self.KERNEL_SYSPATH_PLACEHOLDER, self._replace_kernel_syspath)
        self.register(self.KERNEL_IMPORT_PLACEHOLDER, self._replace_kernel_import)
        self.register(
            self.UTILITY_FUNCTIONS_PLACEHOLDER, self._replace_utility_functions
        )
        self.register(
            self.KERNEL_INVOCATION_PLACEHOLDER, self._replace_kernel_invocation
        )
        self.register(self.KERNEL_NAME_PLACEHOLDER, self._replace_kernel_name)
        # Register new handlers for embed-context mode
        self.register(self.CONTEXT_JSON_PLACEHOLDER, self._replace_context_json)
        self.register(self.COMPILATION_JSON_PLACEHOLDER, self._replace_compilation_json)
        self.register(
            self.LAUNCH_KERNEL_BODY_PLACEHOLDER, self._replace_launch_kernel_body
        )
        # Register handler for reproducer metadata in docstring
        self.register(
            self.REPRODUCER_METADATA_PLACEHOLDER, self._replace_reproducer_metadata
        )

    def _replace_kernel_name(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the kernel name placeholder."""
        kernel_name = context_bundle.kernel_info.function_name
        if not kernel_name:
            raise ValueError("Kernel function name is not available")
        return code.replace(self.KERNEL_NAME_PLACEHOLDER, kernel_name)

    def _replace_json_filename(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the JSON file name placeholder."""
        temp_json_path = kwargs.get("temp_json_path")
        if temp_json_path is None:
            raise ValueError("temp_json_path is required for JSON filename replacement")
        return code.replace(self.JSON_FILE_NAME_PLACEHOLDER, temp_json_path.name)

    def _replace_ir_override_setup(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the IR override setup placeholder."""
        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)

        if kernel_import != KernelImportMode.OVERRIDE_TTIR:
            return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, "")

        comp_json_filename = kwargs.get("comp_json_filename")
        if not comp_json_filename:
            raise ValueError("comp_json_filename is required for OVERRIDE_TTIR mode")

        setup_code = f'''
def create_ttir_tempfile():
    """Extract TTIR from compilation event and create temporary file."""
    script_dir = Path(__file__).resolve().parent
    comp_json_file = script_dir / "{comp_json_filename}"

    with open(comp_json_file, 'r') as f:
        comp_data = json.load(f)

    # Extract TTIR content
    kernel_name = comp_data['payload']['metadata']['name']
    ttir_key = f"{{kernel_name}}.ttir"
    ttir_content = comp_data['payload']['file_content'][ttir_key]

    # Create temporary file
    temp_file = tempfile.NamedTemporaryFile(
        mode='w',
        suffix='.ttir',
        delete=False,
        prefix=f'{{kernel_name}}_'
    )
    temp_file.write(ttir_content)
    temp_file.close()
    return temp_file.name


# Monkeypatch triton.autotune to use our TTIR
_ttir_file = create_ttir_tempfile()
_original_autotune = None

def _patched_autotune(configs, key=None, **kwargs):
    """Patched autotune that uses our TTIR file."""
    import triton
    # Replace configs with our single config using ir_override
    new_configs = [triton.Config(kwargs={{}}, ir_override=_ttir_file)]
    # Call original autotune with our config
    return _original_autotune(new_configs, key=[], **kwargs)

# Apply the monkeypatch before importing the kernel
import triton
_original_autotune = triton.autotune
triton.autotune = _patched_autotune
'''

        return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, setup_code)

    def _replace_kernel_syspath(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the kernel sys.path placeholder."""
        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)

        if kernel_import == KernelImportMode.DEFAULT:
            sys_stmt, _ = _generate_import_statements(context_bundle.kernel_info)
            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, sys_stmt)
        elif kernel_import == KernelImportMode.COPY:
            comment = (
                "# Kernel sys.path setup skipped - kernel source code embedded below"
            )
            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
        elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
            comment = "# Kernel sys.path setup skipped - using IR override mode"
            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
        else:
            raise ValueError(f"Unknown kernel_import mode: {kernel_import}")

    def _replace_kernel_import(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the kernel import placeholder.

        When preserve_autotune is enabled for COPY mode:
        - Does NOT inject _disable_triton_autotune() function
        - Re-extracts kernel WITH its @triton.autotune decorator from source
        - Adds a note that autotuning is enabled
        """
        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)

        if kernel_import == KernelImportMode.DEFAULT:
            _, import_statement = _generate_import_statements(
                context_bundle.kernel_info
            )

            # For DEFAULT mode, always disable autotune (import from original file)
            final_stmt = "\n".join(
                [import_statement, ""] + get_function_source(_disable_triton_autotune)
            )
            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, final_stmt)
        elif kernel_import == KernelImportMode.COPY:
            source_code = context_bundle.kernel_info.source_code
            func_name = context_bundle.kernel_info.function_name

            if not source_code or not source_code.strip():
                raise ValueError("Kernel source code is empty, cannot use 'copy' mode")
            if not func_name:
                raise ValueError(
                    "Cannot determine kernel function name for 'copy' mode"
                )

            # When preserve_autotune is enabled, re-extract kernel with decorators
            if self.preserve_autotune:
                source_with_decorators = extract_function_with_decorators(
                    context_bundle.kernel_info.file_path,
                    func_name,
                    context_bundle.source_repo_dir,
                )
                if source_with_decorators:
                    source_code = source_with_decorators
                    logger.debug("Using kernel source with decorators for autotuning")

            else:
                # Standard behavior: add dependent functions after kernel
                dep_result = get_dependent_source_map(
                    context_bundle.kernel_info.function_name,
                    context_bundle.kernel_info.file_path,
                    context_bundle.source_repo_dir,
                )
                # Only add dependent functions if extraction was successful
                if dep_result:
                    # Add imports required by dependent functions
                    if dep_result.import_statements:
                        dep_imports_code = "\n".join(dep_result.import_statements)
                        source_code = dep_imports_code + "\n\n" + source_code
                    # Add separator and dependent functions
                    dependent_code = (
                        "\n\n# Dependent functions extracted from source file\n\n"
                    )
                    dependent_code += "\n\n".join(dep_result.functions.values())
                    source_code += "\n\n" + dependent_code
                    logger.debug("Appended dependent functions to kernel source code")

            # Add common imports needed for most Triton kernels
            import_lines = list(_BASE_IMPORT_LINES) + [""]

            # Scan kernel and dependent source for additional imports
            # (e.g., "import triton.language.extra.tlx as tlx")
            all_source_for_imports = source_code
            extra_imports = _detect_extra_imports(all_source_for_imports)
            if extra_imports:
                import_lines.extend(extra_imports)
                import_lines.append("")

            # Only add autotune disabling code when not preserving autotune
            if self.preserve_autotune:
                import_lines.append(
                    "# NOTE: Autotuning is ENABLED - kernel will autotune on first run"
                )
                import_lines.append("")

                # Inject dependent functions after imports so they can use @triton.jit
                dep_result = get_dependent_source_map(
                    func_name,
                    context_bundle.kernel_info.file_path,
                    context_bundle.source_repo_dir,
                )
                if dep_result:
                    # Add imports required by dependent functions
                    if dep_result.import_statements:
                        import_lines.extend(dep_result.import_statements)
                        import_lines.append("")
                    import_lines.append("# Dependent functions for autotuning")
                    import_lines.append("")
                    import_lines.extend(dep_result.functions.values())
                    import_lines.append("")
            else:
                import_lines.extend(get_function_source(_disable_triton_autotune))

            # Combine: imports + kernel source code + alias
            embedded_code = "\n".join(import_lines)
            embedded_code += "\n" + source_code
            embedded_code += f"\n\n# Use kernel function directly\nimported_kernel_function = {func_name}"

            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, embedded_code)
        elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
            comment = "# Kernel import skipped - using IR override mode with TTIR"
            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, comment)
        else:
            raise ValueError(f"Unknown kernel_import mode: {kernel_import}")

    def _replace_utility_functions(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the utility functions placeholder with extracted functions."""
        embed_context = kwargs.get("embed_context", False)
        utility_code = extract_utility_functions(embed_context=embed_context)
        code = code.replace(self.UTILITY_FUNCTIONS_PLACEHOLDER, utility_code)

        return code

    def _replace_kernel_invocation(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the kernel invocation placeholder.

        When preserve_autotune is enabled:
        - Does NOT pass compile params (num_warps, num_stages, etc.) - autotuner handles these
        - Does NOT pass params that autotune configs provide (BLOCK_SIZE_*, etc.)
        - Passes remaining constexpr params as keyword args to avoid positional conflicts

        Otherwise, passes all args plus compile params as usual.
        """
        source_code = context_bundle.kernel_info.source_code
        pos_args, kw_args = _parse_kernel_signature(source_code)

        if self.preserve_autotune:
            # Get full kernel source with decorators to find autotune config params
            func_name = context_bundle.kernel_info.function_name
            full_source_code = extract_function_with_decorators(
                context_bundle.kernel_info.file_path,
                func_name,
                context_bundle.source_repo_dir,
            )
            # Fall back to captured source if extraction fails
            if not full_source_code:
                full_source_code = source_code

            # Also get dependent functions - they may contain triton.Config definitions
            dep_result_for_invocation = get_dependent_source_map(
                func_name,
                context_bundle.kernel_info.file_path,
                context_bundle.source_repo_dir,
            )
            all_source_code = full_source_code
            if dep_result_for_invocation:
                all_source_code += "\n\n" + "\n\n".join(
                    dep_result_for_invocation.functions.values()
                )

            # Get params that autotune configs provide
            autotune_config_params = extract_autotune_config_params(all_source_code)

            if not autotune_config_params:
                # No autotune configs found (decorator extraction may have failed).
                # Fall back to injecting compile params so the kernel launches with
                # the correct num_warps / num_stages instead of Triton defaults,
                # which can deadlock warp-specialized kernels.
                logger.debug(
                    "No autotune config params found; injecting compile params"
                )
                compile_params = _get_compile_params_for_invocation(
                    context_bundle.compile, kw_args
                )
                invocation_snippet = _generate_invocation_snippet(
                    pos_args, kw_args, compile_params
                )
            else:
                # Filter out params provided by autotune configs
                # Remaining constexpr params must be passed as keyword args
                filtered_pos_args = []
                extra_kw_args = []
                for arg in pos_args:
                    if arg in autotune_config_params:
                        continue
                    if autotune_config_params and is_constexpr_param(arg, source_code):
                        extra_kw_args.append(arg)
                    else:
                        filtered_pos_args.append(arg)

                filtered_kw_args = [
                    arg for arg in kw_args if arg not in autotune_config_params
                ]
                all_filtered_kw_args = extra_kw_args + filtered_kw_args

                # Get compile params that are NOT provided by autotune configs
                # (e.g., num_warps=8 when configs don't specify num_warps)
                from tritonparse.reproducer.utils import TRITON_COMPILE_PARAMS

                autotune_compile_params = set(autotune_config_params) & set(
                    TRITON_COMPILE_PARAMS
                )
                # Add autotune-handled compile params to the filter list so
                # _get_compile_params_for_invocation excludes them from the result
                extended_kw_args = list(kw_args) + list(autotune_compile_params)
                remaining_compile_params = _get_compile_params_for_invocation(
                    context_bundle.compile, extended_kw_args
                )

                # Generate invocation with remaining compile_params not handled by autotuner
                invocation_snippet = self._generate_autotune_invocation_snippet(
                    filtered_pos_args, all_filtered_kw_args, remaining_compile_params
                )
        else:
            # Standard behavior: pass compile params
            compile_params = _get_compile_params_for_invocation(
                context_bundle.compile, kw_args
            )
            invocation_snippet = _generate_invocation_snippet(
                pos_args, kw_args, compile_params
            )

        return code.replace(self.KERNEL_INVOCATION_PLACEHOLDER, invocation_snippet)

    def _generate_autotune_invocation_snippet(
        self,
        positional_args: List[str],
        keyword_args: List[str],
        compile_params: Optional[List[str]] = None,
    ) -> str:
        """Generate kernel invocation snippet for autotuned kernels.

        Unlike the default generator, this only includes compile params that are
        NOT handled by the autotuner, and passes non-autotune constexpr params
        as keyword arguments to avoid positional conflicts.
        """
        pos_args_str = ", ".join([f'args_dict["{arg}"]' for arg in positional_args])
        kw_args_str = ", ".join([f'{arg}=args_dict["{arg}"]' for arg in keyword_args])

        all_args = []
        if pos_args_str:
            all_args.append(pos_args_str)
        if kw_args_str:
            all_args.append(kw_args_str)

        # Add compile params that are not handled by autotuner
        if compile_params:
            all_args.append(", ".join(compile_params))

        return f"imported_kernel_function[tuple(grid)]({', '.join(all_args)})"

    def _replace_context_json(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the context JSON placeholder with embedded JSON or empty string."""
        embed_context = kwargs.get("embed_context", False)

        if not embed_context:
            return code.replace(self.CONTEXT_JSON_PLACEHOLDER, "")

        # Serialize launch event to JSON
        json_str = json.dumps(context_bundle.raw_launch_event, indent=2)

        # Warn if blob_path detected (external tensor dependencies)
        self._warn_if_blob_path_present(context_bundle.raw_launch_event)

        # Warn if size exceeds threshold
        if len(json_str) > EMBEDDED_JSON_SIZE_THRESHOLD:
            logger.warning(
                f"Embedded JSON is large ({len(json_str) // 1024}KB). "
                "Consider using file mode for better readability."
            )

        # Use raw triple-quoted string to minimize escaping issues
        embedded = f'CONTEXT_JSON = r"""\n{json_str}\n"""'
        return code.replace(self.CONTEXT_JSON_PLACEHOLDER, embedded)

    def _replace_compilation_json(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace compilation JSON placeholder (only for OVERRIDE_TTIR mode)."""
        embed_context = kwargs.get("embed_context", False)
        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)

        # Only embed compilation JSON for OVERRIDE_TTIR mode with embedding enabled
        if not embed_context or kernel_import != KernelImportMode.OVERRIDE_TTIR:
            return code.replace(self.COMPILATION_JSON_PLACEHOLDER, "")

        json_str = json.dumps(context_bundle.raw_comp_event, indent=2)
        embedded = f'COMPILATION_JSON = r"""\n{json_str}\n"""'
        return code.replace(self.COMPILATION_JSON_PLACEHOLDER, embedded)

    def _replace_launch_kernel_body(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace launch kernel body with file-based or embedded loading logic."""
        embed_context = kwargs.get("embed_context", False)
        temp_json_path = kwargs.get("temp_json_path")

        body_lines: list[str] = []

        # Inject triton.set_allocator when the kernel requires scratch memory.
        # Kernels compiled with global_scratch_size > 0 need a custom allocator;
        # without one the kernel launch will fail or hang.
        global_scratch_size = context_bundle.compile.get("global_scratch_size")
        if global_scratch_size and int(global_scratch_size) > 0:
            body_lines.extend(
                [
                    f"# This kernel requires scratch memory (global_scratch_size={global_scratch_size}).",
                    "# A default allocator is provided below. If the kernel hangs or produces",
                    "# incorrect results, replace this with the allocator from the original application.",
                    "def _alloc_fn(size: int, align: int, stream):",
                    "    return torch.empty(size, dtype=torch.int8, device='cuda')",
                    "triton.set_allocator(_alloc_fn)",
                    "",
                ]
            )

        if embed_context:
            # Embedded mode: parse JSON string directly
            body_lines.extend(
                [
                    "data = json.loads(CONTEXT_JSON)",
                    "grid, args_dict = create_args_from_json(data)",
                ]
            )
        else:
            # File mode: load from external JSON file
            json_filename = (
                temp_json_path.name if temp_json_path else "repro_context.json"
            )
            body_lines.extend(
                [
                    "script_dir = Path(__file__).resolve().parent",
                    f'json_file = script_dir / "{json_filename}"',
                    "grid, args_dict = create_args_from_json_file(str(json_file))",
                ]
            )

        body = "\n    ".join(body_lines)
        return code.replace(self.LAUNCH_KERNEL_BODY_PLACEHOLDER, body)

    def _warn_if_blob_path_present(self, raw_launch_event: dict) -> None:
        """Warn if any tensor argument has blob_path (external dependency)."""
        extracted_args = raw_launch_event.get("extracted_args", {})
        blob_args = [
            name
            for name, info in extracted_args.items()
            if isinstance(info, dict) and info.get("blob_path")
        ]
        if blob_args:
            logger.warning(
                f"The following tensor arguments have external blob_path dependencies: "
                f"{blob_args}. The embedded reproducer will NOT be fully standalone."
            )

    def _replace_reproducer_metadata(
        self, code: str, context_bundle: ContextBundle, **kwargs
    ) -> str:
        """Replace the reproducer metadata placeholder with generation info."""
        from datetime import datetime

        input_path = kwargs.get("input_path", "unknown")
        line_index = kwargs.get("line_index", "unknown")
        template = kwargs.get("template", "example")
        kernel_name = context_bundle.kernel_info.function_name
        generated_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        metadata = f"""Kernel: {kernel_name}
Source NDJSON: {input_path}
Line Index: {line_index}  (use --line {line_index} to regenerate)
Template: {template}
Generated: {generated_time}

To regenerate this reproducer:
    python -m tritonparse reproduce <ndjson_file> --line {line_index} --out-dir <output_dir>"""

        return code.replace(self.REPRODUCER_METADATA_PLACEHOLDER, metadata)


@dataclass
class DependentSourceResult:
    """Result of dependent function extraction, including functions and their imports."""

    functions: Dict[str, str]  # qualified_name -> source_code
    import_statements: List[str]  # formatted import statements needed by dependencies


def get_dependent_source_map(
    function_name: str,
    file_path: str,
    source_repo_dir: Optional[str] = None,
) -> Optional[DependentSourceResult]:
    """
    Extract dependent functions and their required imports.

    Returns:
        DependentSourceResult with functions and their import statements,
        or None if extraction fails.
    """
    from pathlib import Path

    source_path = Path(file_path)
    if not source_path.exists() and source_repo_dir:
        source_path = _map_file_path_to_source_repo(file_path, source_repo_dir)
    if not source_path or not source_path.exists():
        return None

    try:
        # Use MultiFileCallGraphAnalyzer for multi-file analysis
        from tritonparse.reproducer.multi_file_analyzer import (
            MultiFileCallGraphAnalyzer,
        )

        analyzer = MultiFileCallGraphAnalyzer(
            entry_file=str(source_path),
            entry_function=function_name,
        )
        result = analyzer.analyze()

        logger.info(
            f"Extracted {result.stats.total_functions_found} dependent functions "
            f"from {result.stats.total_files_analyzed} files with "
            f"{result.stats.total_imports} imports"
        )

        # Print dependent functions' short names
        logger.info("\nDependent functions (short names):")
        for func_name in sorted(result.function_short_names.keys()):
            short_name = result.function_short_names[func_name]
            logger.info(
                "  - %s. %s", short_name, result.functions[func_name].splitlines()[0]
            )

        # Filter out imports already covered by _BASE_IMPORT_LINES and
        # bare submodule names that aren't valid standalone imports.
        import_stmts: List[str] = []
        for imp in result.imports:
            if imp.import_type == "from_import":
                names_str = ", ".join(imp.names)
                stmt = f"from {imp.module} import {names_str}"
            else:
                if imp.names:
                    stmt = f"import {imp.names[0]}"
                else:
                    stmt = f"import {imp.module}"
            # Skip duplicate or already-covered imports
            if stmt in import_stmts or stmt in _SKIP_IMPORTS:
                continue
            # Skip bare module names that are submodule references
            if imp.import_type == "import":
                module = imp.names[0] if imp.names else imp.module
                if module in _SKIP_BARE_MODULES:
                    continue
            import_stmts.append(stmt)

        return DependentSourceResult(
            functions=result.functions,
            import_statements=import_stmts,
        )

    except Exception as e:
        # If AST analysis fails, continue without dependent functions
        logger.warning(f"Failed to extract dependent functions: {e}", exc_info=True)
        return None


def _map_file_path_to_source_repo(
    file_path: str, source_repo_dir: str
) -> Optional[Path]:
    """
    Try to map the given file path to a path in fbsource.

    Returns:
        The mapped file path in fbsource, or None if mapping fails.
    """

    source_repo_path = Path(source_repo_dir)
    if not source_repo_path.exists():
        logger.warning(f"Source repo dir {source_repo_dir} does not exist")
        return None

    prod_file_path = Path(file_path)
    for i in range(1, len(prod_file_path.parts)):
        new_path = source_repo_path / Path("/".join(prod_file_path.parts[i:]))
        if new_path.exists():
            logger.info(f"Map file_path: {file_path} to {new_path}")
            return new_path
    return None
