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

import json
from abc import ABC
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


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
                dependent_source_map = 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 dependent_source_map:
                    # Add separator, import statements, and dependent functions
                    dependent_code = (
                        "\n\n# Dependent functions extracted from source file\n\n"
                    )
                    dependent_code += "\n\n".join(dependent_source_map.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 = [
                "import torch",
                "import numpy as np",
                "import triton",
                "import triton.language as tl",
                "from typing import List, Tuple",
                "",
            ]

            # 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("")
            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.

        When preserve_autotune is enabled for COPY mode, also injects dependent
        functions (like config generators) before the kernel source so they're
        defined when referenced by @triton.autotune decorator.
        """
        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)

        # For preserve_autotune + COPY mode, inject dependent functions before kernel
        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
        if self.preserve_autotune and kernel_import == KernelImportMode.COPY:
            func_name = context_bundle.kernel_info.function_name
            if func_name:
                dependent_source_map = get_dependent_source_map(
                    func_name,
                    context_bundle.kernel_info.file_path,
                    context_bundle.source_repo_dir,
                )
                if dependent_source_map:
                    # Insert dependent functions before "# isort: off"
                    dependent_code = (
                        "\n# Dependent functions for autotuning\n\n"
                        + "\n\n".join(dependent_source_map.values())
                        + "\n\n"
                    )
                    code = code.replace("# isort: off", dependent_code + "# isort: off")

        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
            dependent_source_map = get_dependent_source_map(
                func_name,
                context_bundle.kernel_info.file_path,
                context_bundle.source_repo_dir,
            )
            all_source_code = full_source_code
            if dependent_source_map:
                all_source_code += "\n\n" + "\n\n".join(dependent_source_map.values())

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

            # 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

            # Generate invocation WITHOUT compile_params
            invocation_snippet = self._generate_autotune_invocation_snippet(
                filtered_pos_args, all_filtered_kw_args
            )
        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],
    ) -> str:
        """Generate kernel invocation snippet for autotuned kernels.

        Unlike the default generator, this does NOT include compile params
        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)

        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")

        if embed_context:
            # Embedded mode: parse JSON string directly
            body = """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 = f"""script_dir = Path(__file__).resolve().parent
    json_file = script_dir / "{json_filename}"
    grid, args_dict = create_args_from_json_file(str(json_file))"""

        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)


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

    Returns:
        A tuple of (functions_dict, import_statements_list) or None if extraction fails.
        - functions_dict: Maps qualified function names to their source code
        - import_statements_list: List of formatted import statements needed by dependent functions
    """
    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]
            )

        return result.functions

    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
