import os
import re
import sys
from pathlib import Path
from typing import Any, Literal

import arxiv
from fastmcp import FastMCP
from semanticscholar import AsyncSemanticScholar

from deeppresenter.utils.log import set_logger

mcp = FastMCP(name="Research")
PAGE_SIZE = int(os.getenv("ARXIV_PAGE_SIZE", "5"))
client = arxiv.Client(page_size=PAGE_SIZE)
sch = AsyncSemanticScholar()


@mcp.tool()
def search_papers(
    query: str,
    max_results: int | None = 5,
) -> dict[str, Any]:
    """
    Search for academic papers using arXiv query syntax.
    Best Practice: Start with broad queries and progressively add constraints if too many results are returned.

    QUERY PARAMETERS (field prefixes):
    - ti: Title field, e.g., ti:"attention mechanism"
    - au: Author name, e.g., au:"Hinton"
    - abs: Abstract content, e.g., abs:"transformer"
    - cat: Subject category, e.g., cat:cs.AI

    Args:
        query: Search query using arXiv syntax (keywords, field searches, boolean operators)
        max_results: Maximum number of results to return (default: 5)

    Returns:
        List of papers, each containing: title, authors, abstract, published date and pdf_url
    """
    # Map sort_by string to arxiv SortCriterion
    search = arxiv.Search(
        query=query,
        max_results=PAGE_SIZE,
    )

    papers = [
        {
            "title": paper.title,
            "authors": [author.name for author in paper.authors],
            "summary": paper.summary,
            "published": paper.published.strftime("%Y-%m-%d"),
            "pdf_url": paper.pdf_url,
        }
        for paper in client.results(search)
    ]
    if max_results is not None:
        papers = papers[:max_results]
    return {"total_results": len(papers), "papers": papers}


@mcp.tool()
async def get_paper_authors(arxiv_id: str) -> dict:
    """
    Get the authors of a paper by arxiv id.
    Args:
        arxiv_id: The id of the paper, must be in the format of ARXIV:2501.03936
    Returns:
        List of authors, containing their authorId and details.
    """
    assert re.fullmatch(r"ARXIV:\d{4}\.\d{4,5}(v\d+)?", arxiv_id), (
        "Invalid arxiv_id format. It should be like ARXIV:2501.03936"
    )
    fields = ["name", "citationCount", "affiliations"]
    authors = await sch.get_paper_authors(arxiv_id, fields=fields)
    return {"authors": [author._data for author in authors]}


@mcp.tool()
async def get_scholar_details(
    author_id: str,
    paper_start_index: int = 0,
    sort_by: Literal["citationCount", "year"] | None = None,
) -> dict[str, Any]:
    """
    Get the details of a scholar by author ID.

    Args:
        author_id: The Semantic Scholar author ID of the scholar
        paper_start_index: Starting index for paper pagination (default: 0)
        sort_by: Sort papers by "citationCount" or "year" (default: None - no sorting)

    Returns:
        Dictionary containing scholar details including:
        - Basic info: name, hIndex, homepage, paperCount
        - papers: List of papers with title, year, citationCount, publicationVenue, arxivId
    """
    fields = [
        "hIndex",
        "homepage",
        "paperCount",
        "papers",
        "papers.citationCount",
        "papers.externalIds",
        "papers.publicationDate",
        "papers.publicationVenue",
        "papers.title",
        "papers.year",
        "name",
    ]
    author = await sch.get_author(author_id, fields=fields)
    author, papers = author._data, author._data.pop("papers")
    processed_papers = []
    for p in papers:
        p["arxivId"] = p.pop("externalIds", {}).get("ArXiv", "")
        p.pop("paperId", "")
        processed_papers.append({k: v for k, v in p.items() if v is not None})
    if sort_by == "citationCount":
        processed_papers.sort(key=lambda x: x.get("citationCount", 0), reverse=True)
    elif sort_by == "year":
        processed_papers.sort(key=lambda x: x.get("year", 0), reverse=True)
    author["papers"] = processed_papers[
        paper_start_index : paper_start_index + PAGE_SIZE
    ]
    return author


if __name__ == "__main__":
    assert len(sys.argv) == 2, "Usage: python research.py <workspace>"
    work_dir = Path(sys.argv[1])
    assert work_dir.exists(), f"Workspace {work_dir} does not exist."
    os.chdir(work_dir)
    set_logger(f"research-{work_dir.stem}", work_dir / ".history" / "research.log")

    mcp.run(show_banner=False)
