"""MCP server for the Exploit Intelligence Platform.

Exposes vulnerability and exploit intelligence to AI assistants via
the Model Context Protocol (stdio transport).

Usage:
    python -m eip_mcp.server
"""

from __future__ import annotations

import asyncio
import logging
import os
import sys
import time
from collections import deque
from threading import Lock
from typing import Any

import mcp.server.stdio
from mcp import types
from mcp.server.lowlevel import NotificationOptions, Server
from mcp.server.models import InitializationOptions

from eip_mcp import __version__
from eip_mcp import api_client
from eip_mcp import formatters
from eip_mcp import validators
from eip_mcp.validators import ValidationError

logging.basicConfig(level=logging.WARNING, stream=sys.stderr)
logger = logging.getLogger("eip-mcp")

_INSTRUCTIONS = """You have access to the Exploit Intelligence Platform (EIP) — 370K+ vulnerabilities and 105K+ exploits from NVD, CISA KEV, EPSS, ExploitDB, Metasploit, GitHub, and nomi-sec.

TOOL SELECTION GUIDE:
- User asks about vulnerabilities for a vendor/product/technology → search_vulnerabilities
- User mentions a specific CVE ID → get_vulnerability (full brief with ranked exploits, AI analysis, MITRE techniques)
- User asks about exploits by source/language/author/classification → search_exploits
- User asks "find reliable RCE exploits" or filters by attack type/complexity/reliability → search_exploits with attack_type, complexity, reliability, requires_auth filters
- User asks about trojans or backdoored exploits → search_exploits with llm_classification=trojan
- User asks for exploits targeting a specific CVE → search_exploits with cve= filter
- User asks for all exploits from a vendor → search_exploits with vendor= filter
- User asks about exploit authors → list_authors or get_author
- User asks about vulnerability types or CWEs → list_cwes or get_cwe
- User asks which vendors have the most vulns → list_vendors
- User describes their tech stack and wants risk assessment → audit_stack
- User wants Nuclei scanner templates or Shodan/FOFA/Google dorks → get_nuclei_templates
- User needs a pentest report finding → generate_finding
- User wants to read exploit source code → get_exploit_code

KEY CONCEPTS:
- EPSS: Probability (0-100%) that a CVE will be exploited in the next 30 days. >90% is very high risk.
- CVSS: Severity score (0-10). Critical >= 9.0, High >= 7.0, Medium >= 4.0.
- KEV: CISA Known Exploited Vulnerabilities — confirmed actively exploited in the wild.
- Exploit ranking: Metasploit modules > verified ExploitDB > GitHub by stars. Trojans flagged at bottom.
- AI Analysis: 61K+ exploits analyzed. Results include attack_type (RCE, SQLi, XSS, DoS, LPE, auth_bypass, info_leak), complexity (trivial/simple/moderate/complex), reliability (reliable/unreliable/untested), target_software, MITRE ATT&CK techniques, and deception_indicators for trojans.
MULTI-TOOL WORKFLOWS:
- "Tell me about CVE-XXXX": call get_vulnerability for full brief including exploits with MITRE techniques
- "What should I worry about?": call audit_stack with their technologies, then get_vulnerability on top findings
- "Is there a Metasploit module for X?": call get_vulnerability, check METASPLOIT MODULES section
- "Write a finding for my report": call generate_finding with CVE, target, and notes
"""

server = Server("eip-mcp", instructions=_INSTRUCTIONS)


# ---------------------------------------------------------------------------
# Local safety controls
# ---------------------------------------------------------------------------


class LocalRateLimitError(Exception):
    """Raised when local (process-level) rate limiting is exceeded."""

    def __init__(self, retry_after_s: float):
        self.retry_after_s = max(0.0, float(retry_after_s))
        super().__init__(f"Rate limited. Retry after {self.retry_after_s:.1f}s")


