# -*- coding: utf-8 -*-
"""Location: ./tests/performance/test_plugins_performance.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Fred Araujo

Performance profiling for plugins using cProfile.

This script profiles each plugin's hooks individually using the plugin manager
and cProfile. It generates detailed performance profiles and a summary table
showing average execution times per hook type per plugin.

Usage:
    python tests/performance/test_plugins_performance.py [--details]

Options:
    --details    Print detailed profile for each plugin-hook combination

Output:
    - Individual .prof files in prof/ directory for each plugin-hook combination
    - Summary table printed to stdout showing average times per plugin per hook
    - Detailed profiles (if --details is specified)
"""

# Standard
import argparse
import asyncio
from collections import defaultdict
import cProfile
import importlib.util
import io
import logging
import os
import pstats
from pstats import SortKey
import sys
from typing import Any, Dict, Tuple

# Disable security warnings
logging.getLogger("mcpgateway.config").setLevel(logging.ERROR)
logging.getLogger("mcpgateway.observability").setLevel(logging.ERROR)

# Add repo root to PYTHONPATH
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.join(SCRIPT_DIR, "..", "..")
sys.path.insert(0, ROOT_DIR)

# First-Party
from mcpgateway.common.models import Message, PromptResult, ResourceContent, Role, TextContent  # noqa: E402
from cpex.framework import (  # noqa: E402
    GlobalContext,
    PluginManager,
    PromptHookType,
    PromptPosthookPayload,
    PromptPrehookPayload,
    ResourceHookType,
    ResourcePostFetchPayload,
    ResourcePreFetchPayload,
    ToolHookType,
    ToolPostInvokePayload,
    ToolPreInvokePayload,
)

# Configuration
CONFIG_PATH = os.path.join(SCRIPT_DIR, "plugins", "config.yaml")
PROFILE_OUTPUT_DIR = os.path.join(SCRIPT_DIR, "plugins", "prof")
ITERATIONS = 1000  # Number of iterations per hook

# Optional cpex-* plugin packages referenced by tests/performance/plugins/config.yaml.
# These ship in the [plugins] extra; if any are missing, the perf script skips
# cleanly rather than failing inside PluginManager initialization.
REQUIRED_CPEX_PACKAGES = (
    "cpex_pii_filter",
    "cpex_rate_limiter",
    "cpex_retry_with_backoff",
    "cpex_secrets_detection",
    "cpex_url_reputation",
)


def ensure_profile_dir() -> None:
    """Ensure the profile output directory exists."""
    os.makedirs(PROFILE_OUTPUT_DIR, exist_ok=True)


def create_sample_payloads() -> Dict[str, Any]:
    """Create sample payloads for each hook type.

    Returns:
        Dictionary mapping hook type names to sample payload objects
    """
    # Prompt pre-fetch payload
    prompt_pre = PromptPrehookPayload(
        prompt_id="test_prompt",
        args={
            "user": "test_user",
            "query": "What is the capital of France? My email is test@example.com",
            "context": "Additional context with SSN: 123-45-6789 and credit card: 4111-1111-1111-1111",
        },
    )

    # Prompt post-fetch payload
    message = Message(
        content=TextContent(type="text", text="The capital of France is Paris. Contact me at admin@test.com or call 555-123-4567."),
        role=Role.ASSISTANT,
    )
    prompt_result = PromptResult(messages=[message])
    prompt_post = PromptPosthookPayload(prompt_id="test_prompt", result=prompt_result)

    # Tool pre-invoke payload
    tool_pre = ToolPreInvokePayload(
        name="search_tool",
        args={
            "query": "search for crap information",
            "limit": 10,
            "filter": "type:document",
            "user_email": "user@example.com",
        },
    )

    # Tool post-invoke payload
    tool_post = ToolPostInvokePayload(
        name="search_tool",
        result={
            "status": "success",
            "results": [
                {"title": "Result 1", "snippet": "This is crap content with crud data"},
                {"title": "Result 2", "snippet": "Another crud result with sensitive info"},
            ],
            "count": 2,
            "metadata": {"query_time": 0.123, "source": "database"},
        },
    )

    # Resource pre-fetch payload
    resource_pre = ResourcePreFetchPayload(uri="https://example.com/document.txt", metadata={"version": "1.0", "encoding": "utf-8"})

    # Resource post-fetch payload
    resource_content = ResourceContent(
        type="resource",
        id="res-1",
        uri="https://example.com/document.txt",
        text="This document contains crap information and crud data. Contact admin@example.com for more details.",
    )
    resource_post = ResourcePostFetchPayload(uri="https://example.com/document.txt", content=resource_content)

    return {
        "prompt_pre_fetch": prompt_pre,
        "prompt_post_fetch": prompt_post,
        "tool_pre_invoke": tool_pre,
        "tool_post_invoke": tool_post,
        "resource_pre_fetch": resource_pre,
        "resource_post_fetch": resource_post,
    }


