#!/usr/bin/env python3
"""Performance profiling for MCP tool calls.

Measures actual response times for every MCP tool call via Streamable HTTP transport.
Outputs a timing report to justify (or not) caching, batching, or other optimizations.

Run: python3 test_mcp_timing.py [--base-url http://host:port/mcp/] [--rounds N]
"""

import json
import sys
import time
import statistics
import requests
import argparse

# ─── Config ───────────────────────────────────────────────────────────────

DEFAULT_BASE_URL = "http://127.0.0.1:8080/mcp/"
HEADERS = {
    "Content-Type": "application/json",
    "Accept": "text/event-stream, application/json",
}

# ─── Helpers ──────────────────────────────────────────────────────────────

_id_counter = 0
_session_id = None
_timings = {}  # {tool_name: [duration_ms, ...]}


def _next_id():
    global _id_counter
    _id_counter += 1
    return _id_counter


def _send(method, params=None, base_url=DEFAULT_BASE_URL):
    """Send a JSON-RPC request and parse the SSE response."""
    global _session_id
    payload = {
        "jsonrpc": "2.0",
        "id": _next_id(),
        "method": method,
        "params": params or {},
    }
    hdrs = dict(HEADERS)
    if _session_id:
        hdrs["Mcp-Session-Id"] = _session_id

    resp = requests.post(base_url, json=payload, headers=hdrs, timeout=120, stream=True)

    if "Mcp-Session-Id" in resp.headers:
        _session_id = resp.headers["Mcp-Session-Id"]

    if resp.status_code != 200:
        return {"error": f"HTTP {resp.status_code}: {resp.text[:200]}"}

    result = None
    raw = resp.text
    events = raw.split("\n\n")
    for event in events:
        data_parts = []
        for line in event.split("\n"):
            if line.startswith("data: "):
                data_parts.append(line[6:])
            elif line.startswith("data:"):
                data_parts.append(line[5:])
        if data_parts:
            full_data = "\n".join(data_parts)
            try:
                parsed = json.loads(full_data)
                if "result" in parsed:
                    result = parsed["result"]
                elif "error" in parsed:
                    result = {"error": parsed["error"]}
            except json.JSONDecodeError:
                continue
    return result


def call_tool_timed(name, arguments=None, base_url=DEFAULT_BASE_URL):
    """Call an MCP tool and return (text, error_text, duration_ms)."""
    t0 = time.monotonic()
    result = _send("tools/call", {"name": name, "arguments": arguments or {}}, base_url)
    duration_ms = (time.monotonic() - t0) * 1000

    if result is None:
        return None, "No response", duration_ms
    if isinstance(result, dict) and "error" in result:
        return None, str(result["error"]), duration_ms

    content = result.get("content", [])
    is_error = result.get("isError", False)
    text = "\n".join(c.get("text", "") for c in content if c.get("type") == "text")
    if is_error:
        return text, text, duration_ms
    return text, None, duration_ms


def record(tool_name, duration_ms):
    """Record a timing measurement."""
    _timings.setdefault(tool_name, []).append(duration_ms)


# ─── Timing Tests ─────────────────────────────────────────────────────────

def init_session(base_url):
    """Initialize MCP session."""
    _send("initialize", {
        "protocolVersion": "2025-03-26",
        "capabilities": {},
        "clientInfo": {"name": "timing-suite", "version": "1.0"},
    }, base_url)
    _send("notifications/initialized", {}, base_url)
    print("Session initialized.\n")