class RateLimiter:
    """Simple sliding-window limiter for tool calls.

    This limits *tool invocations* (not HTTP requests) per process.
    """

    def __init__(self, max_calls: int, period_s: float):
        self.max_calls = int(max_calls)
        self.period_s = float(period_s)
        self._times: deque[float] = deque()
        # Use a thread lock to avoid binding this limiter to a specific event loop.
        self._lock = Lock()

    def enabled(self) -> bool:
        return self.max_calls > 0 and self.period_s > 0

    def check_or_raise(self) -> None:
        if not self.enabled():
            return
        now = time.monotonic()
        with self._lock:
            cutoff = now - self.period_s
            while self._times and self._times[0] < cutoff:
                self._times.popleft()
            if len(self._times) >= self.max_calls:
                retry_after = (self._times[0] + self.period_s) - now
                raise LocalRateLimitError(retry_after)
            self._times.append(now)


def _int_env(name: str, default: int) -> int:
    raw = os.environ.get(name)
    if raw is None or raw == "":
        return default
    try:
        return int(raw)
    except ValueError:
        return default


def _float_env(name: str, default: float) -> float:
    raw = os.environ.get(name)
    if raw is None or raw == "":
        return default
    try:
        return float(raw)
    except ValueError:
        return default


_MAX_CONCURRENCY = max(0, _int_env("EIP_MCP_MAX_CONCURRENCY", 4))
_CONCURRENCY_SEM: asyncio.Semaphore | None = None

_MAX_CALLS = max(0, _int_env("EIP_MCP_MAX_CALLS", 60))
_CALL_PERIOD_S = max(0.0, _float_env("EIP_MCP_CALL_PERIOD_SECONDS", 60.0))
_RATE_LIMITER = RateLimiter(_MAX_CALLS, _CALL_PERIOD_S)

# ---------------------------------------------------------------------------
# Tool definitions
# ---------------------------------------------------------------------------