async def profile_plugin_hook(manager: PluginManager, plugin_name: str, hook_type: str, payload: Any, global_context: GlobalContext, iterations: int = ITERATIONS) -> Tuple[float, str]:
    """Profile a specific plugin's hook invocation.

    Args:
        manager: Plugin manager instance
        plugin_name: Name of the plugin to profile
        hook_type: Hook type (e.g., "prompt_pre_fetch")
        payload: Payload to pass to the hook
        global_context: Global context for hook invocation
        iterations: Number of iterations to run

    Returns:
        Tuple of (average_time_ms, profile_file_path)
    """
    # Map hook type strings to enum values
    hook_type_map = {
        "prompt_pre_fetch": PromptHookType.PROMPT_PRE_FETCH,
        "prompt_post_fetch": PromptHookType.PROMPT_POST_FETCH,
        "tool_pre_invoke": ToolHookType.TOOL_PRE_INVOKE,
        "tool_post_invoke": ToolHookType.TOOL_POST_INVOKE,
        "resource_pre_fetch": ResourceHookType.RESOURCE_PRE_FETCH,
        "resource_post_fetch": ResourceHookType.RESOURCE_POST_FETCH,
    }

    hook_enum = hook_type_map[hook_type]

    # Create profiler
    profiler = cProfile.Profile()

    # Profile the hook invocations
    profiler.enable()
    for _ in range(iterations):
        await manager.invoke_hook_for_plugin(plugin_name, hook_enum, payload, context=global_context)
    profiler.disable()

    # Save profile to file
    profile_filename = f"{plugin_name}_{hook_type}.prof"
    profile_path = os.path.join(PROFILE_OUTPUT_DIR, profile_filename)
    profiler.dump_stats(profile_path)

    # Calculate average time from stats
    stats = pstats.Stats(profiler)
    stats.strip_dirs()

    # Get cumulative time for invoke_hook_for_plugin calls
    # Note: ct (cumulative time) is in SECONDS
    # We need to find the invoke_hook_for_plugin function and get its average time
    avg_time_ms = 0.0
    for func_info, (cc, nc, tt, ct, callers) in stats.stats.items():
        # func_info is a tuple: (filename, line_number, function_name)
        filename, line_num, func_name = func_info
        if func_name == "invoke_hook_for_plugin":
            # ct is cumulative time in SECONDS for ALL calls
            # cc is the actual number of calls to this function
            # Average time per call = ct / cc (in seconds)
            # Convert to milliseconds by multiplying by 1000
            avg_time_ms = (ct / cc) * 1000 if cc > 0 else 0.0
            break

    return avg_time_ms, profile_path


async def profile_all_plugins(manager: PluginManager, show_details: bool = False) -> Dict[str, Dict[str, float]]:
    """Profile all enabled plugins and their hooks.

    Args:
        manager: Initialized plugin manager
        show_details: If True, print detailed profile for each plugin-hook combination

    Returns:
        Dictionary mapping plugin names to dictionaries of hook types to average times in ms
    """
    # Create sample payloads
    payloads = create_sample_payloads()

    # Create global context
    global_context = GlobalContext(request_id="perf-test-request", server_id="perf-test-server", user={"email": "admin@example.com", "is_admin": True})

    # Results storage: plugin_name -> hook_type -> avg_time_ms
    results: Dict[str, Dict[str, float]] = defaultdict(dict)

    # Get all plugins from manager
    plugins_info = []
    for plugin_config in manager.config.plugins:
        if plugin_config.mode != "disabled":
            plugin_ref = manager.get_plugin(plugin_config.name)
            if plugin_ref is None:
                print(f"  Warning: Enabled plugin '{plugin_config.name}' not found in registry, skipping")
                continue
            plugins_info.append((plugin_config.name, plugin_ref.hooks))

    print(f"\nProfiling {len(plugins_info)} enabled plugins...")
    print(f"Iterations per hook: {ITERATIONS}")
    print("=" * 80)

    # Profile each plugin's hooks
    for plugin_name, hooks in plugins_info:
        print(f"\nProfiling plugin: {plugin_name}")
        print(f"  Hooks: {', '.join(hooks)}")

        for hook_type in hooks:
            if hook_type not in payloads:
                print(f"  ⚠ Skipping {hook_type} - no sample payload defined")
                continue

            try:
                print(f"  • Profiling {hook_type}...", end=" ", flush=True)
                avg_time_ms, profile_path = await profile_plugin_hook(manager, plugin_name, hook_type, payloads[hook_type], global_context, ITERATIONS)
                results[plugin_name][hook_type] = avg_time_ms
                print(f"✓ {avg_time_ms:.3f}ms avg (saved to {os.path.relpath(profile_path, ROOT_DIR)})")

                # Print detailed profile if requested
                if show_details:
                    print_detailed_profile(profile_path, limit=20)
            except Exception as e:
                print(f"✗ Error: {e}")
                results[plugin_name][hook_type] = -1.0

    return dict(results)


