import json
from typing import Any, Callable
from mcp.server.fastmcp import FastMCP
from paperscraper.pubmed.pubmed import get_pubmed_papers
from paperscraper.scholar.scholar import get_scholar_papers

# Initialize FastMCP server
mcp = FastMCP("paperscraper-mcpserver")

def _safe_execute(func: Callable, *args: Any, **kwargs: Any) -> str:
    """
    Executes a function safely, catches any exceptions, and returns a JSON string.
    """
    try:
        df = func(*args, **kwargs)
        if df.empty:
            return json.dumps({"status": "success", "data": []})
        records = df.to_dict(orient="records")
        return json.dumps({"status": "success", "data": records}, ensure_ascii=False)
    except Exception as e:
        return json.dumps({"status": "error", "message": str(e)}, ensure_ascii=False)

@mcp.tool()
def search_pubmed_papers(query: str, max_results: int = 10) -> str:
    """
    Search for papers on PubMed.
    
    Args:
        query: The search query formatted for the PubMed API.
        max_results: Maximum number of papers to retrieve (default: 10).
        
    Returns:
        JSON string containing the list of papers.
    """
    return _safe_execute(get_pubmed_papers, query=query, max_results=max_results)

@mcp.tool()
def search_scholar_papers(title: str) -> str:
    """
    Search for a specific paper title on Google Scholar and retrieve its metadata.
    Note that Google Scholar does not support complex Boolean logic and may easily rate limit or block bots.
    
    Args:
        title: The specific title or simple search string to look up.
        
    Returns:
        JSON string containing the list of papers.
    """
    # Scholar API might easily get blocked, we only request 1 or a few fields if possible, but the original API 
    # doesn't support max_results. We'll just call the default.
    return _safe_execute(get_scholar_papers, title=title)

def main():
    """Main entry point for the MCP server."""
    mcp.run()

if __name__ == "__main__":
    main()