TOOLS = [
    types.Tool(
        name="search_vulnerabilities",
        description=(
            "Search the Exploit Intelligence Platform for vulnerabilities (CVEs). "
            "Returns a list of matching CVEs with CVSS scores, EPSS exploitation "
            "probability, exploit counts, and CISA KEV status. Supports full-text "
            "search, severity/vendor/product/ecosystem/CWE filters, and CVSS/EPSS "
            "thresholds. Use this for broad searches like 'apache vulnerabilities' "
            "or filtered queries like 'critical Fortinet CVEs with exploits'. "
            "Examples: query='apache httpd' with has_exploits=true; "
            "vendor='fortinet' with severity='critical' and is_kev=true sorted by epss_desc; "
            "cwe='89' with min_cvss=9 for critical SQL injection CVEs."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Search keywords (e.g. 'apache httpd', 'log4j'). Optional if filters are provided."},
                "severity": {"type": "string", "enum": ["critical", "high", "medium", "low"], "description": "Filter by severity level"},
                "has_exploits": {"type": "boolean", "description": "Only return CVEs with public exploit code"},
                "is_kev": {"type": "boolean", "description": "Only return CISA Known Exploited Vulnerabilities"},
                "has_nuclei": {"type": "boolean", "description": "Only return CVEs with Nuclei scanner templates"},
                "vendor": {"type": "string", "description": "Filter by vendor name (e.g. 'microsoft', 'fortinet')"},
                "product": {"type": "string", "description": "Filter by product name (e.g. 'exchange', 'pan-os')"},
                "ecosystem": {"type": "string", "enum": ["npm", "pip", "maven", "go", "crates", "nuget", "rubygems", "composer"], "description": "Filter by package ecosystem"},
                "cwe": {"type": "string", "description": "Filter by CWE ID (e.g. '79' or 'CWE-79')"},
                "min_cvss": {"type": "number", "description": "Minimum CVSS score (0-10)"},
                "min_epss": {"type": "number", "description": "Minimum EPSS score (0-1)"},
                "sort": {"type": "string", "enum": ["newest", "oldest", "cvss_desc", "epss_desc", "relevance"], "description": "Sort order (default: newest)"},
                "per_page": {"type": "integer", "description": "Results per page (1-25, default: 10)"},
            },
        },
    ),
    types.Tool(
        name="get_vulnerability",
        description=(
            "Get a full intelligence brief for a specific CVE. Returns detailed "
            "information including CVSS score and vector, EPSS exploitation probability, "
            "CISA KEV status, description, affected products, ranked exploits (grouped "
            "by Metasploit modules, verified ExploitDB, GitHub PoCs, and trojans), "
            "Nuclei scanner templates with recon dorks, alternate identifiers, and "
            "references. Use this when you need deep detail on a specific vulnerability. "
            "Exploits are ranked by quality: Metasploit modules first (peer-reviewed), "
            "then verified ExploitDB, then GitHub by stars. Trojans are flagged at the bottom."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "cve_id": {"type": "string", "description": "CVE identifier (e.g. 'CVE-2024-3400')"},
            },
            "required": ["cve_id"],
        },
    ),
    types.Tool(
        name="get_exploit_code",
        description=(
            "Retrieve the source code of a specific exploit by its exploit ID. "
            "Returns syntax-highlighted code from the exploit archive. If no file_path "
            "is specified, auto-selects the most relevant code file. Use this to "
            "analyze exploit mechanics, understand attack techniques, or review PoC code. "
            "Get the exploit_id from a get_vulnerability result (listed in the exploits section)."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "exploit_id": {"type": "integer", "description": "Exploit ID (from search or vulnerability detail results)"},
                "file_path": {"type": "string", "description": "Specific file path to view (optional — auto-selects if omitted)"},
            },
            "required": ["exploit_id"],
        },
    ),
    types.Tool(
        name="get_nuclei_templates",
        description=(
            "Get Nuclei scanner templates and recon dorks for a CVE. Returns template "
            "metadata, severity, verification status, tags, and ready-to-use Shodan, "
            "FOFA, and Google dork queries for target identification. Use this to plan "
            "scanning or reconnaissance for a specific vulnerability."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "cve_id": {"type": "string", "description": "CVE identifier (e.g. 'CVE-2024-27198')"},
            },
            "required": ["cve_id"],
        },
    ),
    types.Tool(
        name="get_platform_stats",
        description=(
            "Get platform-wide statistics from the Exploit Intelligence Platform. "
            "Returns total counts of vulnerabilities, exploits, KEV entries, Nuclei "
            "templates, vendors, and authors, plus the last data update timestamp."
        ),
        inputSchema={"type": "object", "properties": {}},
    ),
    types.Tool(
        name="check_health",
        description=(
            "Check the EIP API health and data freshness. Returns database status "
            "and timestamps for each of the 11 ingestion sources (NVD, KEV, EPSS, "
            "ExploitDB, GitHub, Metasploit, etc.)."
        ),
        inputSchema={"type": "object", "properties": {}},
    ),
    types.Tool(
        name="search_exploits",
        description=(
            "Search and browse exploits across the entire platform. Filter by source "
            "(github, metasploit, exploitdb, nomisec, ghsa), language (python, ruby, etc.), "
            "LLM classification (working_poc, trojan, suspicious, scanner, stub, writeup), "
            "author name, minimum GitHub stars, code availability, CVE ID, vendor, or product. "
            "Also filter by AI analysis fields: attack_type (RCE, SQLi, XSS, DoS, LPE, "
            "auth_bypass, info_leak), complexity (trivial, simple, moderate, complex), "
            "reliability (reliable, unreliable, untested), and requires_auth. "
            "Examples: llm_classification='trojan' for backdoored exploits; "
            "source='metasploit' for all Metasploit modules; "
            "attack_type='RCE' with reliability='reliable' for weaponizable RCE exploits; "
            "vendor='mitel' for all Mitel exploits; "
            "cve='CVE-2024-3400' for all exploits targeting a specific CVE."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "source": {"type": "string", "description": "Filter by source: github, metasploit, exploitdb, nomisec, ghsa, writeup"},
                "language": {"type": "string", "description": "Filter by language: python, ruby, go, c, etc."},
                "llm_classification": {"type": "string", "description": "Filter by LLM classification: working_poc, trojan, suspicious, scanner, stub, writeup, tool"},
                "attack_type": {"type": "string", "description": "Filter by attack type from AI analysis: RCE, SQLi, XSS, DoS, LPE, auth_bypass, info_leak, deserialization, other"},
                "complexity": {"type": "string", "description": "Filter by exploit complexity: trivial, simple, moderate, complex"},
                "reliability": {"type": "string", "description": "Filter by exploit reliability: reliable, unreliable, untested"},
                "requires_auth": {"type": "boolean", "description": "Filter by whether exploit requires authentication"},
                "author": {"type": "string", "description": "Filter by author name"},
                "min_stars": {"type": "integer", "description": "Minimum GitHub stars"},
                "has_code": {"type": "boolean", "description": "Only exploits with downloadable code"},
                "cve": {"type": "string", "description": "Filter by CVE ID (e.g. 'CVE-2024-3400') — returns all exploits for that CVE"},
                "vendor": {"type": "string", "description": "Filter by vendor name (e.g. 'mitel', 'fortinet') — returns exploits for all CVEs affecting that vendor"},
                "product": {"type": "string", "description": "Filter by product name (e.g. 'micollab', 'pan-os')"},
                "sort": {"type": "string", "enum": ["newest", "stars_desc"], "description": "Sort order"},
                "per_page": {"type": "integer", "description": "Results per page (1-25, default: 10)"},
            },
        },
    ),
    types.Tool(
        name="list_authors",
        description=(
            "List exploit authors/researchers ranked by exploit count. Returns the top "
            "security researchers with their exploit counts and handles. "
            "Use this when asked 'who are the top exploit authors?' or 'who writes the most exploits?'"
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "per_page": {"type": "integer", "description": "Results per page (1-50, default: 25)"},
            },
        },
    ),
    types.Tool(
        name="get_author",
        description=(
            "Get an exploit author's profile with all their exploits. Returns author name, "
            "handle, total exploit count, activity start date, and a paginated list of their "
            "exploits with CVE context. Use this when asked about a specific researcher "
            "like 'show me all exploits by Chocapikk'."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "author_name": {"type": "string", "description": "Author name (e.g. 'Chocapikk')"},
            },
            "required": ["author_name"],
        },
    ),
    types.Tool(
        name="list_cwes",
        description=(
            "List CWE (Common Weakness Enumeration) categories ranked by vulnerability count. "
            "Returns CWE IDs, names, short labels, exploit likelihood, and how many CVEs "
            "have that weakness. Use this when asked 'what are the most common vulnerability types?'"
        ),
        inputSchema={"type": "object", "properties": {}},
    ),
    types.Tool(
        name="get_cwe",
        description=(
            "Get details for a specific CWE including full name, description, exploit "
            "likelihood, parent CWE, and total vulnerability count. "
            "Example: cwe_id='CWE-79' returns details about Cross-Site Scripting."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "cwe_id": {"type": "string", "description": "CWE identifier (e.g. 'CWE-79' or '79')"},
            },
            "required": ["cwe_id"],
        },
    ),
    types.Tool(
        name="list_vendors",
        description=(
            "List software vendors ranked by vulnerability count. Returns the top 200 "
            "vendors with their total CVE counts. Use this when asked 'which vendors have "
            "the most vulnerabilities?' or to understand the threat landscape by vendor."
        ),
        inputSchema={"type": "object", "properties": {}},
    ),
    types.Tool(
        name="audit_stack",
        description=(
            "Audit a technology stack for exploitable vulnerabilities. Accepts a "
            "comma-separated list of technologies (max 5) and searches for critical/ "
            "high severity CVEs with public exploits for each one, sorted by EPSS "
            "exploitation probability. Use this when a user describes their "
            "infrastructure and wants to know what to patch first. "
            "Example: technologies='nginx, postgresql, node.js' returns a risk-sorted "
            "list of exploitable CVEs grouped by technology."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "technologies": {
                    "type": "string",
                    "description": "Comma-separated list of technologies (e.g. 'nginx, postgresql, node.js'). Max 5.",
                },
            },
            "required": ["technologies"],
        },
    ),
    types.Tool(
        name="generate_finding",
        description=(
            "Generate a pentest report finding in Markdown format for a specific CVE. "
            "Fetches full vulnerability detail and formats it as a professional finding "
            "with severity, CVSS, description, affected products, exploit availability, "
            "and references. Optionally include the target system tested and tester notes. "
            "The output is ready to paste into a pentest report. "
            "Example: cve_id='CVE-2024-3400', target='fw.corp.example.com', "
            "notes='Confirmed RCE via GlobalProtect gateway'."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "cve_id": {"type": "string", "description": "CVE identifier (e.g. 'CVE-2024-3400')"},
                "target": {"type": "string", "description": "Target system tested (e.g. 'fw.corp.example.com'). Optional."},
                "notes": {"type": "string", "description": "Tester notes to include in the finding. Optional."},
            },
            "required": ["cve_id"],
        },
    ),
]


