import json
import argparse
import logging
from mcp.server.fastmcp import FastMCP
from cve_mcp_server.cve import CVE

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
LOGGER = logging.getLogger("mcp-cve")

mcp = FastMCP("mcp-cve")


@mcp.tool()
def query_cve(cve_id: str) -> str:
    cve = CVE(cve_id=cve_id)
    return json.dumps(cve.get_cve())


if __name__ == "__main__":
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="CVE MCP Server")
    parser.add_argument('transport', nargs='?', default='stdio', choices=['stdio', 'sse'],
                        help='Transport type (stdio or sse)')
    args = parser.parse_args()
    
    # Run the MCP server with the specified transport
    mcp.run(transport=args.transport)