def run_timing_tests(base_url, rounds=3):
    """Run each tool multiple times and collect timing data."""

    # Define test cases: (tool_name, arguments, description)
    test_cases = [
        # --- Lightweight lookups ---
        ("get_platform_stats", {}, "Platform stats (simple aggregation)"),
        ("check_health", {}, "Health check"),
        ("list_vendors", {}, "Vendor list (cached in DB?)"),
        ("list_cwes", {}, "CWE list"),

        # --- Searches (varying complexity) ---
        ("search_vulnerabilities", {"query": "log4j", "per_page": 5}, "Search: log4j (text search)"),
        ("search_vulnerabilities", {"severity": "critical", "has_exploits": True, "per_page": 5}, "Search: critical+exploits"),
        ("search_vulnerabilities", {"vendor": "fortinet", "is_kev": True, "per_page": 5}, "Search: fortinet+kev"),
        ("search_vulnerabilities", {"min_cvss": 9.0, "min_epss": 0.9, "per_page": 3}, "Search: high CVSS+EPSS"),
        ("search_vulnerabilities", {"vendor": "microsoft", "sort": "epss_desc", "per_page": 10}, "Search: MS sorted by EPSS"),
        ("search_exploits", {"source": "metasploit", "per_page": 5}, "Exploits: metasploit source"),
        ("search_exploits", {"attack_type": "RCE", "reliability": "reliable", "per_page": 5}, "Exploits: reliable RCE"),
        ("search_exploits", {"cve": "CVE-2024-3400", "per_page": 5}, "Exploits: by CVE"),
        ("search_exploits", {"min_stars": 100, "sort": "stars_desc", "per_page": 5}, "Exploits: popular (100+ stars)"),

        # --- Detail lookups (single resource) ---
        ("get_vulnerability", {"cve_id": "CVE-2024-3400"}, "Detail: CVE-2024-3400 (PAN-OS)"),
        ("get_vulnerability", {"cve_id": "CVE-2021-44228"}, "Detail: Log4Shell"),
        ("get_cwe", {"cwe_id": "79"}, "CWE detail: XSS"),
        ("get_author", {"author_name": "Chocapikk"}, "Author detail"),
        ("list_products", {"vendor": "microsoft"}, "Products for Microsoft"),
        ("lookup_alt_id", {"alt_id": "EDB-48537"}, "Lookup ExploitDB alt ID"),
        ("get_nuclei_templates", {"cve_id": "CVE-2024-3400"}, "Nuclei templates"),

        # --- Heavy operations ---
        ("audit_stack", {"technologies": "nginx, postgresql, node.js"}, "Audit stack (3 techs)"),
        ("generate_finding", {"cve_id": "CVE-2024-3400", "target": "fw.example.com"}, "Generate finding"),
    ]

    for round_num in range(1, rounds + 1):
        print(f"=== Round {round_num}/{rounds} ===")
        for tool_name, args, desc in test_cases:
            text, err, dur = call_tool_timed(tool_name, args, base_url)
            ok = "OK" if not err else "ERR"
            resp_size = len(text or "")
            print(f"  {dur:7.0f}ms  {ok}  {resp_size:>6}B  {desc}")
            record(f"{tool_name}:{desc}", dur)
            time.sleep(0.3)  # Small delay to avoid rate limiting
        print()
        if round_num < rounds:
            print("  (pausing 3s between rounds...)\n")
            time.sleep(3)

    # Also time a duplicate call immediately to measure consistency
    print("=== Duplicate Calls (same data, back-to-back) ===")
    dup_tests = [
        ("get_vulnerability", {"cve_id": "CVE-2024-3400"}, "Detail: CVE-2024-3400 (dup)"),
        ("search_vulnerabilities", {"query": "log4j", "per_page": 5}, "Search: log4j (dup)"),
        ("list_vendors", {}, "Vendor list (dup)"),
        ("get_platform_stats", {}, "Platform stats (dup)"),
    ]
    for tool_name, args, desc in dup_tests:
        # Call twice in quick succession
        _, _, dur1 = call_tool_timed(tool_name, args, base_url)
        _, _, dur2 = call_tool_timed(tool_name, args, base_url)
        print(f"  Call 1: {dur1:7.0f}ms   Call 2: {dur2:7.0f}ms   Δ: {dur2-dur1:+.0f}ms   {desc}")
        record(f"dup:{desc}:1", dur1)
        record(f"dup:{desc}:2", dur2)
        time.sleep(0.5)