@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    return TOOLS


@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any] | None):
    """Route tool calls to handlers with validation and error handling.

    Notes:
    - Tool calls run with local concurrency and rate limiting.
    - Sync work (HTTP calls) executes in a worker thread to avoid blocking the event loop.
    """
    args = arguments or {}

    try:
        _RATE_LIMITER.check_or_raise()

        global _CONCURRENCY_SEM
        if _MAX_CONCURRENCY and _CONCURRENCY_SEM is None:
            _CONCURRENCY_SEM = asyncio.Semaphore(_MAX_CONCURRENCY)

        if _CONCURRENCY_SEM is None:
            result = await asyncio.to_thread(_dispatch, name, args)
        else:
            async with _CONCURRENCY_SEM:
                result = await asyncio.to_thread(_dispatch, name, args)

        return _tool_result(result, is_error=False)

    except LocalRateLimitError as exc:
        return _tool_result(
            f"Rate limited locally to protect reliability. Retry after {exc.retry_after_s:.1f}s.",
            is_error=True,
        )
    except ValidationError as exc:
        return _tool_result(f"Validation error: {exc}", is_error=True)
    except api_client.APIError as exc:
        if exc.status_code == 404:
            ident = args.get("cve_id", args.get("exploit_id", args.get("author_name", args.get("cwe_id", "?"))))
            return _tool_result(f"Not found: {ident}", is_error=True)
        if exc.status_code == 429:
            return _tool_result(f"Rate limited by the EIP API. {exc.message}", is_error=True)
        if exc.status_code == 0:
            return _tool_result(f"Network error contacting the EIP API. {exc.message}", is_error=True)
        return _tool_result(f"API error: {exc.message}", is_error=True)
    except Exception:
        logger.exception("Unexpected error in tool %s", name)
        return _tool_result("Internal error processing request. Please try again.", is_error=True)