def print_summary_table(results: Dict[str, Dict[str, float]]) -> None:
    """Print a formatted summary table of profiling results.

    Args:
        results: Dictionary mapping plugin names to hook type to average times
    """
    # Get all unique hook types
    all_hooks = set()
    for plugin_hooks in results.values():
        all_hooks.update(plugin_hooks.keys())
    hook_types = sorted(all_hooks)

    # Calculate column widths
    plugin_col_width = max(len(name) for name in results.keys()) + 2
    hook_col_width = 12  # Fixed width for hook columns

    # Print header
    print("\n" + "=" * 80)
    print("PERFORMANCE SUMMARY TABLE")
    print("=" * 80)
    print(f"\n{'Plugin':<{plugin_col_width}}", end="")
    for hook_type in hook_types:
        # Abbreviate hook names for table
        abbrev = hook_type.replace("prompt_", "P:").replace("tool_", "T:").replace("resource_", "R:").replace("_fetch", "").replace("_invoke", "")
        print(f"{abbrev:>{hook_col_width}}", end="")
    print()

    # Print separator
    print("-" * plugin_col_width, end="")
    for _ in hook_types:
        print("-" * hook_col_width, end="")
    print()

    # Print data rows
    for plugin_name in sorted(results.keys()):
        print(f"{plugin_name:<{plugin_col_width}}", end="")
        for hook_type in hook_types:
            avg_time = results[plugin_name].get(hook_type)
            if avg_time is None:
                print(f"{'—':>{hook_col_width}}", end="")  # Not implemented
            elif avg_time < 0:
                print(f"{'ERROR':>{hook_col_width}}", end="")  # Error
            else:
                print(f"{avg_time:>{hook_col_width - 2}.3f}ms", end="")
        print()

    # Print legend
    print("\n" + "=" * 80)
    print("LEGEND:")
    print("  P: = Prompt hooks   T: = Tool hooks   R: = Resource hooks")
    print("  pre/post = Hook timing   — = Not implemented   ERROR = Profiling failed")
    print(f"  All times are average per invocation over {ITERATIONS} iterations")
    print("=" * 80)


def print_detailed_profile(profile_path: str, limit: int = 20) -> None:
    """Print detailed statistics from a profile file.

    Args:
        profile_path: Path to the .prof file
        limit: Number of top functions to display
    """
    print(f"\n{'=' * 80}")
    print(f"Detailed Profile: {os.path.basename(profile_path)}")
    print("=" * 80)

    stats = pstats.Stats(profile_path)
    stats.strip_dirs()
    stats.sort_stats(SortKey.CUMULATIVE)

    # Create a string buffer to capture output
    s = io.StringIO()
    ps = pstats.Stats(profile_path, stream=s)
    ps.strip_dirs()
    ps.sort_stats(SortKey.CUMULATIVE)
    ps.print_stats(limit)

    print(s.getvalue())


async def main():
    """Main execution function."""
    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description="Profile plugin performance using cProfile",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run with summary table only
  python tests/performance/test_plugins_performance.py

  # Run with detailed profiles for each plugin-hook
  python tests/performance/test_plugins_performance.py --details
        """,
    )
    parser.add_argument("--details", action="store_true", help="Print detailed profile for each plugin-hook combination")
    args = parser.parse_args()

    print("=" * 80)
    print("MCP GATEWAY PLUGIN PERFORMANCE PROFILER")
    print("=" * 80)
    print(f"Config: {CONFIG_PATH}")
    print(f"Output: {PROFILE_OUTPUT_DIR}/")
    if args.details:
        print("Mode: Detailed profiles enabled")

    # Skip cleanly if optional cpex-* plugin packages are not installed.
    # The perf config references plugins from the [plugins] extra; without them
    # PluginManager.initialize() would raise inside the loader.
    missing = [pkg for pkg in REQUIRED_CPEX_PACKAGES if importlib.util.find_spec(pkg) is None]
    if missing:
        print(f"\n⏭️  Skipping plugin perf profiling — missing optional packages: {', '.join(missing)}")
        print("   Install with: pip install -e '.[plugins]'")
        return

    # Ensure output directory exists
    ensure_profile_dir()

    # Initialize plugin manager (do NOT profile this)
    print("\nInitializing plugin manager...")
    manager = PluginManager(CONFIG_PATH)
    await manager.initialize()

    if not manager.initialized:
        print("✗ Plugin manager failed to initialize")
        return

    print(f"✓ Plugin manager initialized with {manager.plugin_count} plugins")

    # Profile all plugins
    results = await profile_all_plugins(manager, show_details=args.details)

    # Print summary table
    print_summary_table(results)

    # Shutdown manager
    await manager.shutdown()
    print("\n✓ Performance profiling complete")
    print(f"✓ Profile files saved to: {os.path.relpath(PROFILE_OUTPUT_DIR, ROOT_DIR)}/")


if __name__ == "__main__":
    asyncio.run(main())