def print_report():
    """Print a structured timing report."""
    print("\n" + "=" * 80)
    print("TIMING REPORT")
    print("=" * 80)
    print(f"{'Tool Call':<55} {'Min':>7} {'Avg':>7} {'Max':>7} {'P90':>7} {'N':>3}")
    print("-" * 80)

    # Group by category
    categories = {
        "LIGHTWEIGHT LOOKUPS": [],
        "SEARCH QUERIES": [],
        "DETAIL LOOKUPS": [],
        "HEAVY OPERATIONS": [],
        "DUPLICATE CALLS": [],
    }

    for key, durations in sorted(_timings.items()):
        if key.startswith("dup:"):
            categories["DUPLICATE CALLS"].append((key, durations))
        elif any(k in key for k in ["stats", "health", "list_vendors", "list_cwes"]):
            categories["LIGHTWEIGHT LOOKUPS"].append((key, durations))
        elif "Search" in key or "Exploit" in key:
            categories["SEARCH QUERIES"].append((key, durations))
        elif "Detail" in key or "CWE" in key or "Author" in key or "Products" in key or "Lookup" in key or "Nuclei" in key:
            categories["DETAIL LOOKUPS"].append((key, durations))
        else:
            categories["HEAVY OPERATIONS"].append((key, durations))

    all_durations = []
    for cat_name, items in categories.items():
        if not items:
            continue
        print(f"\n  {cat_name}:")
        for key, durations in items:
            # Clean up the key for display
            display_name = key.split(":", 1)[1] if ":" in key else key
            display_name = display_name[:52]
            mn = min(durations)
            avg = statistics.mean(durations)
            mx = max(durations)
            p90 = sorted(durations)[int(len(durations) * 0.9)] if len(durations) > 1 else mx
            print(f"    {display_name:<53} {mn:>6.0f} {avg:>6.0f} {mx:>6.0f} {p90:>6.0f} {len(durations):>3}")
            all_durations.extend(durations)

    print("-" * 80)
    if all_durations:
        total_calls = len(all_durations)
        overall_avg = statistics.mean(all_durations)
        overall_p50 = sorted(all_durations)[total_calls // 2]
        overall_p90 = sorted(all_durations)[int(total_calls * 0.9)]
        overall_p99 = sorted(all_durations)[int(total_calls * 0.99)]
        print(f"\n  Total calls: {total_calls}")
        print(f"  Overall avg: {overall_avg:.0f}ms")
        print(f"  P50: {overall_p50:.0f}ms  |  P90: {overall_p90:.0f}ms  |  P99: {overall_p99:.0f}ms")
        print(f"  Range: {min(all_durations):.0f}ms - {max(all_durations):.0f}ms")

        # Identify the slow tools (avg > 1s)
        slow = []
        for key, durations in _timings.items():
            if not key.startswith("dup:") and statistics.mean(durations) > 1000:
                slow.append((key, statistics.mean(durations)))
        if slow:
            print("\n  ⚠ SLOW TOOLS (avg > 1s):")
            for name, avg in sorted(slow, key=lambda x: -x[1]):
                display = name.split(":", 1)[1] if ":" in name else name
                print(f"    {display:<50} {avg:.0f}ms avg")

        # Check for caching opportunity: are duplicate calls the same speed?
        print("\n  CACHING ANALYSIS:")
        dup_items = [k for k in _timings if k.startswith("dup:")]
        pairs = {}
        for k in dup_items:
            base = k.replace(":1", "").replace(":2", "")
            pairs.setdefault(base, []).append((k, _timings[k][0]))
        for base, vals in pairs.items():
            if len(vals) == 2:
                d1 = vals[0][1]
                d2 = vals[1][1]
                saving = d1 - d2 if d2 < d1 else 0
                display = base.replace("dup:", "")
                verdict = f"No speedup (both ~{(d1+d2)/2:.0f}ms)" if abs(d1-d2) < 50 else f"Δ={d2-d1:+.0f}ms"
                print(f"    {display:<45} 1st: {d1:.0f}ms  2nd: {d2:.0f}ms  {verdict}")

    print("\n" + "=" * 80)


# ─── Main ─────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
    parser.add_argument("--rounds", type=int, default=3, help="Number of rounds per tool")
    args = parser.parse_args()

    print(f"EIP MCP Performance Profiler")
    print(f"Target: {args.base_url}")
    print(f"Rounds: {args.rounds}")
    print("=" * 60)

    init_session(args.base_url)
    start = time.time()
    run_timing_tests(args.base_url, args.rounds)
    elapsed = time.time() - start

    print(f"\nTotal profiling time: {elapsed:.1f}s")
    print_report()


if __name__ == "__main__":
    main()