def _tool_result(text: str, *, is_error: bool) -> Any:
    """Build a CallTool result with explicit error signaling when supported."""
    content = [types.TextContent(type="text", text=text)]

    # Newer MCP SDKs expose CallToolResult (with an isError field) which clients
    # use to distinguish normal output from failures.
    if hasattr(types, "CallToolResult"):
        try:
            return types.CallToolResult(content=content, isError=is_error)
        except TypeError:
            # Some SDK versions may use snake_case; keep compatibility.
            return types.CallToolResult(content=content, is_error=is_error)  # type: ignore[arg-type]

    # Fallback for older SDKs.
    return content


# ---------------------------------------------------------------------------
# Tool dispatch
# ---------------------------------------------------------------------------

def _dispatch(name: str, args: dict[str, Any]) -> str:
    if name == "search_vulnerabilities":
        return _tool_search(args)
    elif name == "get_vulnerability":
        return _tool_get_vuln(args)
    elif name == "get_exploit_code":
        return _tool_get_code(args)
    elif name == "get_nuclei_templates":
        return _tool_get_nuclei(args)
    elif name == "get_platform_stats":
        return _tool_stats()
    elif name == "check_health":
        return _tool_health()
    elif name == "search_exploits":
        return _tool_search_exploits(args)
    elif name == "list_authors":
        return _tool_list_authors(args)
    elif name == "get_author":
        return _tool_get_author(args)
    elif name == "list_cwes":
        return _tool_list_cwes()
    elif name == "get_cwe":
        return _tool_get_cwe(args)
    elif name == "list_vendors":
        return _tool_list_vendors()
    elif name == "audit_stack":
        return _tool_audit_stack(args)
    elif name == "generate_finding":
        return _tool_generate_finding(args)
    else:
        raise ValidationError(f"Unknown tool: {name}")


# ---------------------------------------------------------------------------
# Core tool handlers
# ---------------------------------------------------------------------------

def _tool_search(args: dict) -> str:
    params: dict[str, Any] = {}
    if "query" in args and args["query"]:
        params["q"] = validators.validate_query(args["query"])
    if "severity" in args and args["severity"]:
        params["severity"] = validators.validate_severity(args["severity"])
    if args.get("has_exploits"):
        params["has_exploits"] = True
    if args.get("is_kev"):
        params["is_kev"] = True
    if args.get("has_nuclei"):
        params["has_nuclei"] = True
    if "vendor" in args and args["vendor"]:
        params["vendor"] = validators.validate_vendor(args["vendor"])
    if "product" in args and args["product"]:
        params["product"] = validators.validate_product(args["product"])
    if "ecosystem" in args and args["ecosystem"]:
        params["ecosystem"] = validators.validate_ecosystem(args["ecosystem"])
    if "cwe" in args and args["cwe"]:
        params["cwe"] = validators.validate_cwe(args["cwe"])
    if "min_cvss" in args and args["min_cvss"] is not None:
        params["min_cvss"] = validators.validate_cvss(args["min_cvss"])
    if "min_epss" in args and args["min_epss"] is not None:
        params["min_epss"] = validators.validate_epss(args["min_epss"])
    if "sort" in args and args["sort"]:
        params["sort"] = validators.validate_sort(args["sort"])
    if "per_page" in args and args["per_page"] is not None:
        params["per_page"] = validators.validate_per_page(args["per_page"])
    else:
        params["per_page"] = 10

    if not params or (len(params) == 1 and "per_page" in params):
        raise ValidationError("Provide a search query or at least one filter.")

    data = api_client.search_vulns(params)
    return formatters.format_search_results(data)


def _tool_get_vuln(args: dict) -> str:
    cve_id = validators.validate_cve_id(args.get("cve_id", ""))
    data = api_client.get_vuln_detail(cve_id)
    return formatters.format_vuln_detail(data)


def _tool_get_code(args: dict) -> str:
    exploit_id = validators.validate_exploit_id(args.get("exploit_id", ""))
    file_path = args.get("file_path")

    if file_path:
        file_path = validators.validate_file_path(file_path)
    else:
        # Auto-select: list files and pick the best one
        files = api_client.list_exploit_files(exploit_id)
        if not files:
            return f"No code files found for exploit {exploit_id}."
        file_path = _pick_main_file(files)
        if not file_path:
            return formatters.format_exploit_files(files, exploit_id) + "\n\nSpecify file_path to view a specific file."

    code = api_client.get_exploit_code(exploit_id, file_path)
    return formatters.format_exploit_code(code, file_path)


def _tool_get_nuclei(args: dict) -> str:
    cve_id = validators.validate_cve_id(args.get("cve_id", ""))
    data = api_client.get_vuln_detail(cve_id)
    templates = data.get("nuclei_templates") or []
    if not templates:
        return f"No Nuclei templates found for {cve_id}."
    return formatters.format_nuclei_templates(templates, header=True)


def _tool_stats() -> str:
    data = api_client.get_stats()
    return formatters.format_stats(data)


def _tool_health() -> str:
    data = api_client.get_health()
    return formatters.format_health(data)


# ---------------------------------------------------------------------------
# Browse tool handlers (new endpoints)
# ---------------------------------------------------------------------------

def _tool_search_exploits(args: dict) -> str:
    params: dict[str, Any] = {}
    if args.get("source"):
        params["source"] = validators.validate_source(args["source"])
    if args.get("language"):
        params["language"] = validators.validate_language(args["language"])
    if args.get("llm_classification"):
        params["llm_classification"] = validators.validate_llm_classification(args["llm_classification"])
    if args.get("attack_type"):
        params["attack_type"] = validators.validate_attack_type(args["attack_type"])
    if args.get("complexity"):
        params["complexity"] = validators.validate_complexity(args["complexity"])
    if args.get("reliability"):
        params["reliability"] = validators.validate_reliability(args["reliability"])
    if "requires_auth" in args and args["requires_auth"] is not None:
        if not isinstance(args["requires_auth"], bool):
            raise ValidationError("requires_auth must be true or false")
        params["requires_auth"] = args["requires_auth"]
    if args.get("author"):
        params["author"] = validators.validate_query(args["author"])
    if args.get("min_stars") is not None:
        params["min_stars"] = validators.validate_min_stars(args["min_stars"])
    if args.get("has_code") is not None:
        params["has_code"] = args["has_code"]
    if args.get("cve"):
        params["cve"] = validators.validate_cve_id(args["cve"])
    if args.get("vendor"):
        params["vendor"] = validators.validate_vendor(args["vendor"])
    if args.get("product"):
        params["product"] = validators.validate_product(args["product"])
    if args.get("sort"):
        s = args["sort"]
        if s not in ("newest", "stars_desc"):
            raise ValidationError("sort must be 'newest' or 'stars_desc'")
        params["sort"] = s
    if args.get("per_page") is not None:
        params["per_page"] = validators.validate_per_page(args["per_page"])
    else:
        params["per_page"] = 10

    data = api_client.search_exploits(params)
    return formatters.format_exploit_search(data)


def _tool_list_authors(args: dict) -> str:
    params: dict[str, Any] = {}
    if args.get("per_page") is not None:
        val = validators.validate_per_page(args["per_page"])
        params["per_page"] = val
    else:
        params["per_page"] = 25
    data = api_client.list_authors(params)
    return formatters.format_authors_list(data)


def _tool_get_author(args: dict) -> str:
    name = validators.validate_query(args.get("author_name", ""))
    if not name:
        raise ValidationError("author_name is required")
    data = api_client.get_author(name)
    return formatters.format_author_detail(data)


def _tool_list_cwes() -> str:
    data = api_client.list_cwes()
    return formatters.format_cwe_list(data)


def _tool_get_cwe(args: dict) -> str:
    cwe_id = validators.validate_cwe(args.get("cwe_id", ""))
    data = api_client.get_cwe(f"CWE-{cwe_id}")
    return formatters.format_cwe_detail(data)


def _tool_list_vendors() -> str:
    data = api_client.list_vendors()
    return formatters.format_vendors_list(data)


# ---------------------------------------------------------------------------
# Smart composite tools
# ---------------------------------------------------------------------------

def _tool_audit_stack(args: dict) -> str:
    techs = validators.validate_technologies(args.get("technologies", ""))
    results: dict[str, dict] = {}
    for tech in techs:
        try:
            data = api_client.search_vulns({
                "product": tech,
                "has_exploits": True,
                "sort": "epss_desc",
                "per_page": 10,
            })
            results[tech] = data
        except api_client.APIError:
            results[tech] = {"total": 0, "items": []}
    return formatters.format_stack_audit(results)


def _tool_generate_finding(args: dict) -> str:
    cve_id = validators.validate_cve_id(args.get("cve_id", ""))
    target = args.get("target")
    notes = args.get("notes")
    if target:
        target = validators.validate_query(target)
    if notes:
        notes = validators.validate_query(notes)
    data = api_client.get_vuln_detail(cve_id)
    return formatters.format_finding(data, target, notes)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _pick_main_file(files: list[dict]) -> str | None:
    """Auto-select the most relevant exploit file."""
    code_exts = {".py", ".rb", ".go", ".c", ".cpp", ".java", ".js", ".pl", ".php", ".sh", ".ps1"}
    code_files = [f for f in files if any(f["path"].endswith(ext) for ext in code_exts)]

    for pattern in ("exploit", "poc", "main", "scan", "vuln", "rce"):
        for f in code_files:
            if pattern in f["path"].lower():
                return f["path"]

    if code_files:
        return max(code_files, key=lambda f: f.get("size", 0))["path"]

    non_readme = [f for f in files if "readme" not in f["path"].lower()]
    if non_readme:
        return max(non_readme, key=lambda f: f.get("size", 0))["path"]

    return files[0]["path"] if files else None


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

async def _run():
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="eip-mcp",
                server_version=__version__,
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )


def main():
    """Entry point for the MCP server."""
    asyncio.run(_run())


if __name__ == "__main__":
    main()
