#!/usr/bin/env python3
"""
OpenAlex MCP Server for Author Disambiguation

Provides Claude Agent SDK-compliant tools for searching and retrieving author
information from OpenAlex with domain awareness and filtering.

This is an embedded MCP server that integrates with Claude Agent SDK.
"""

import os
import json
import asyncio
from typing import Any, Optional, List, Dict

import pyalex
from pyalex import Authors, Works, Institutions

from claude_agent_sdk import (
    tool,
    create_sdk_mcp_server,
    ClaudeSDKClient,
    ClaudeAgentOptions,
)


# Domain classification keywords
DOMAIN_KEYWORDS = {
    "life_sciences": [
        "biology", "genetics", "molecular biology", "biochemistry", "cell biology",
        "developmental biology", "microbiology", "immunology", "neuroscience",
        "ecology", "evolutionary biology", "botany", "zoology", "physiology",
        "biophysics", "biotechnology", "genomics", "proteomics", "bioinformatics"
    ],
    "health_sciences": [
        "medicine", "clinical medicine", "pathology", "pharmacology", "epidemiology",
        "public health", "nursing", "dentistry", "veterinary medicine", "psychiatry",
        "surgery", "pediatrics", "oncology", "cardiology", "neurology"
    ],
    "physical_sciences": [
        "physics", "chemistry", "astronomy", "astrophysics", "materials science",
        "quantum mechanics", "thermodynamics", "optics", "nuclear physics",
        "particle physics", "geophysics", "atmospheric science"
    ],
    "earth_environmental_sciences": [
        "geology", "geography", "environmental science", "oceanography",
        "meteorology", "climatology", "soil science", "hydrology"
    ],
    "engineering": [
        "engineering", "mechanical engineering", "electrical engineering",
        "civil engineering", "chemical engineering", "computer engineering",
        "biomedical engineering", "aerospace engineering"
    ],
    "computer_science": [
        "computer science", "artificial intelligence", "machine learning",
        "software engineering", "data science", "algorithms", "programming"
    ],
    "mathematics": [
        "mathematics", "statistics", "applied mathematics", "pure mathematics",
        "algebra", "calculus", "geometry", "topology"
    ],
    "social_sciences": [
        "psychology", "sociology", "anthropology", "economics", "political science",
        "education", "linguistics", "archaeology"
    ],
    "arts_humanities": [
        "history", "philosophy", "literature", "art", "music", "religion",
        "cultural studies", "ethics"
    ]
}


# Initialize PyAlex configuration
def _init_pyalex():
    """Initialize PyAlex with email configuration."""
    email = os.getenv('OPENALEX_API_KEY')
    if not email:
        raise ValueError(
            "OPENALEX_API_KEY environment variable must be set. "
            "This should be your email address for polite pool access."
        )
    pyalex.config.email = email


def determine_primary_domain(x_concepts: List[Dict]) -> str:
    """
    Determine primary research domain based on x_concepts.

    Uses keyword matching against concept names to classify research domain.
    Returns the domain with the highest weighted score.

    Args:
        x_concepts: List of concept dictionaries with 'display_name' and 'score'

    Returns:
        Primary domain string (e.g., "life_sciences") or "other" if no clear match
    """
    if not x_concepts:
        return "other"

    domain_scores = {domain: 0.0 for domain in DOMAIN_KEYWORDS}

    for concept in x_concepts:
        concept_name = concept.get("display_name", "").lower()
        concept_score = concept.get("score", 0)

        for domain, keywords in DOMAIN_KEYWORDS.items():
            for keyword in keywords:
                if keyword in concept_name:
                    domain_scores[domain] += concept_score
                    break

    if max(domain_scores.values()) > 0:
        return max(domain_scores, key=domain_scores.get)
    return "other"


def _search_authors_by_name(
    name: str,
    per_page: int = 200,
    preferred_domain: Optional[str] = None
) -> Dict[str, Any]:
    """
    Search for authors by name with optional domain-aware ranking.

    Args:
        name: Author name to search for
        per_page: Number of results to return (max 200)
        preferred_domain: Optional domain to prefer in ranking

    Returns:
        Dictionary with count, preferred_domain, and results list
    """
    _init_pyalex()

    # Search authors
    authors = Authors().search(name).get(per_page=per_page)

    results = []
    for author in authors:
        # Determine primary domain
        x_concepts = author.get("x_concepts", [])
        primary_domain = determine_primary_domain(x_concepts)

        # Calculate domain match score for ranking
        domain_match_score = 0
        if preferred_domain and primary_domain == preferred_domain:
            domain_match_score = 100

        # Get affiliation info
        affiliations = []
        last_known_institutions = author.get("last_known_institutions", [])
        for inst in last_known_institutions:
            affiliations.append({
                "institution_id": inst.get("id", ""),
                "institution_name": inst.get("display_name", ""),
                "country_code": inst.get("country_code", "")
            })

        results.append({
            "openalex_id": author.get("id", "").split("/")[-1],
            "openalex_url": author.get("id", ""),
            "display_name": author.get("display_name", ""),
            "orcid": author.get("orcid"),
            "works_count": author.get("works_count", 0),
            "cited_by_count": author.get("cited_by_count", 0),
            "primary_domain": primary_domain,
            "domain_match_score": domain_match_score,
            "affiliations": affiliations,
            "x_concepts": x_concepts[:5] if x_concepts else []
        })

    # Sort by domain match score (if preferred_domain specified), then by citation count
    if preferred_domain:
        results.sort(key=lambda x: (x["domain_match_score"], x["cited_by_count"]), reverse=True)
    else:
        results.sort(key=lambda x: x["cited_by_count"], reverse=True)

    return {
        "count": len(results),
        "preferred_domain": preferred_domain,
        "results": results
    }


def _search_authors_by_orcid(orcid: str) -> Dict[str, Any]:
    """
    Search for author by ORCID.

    Args:
        orcid: ORCID identifier (with or without prefix)

    Returns:
        Dictionary with found flag and author details
    """
    _init_pyalex()

    # Normalize ORCID
    orcid = orcid.replace("https://orcid.org/", "")

    # Search by ORCID filter
    authors = Authors().filter(orcid=orcid).get()

    if not authors:
        return {
            "found": False,
            "message": f"No author found with ORCID {orcid}"
        }

    author = authors[0]

    # Determine primary domain
    x_concepts = author.get("x_concepts", [])
    primary_domain = determine_primary_domain(x_concepts)

    # Get affiliation info
    affiliations = []
    last_known_institutions = author.get("last_known_institutions", [])
    for inst in last_known_institutions:
        affiliations.append({
            "institution_id": inst.get("id", ""),
            "institution_name": inst.get("display_name", ""),
            "country_code": inst.get("country_code", "")
        })

    # Get summary stats
    summary_stats = author.get("summary_stats", {})

    return {
        "found": True,
        "openalex_id": author.get("id", "").split("/")[-1],
        "openalex_url": author.get("id", ""),
        "display_name": author.get("display_name", ""),
        "orcid": author.get("orcid"),
        "works_count": author.get("works_count", 0),
        "cited_by_count": author.get("cited_by_count", 0),
        "h_index": summary_stats.get("h_index", 0),
        "i10_index": summary_stats.get("i10_index", 0),
        "primary_domain": primary_domain,
        "affiliations": affiliations,
        "x_concepts": x_concepts[:10] if x_concepts else []
    }


def _search_authors_by_name_and_institution(
    name: str,
    institution_name: str,
    per_page: int = 200,
    preferred_domain: Optional[str] = None
) -> Dict[str, Any]:
    """
    Search for authors by name filtered by institution.

    Args:
        name: Author name to search for
        institution_name: Institution name to filter by
        per_page: Number of results to return (max 200)
        preferred_domain: Optional domain to prefer in ranking

    Returns:
        Dictionary with institution_searched, count, and results list
    """
    _init_pyalex()

    # First, find the institution
    institutions = Institutions().search(institution_name).get(per_page=1)

    if not institutions:
        return {
            "error": "Institution not found",
            "searched_for": institution_name
        }

    institution = institutions[0]
    institution_id = institution.get("id", "").split("/")[-1]

    # Search authors by name and filter by institution
    authors = Authors().search(name).filter(
        last_known_institutions={"id": institution["id"]}
    ).get(per_page=per_page)

    results = []
    for author in authors:
        # Determine primary domain
        x_concepts = author.get("x_concepts", [])
        primary_domain = determine_primary_domain(x_concepts)

        # Calculate domain match score
        domain_match_score = 0
        if preferred_domain and primary_domain == preferred_domain:
            domain_match_score = 100

        # Get affiliation info
        affiliations = []
        last_known_institutions = author.get("last_known_institutions", [])
        for inst in last_known_institutions:
            affiliations.append({
                "institution_id": inst.get("id", ""),
                "institution_name": inst.get("display_name", ""),
                "country_code": inst.get("country_code", "")
            })

        results.append({
            "openalex_id": author.get("id", "").split("/")[-1],
            "openalex_url": author.get("id", ""),
            "display_name": author.get("display_name", ""),
            "orcid": author.get("orcid"),
            "works_count": author.get("works_count", 0),
            "cited_by_count": author.get("cited_by_count", 0),
            "primary_domain": primary_domain,
            "domain_match_score": domain_match_score,
            "affiliations": affiliations,
            "x_concepts": x_concepts[:5] if x_concepts else []
        })

    # Sort by domain match score, then citation count
    if preferred_domain:
        results.sort(key=lambda x: (x["domain_match_score"], x["cited_by_count"]), reverse=True)
    else:
        results.sort(key=lambda x: x["cited_by_count"], reverse=True)

    return {
        "institution_searched": {
            "id": institution_id,
            "display_name": institution.get("display_name", "")
        },
        "count": len(results),
        "preferred_domain": preferred_domain,
        "results": results
    }


def _get_author_details(openalex_author_id: str) -> Dict[str, Any]:
    """
    Get detailed information for a specific author.

    Args:
        openalex_author_id: OpenAlex author ID (with or without prefix)

    Returns:
        Dictionary with found flag and comprehensive author details
    """
    _init_pyalex()

    # Normalize ID
    if not openalex_author_id.startswith("https://"):
        openalex_author_id = f"https://openalex.org/{openalex_author_id}"

    try:
        author = Authors()[openalex_author_id]
    except Exception:
        return {
            "found": False,
            "message": f"Author not found with ID {openalex_author_id}"
        }

    # Determine primary domain
    x_concepts = author.get("x_concepts", [])
    primary_domain = determine_primary_domain(x_concepts)

    # Get last known institution
    last_known_institutions = author.get("last_known_institutions", [])
    last_known_institution = None
    if last_known_institutions:
        inst = last_known_institutions[0]
        last_known_institution = {
            "id": inst.get("id", ""),
            "display_name": inst.get("display_name", ""),
            "country_code": inst.get("country_code", "")
        }

    # Get all affiliations
    affiliations = []
    for inst in last_known_institutions:
        affiliations.append({
            "institution_id": inst.get("id", ""),
            "institution_name": inst.get("display_name", ""),
            "country_code": inst.get("country_code", "")
        })

    # Get summary stats
    summary_stats = author.get("summary_stats", {})

    # Get counts by year
    counts_by_year = author.get("counts_by_year", [])

    return {
        "found": True,
        "openalex_id": author.get("id", "").split("/")[-1],
        "openalex_url": author.get("id", ""),
        "display_name": author.get("display_name", ""),
        "orcid": author.get("orcid"),
        "works_count": author.get("works_count", 0),
        "cited_by_count": author.get("cited_by_count", 0),
        "h_index": summary_stats.get("h_index", 0),
        "i10_index": summary_stats.get("i10_index", 0),
        "primary_domain": primary_domain,
        "last_known_institution": last_known_institution,
        "affiliations": affiliations,
        "x_concepts": x_concepts[:10] if x_concepts else [],
        "counts_by_year": counts_by_year[:10] if counts_by_year else []
    }


def _get_author_recent_works(
    openalex_author_id: str,
    per_page: int = 10
) -> Dict[str, Any]:
    """
    Get recent publications for a specific author.

    Args:
        openalex_author_id: OpenAlex author ID (with or without prefix)
        per_page: Number of works to return (max 200)

    Returns:
        Dictionary with author_id, count, and publications list
    """
    _init_pyalex()

    # Normalize ID
    if not openalex_author_id.startswith("https://"):
        openalex_author_id = f"https://openalex.org/{openalex_author_id}"

    # Get author's works, sorted by publication date (most recent first)
    works = Works().filter(
        authorships={"author": {"id": openalex_author_id}}
    ).sort(publication_date="desc").get(per_page=per_page)

    publications = []
    for work in works:
        # Get primary location (journal)
        primary_location = work.get("primary_location", {})
        journal = None
        if primary_location:
            source = primary_location.get("source", {})
            if source:
                journal = source.get("display_name")

        # Get author's affiliation for this work
        author_affiliations = []
        authorships = work.get("authorships", [])
        for authorship in authorships:
            if authorship.get("author", {}).get("id") == openalex_author_id:
                institutions = authorship.get("institutions", [])
                for inst in institutions:
                    author_affiliations.append({
                        "institution_id": inst.get("id", ""),
                        "institution_name": inst.get("display_name", "")
                    })
                break

        publications.append({
            "id": work.get("id", "").split("/")[-1],
            "title": work.get("title", ""),
            "publication_year": work.get("publication_year"),
            "publication_date": work.get("publication_date"),
            "doi": work.get("doi"),
            "type": work.get("type"),
            "cited_by_count": work.get("cited_by_count", 0),
            "is_oa": work.get("open_access", {}).get("is_oa", False),
            "journal": journal,
            "author_affiliations": author_affiliations
        })

    return {
        "author_id": openalex_author_id.split("/")[-1],
        "count": len(publications),
        "publications": publications
    }


def to_message(content: Any) -> dict[str, Any]:
    """Convert content to MCP message format."""
    return {
        "content": [{
            "type": "text",
            "text": json.dumps(content, indent=2)
        }]
    }


@tool(
    "search_authors_by_name",
    "Search for authors by name with optional domain-aware ranking. Returns comprehensive "
    "profile information including affiliations, publication metrics, and research domains.",
    {
        "name": str,
        "per_page": int,
        "preferred_domain": str,
    }
)
async def search_authors_by_name(args: dict[str, Any]) -> dict[str, Any]:
    """
    Search for authors by name with optional domain-aware ranking.

    Args (from args dict):
        name: The researcher's name to search for (e.g., "Yves Barde", "J. Smith")
        per_page: Number of results to return (default: 200, max: 200)
        preferred_domain: Optional domain filter for ranking results. One of:
            - "life_sciences": Biology, genetics, molecular biology, etc.
            - "health_sciences": Medicine, clinical research, public health, etc.
            - "physical_sciences": Physics, chemistry, astronomy, etc.
            - "social_sciences": Economics, sociology, psychology, etc.
            When specified, results matching this domain are ranked higher.

    Returns:
        MCP message with JSON containing:
        - count: Number of results found
        - preferred_domain: The domain filter applied (if any)
        - results: List of author profiles with openalex_id, display_name, orcid,
                   works_count, cited_by_count, primary_domain, affiliations, etc.
    """
    result = await asyncio.to_thread(
        _search_authors_by_name,
        name=args["name"],
        per_page=args.get("per_page", 200),
        preferred_domain=args.get("preferred_domain")
    )
    return to_message(result)


@tool(
    "search_authors_by_orcid",
    "Search for a specific author by ORCID identifier (most reliable method). "
    "ORCID is a unique persistent identifier for researchers.",
    {
        "orcid": str,
    }
)
async def search_authors_by_orcid(args: dict[str, Any]) -> dict[str, Any]:
    """
    Search for a specific author by ORCID identifier.

    Args (from args dict):
        orcid: The ORCID identifier, with or without prefix
               Examples: "0000-0002-7627-461X" or "https://orcid.org/0000-0002-7627-461X"

    Returns:
        MCP message with JSON containing:
        - found: Boolean indicating if author was found
        - openalex_id, display_name, orcid, works_count, cited_by_count, h_index, etc.
        If not found: {found: False, message: "..."}
    """
    result = await asyncio.to_thread(
        _search_authors_by_orcid,
        orcid=args["orcid"]
    )
    return to_message(result)


@tool(
    "search_authors_by_name_and_institution",
    "Search for authors by name filtered by institutional affiliation. "
    "Useful when you have partial name information but know the institution.",
    {
        "name": str,
        "institution_name": str,
        "per_page": int,
        "preferred_domain": str,
    }
)
async def search_authors_by_name_and_institution(args: dict[str, Any]) -> dict[str, Any]:
    """
    Search for authors by name filtered by institutional affiliation.

    Args (from args dict):
        name: The researcher's name to search for
        institution_name: Institution name to filter by (e.g., "EMBO", "Harvard", "Max Planck")
        per_page: Number of results to return (default: 200, max: 200)
        preferred_domain: Optional domain filter (see search_authors_by_name for options)

    Returns:
        MCP message with JSON containing:
        - institution_searched: Info about the matched institution (id, display_name)
        - count: Number of matching authors
        - preferred_domain: Domain filter applied (if any)
        - results: List of author profiles
        Or if institution not found: {error: "...", searched_for: "..."}
    """
    result = await asyncio.to_thread(
        _search_authors_by_name_and_institution,
        name=args["name"],
        institution_name=args["institution_name"],
        per_page=args.get("per_page", 200),
        preferred_domain=args.get("preferred_domain")
    )
    return to_message(result)


@tool(
    "get_author_details",
    "Get complete profile for a specific author by OpenAlex ID. "
    "Returns detailed information about profile, affiliations, metrics, and research topics.",
    {
        "openalex_author_id": str,
    }
)
async def get_author_details(args: dict[str, Any]) -> dict[str, Any]:
    """
    Get complete profile for a specific author by OpenAlex ID.

    Args (from args dict):
        openalex_author_id: OpenAlex author ID, with or without prefix
                           Examples: "A5074091984" or "https://openalex.org/A5074091984"

    Returns:
        MCP message with JSON containing:
        - found: Boolean indicating if author was found
        - openalex_id, display_name, orcid, works_count, cited_by_count, h_index, i10_index
        - primary_domain, last_known_institution, affiliations, x_concepts, counts_by_year
        If not found: {found: False, message: "..."}
    """
    result = await asyncio.to_thread(
        _get_author_details,
        openalex_author_id=args["openalex_author_id"]
    )
    return to_message(result)


@tool(
    "get_author_recent_works",
    "Get recent publications for a specific author to verify their identity and research activity. "
    "Useful for verifying correct author and understanding current research focus.",
    {
        "openalex_author_id": str,
        "per_page": int,
    }
)
async def get_author_recent_works(args: dict[str, Any]) -> dict[str, Any]:
    """
    Get recent publications for a specific author.

    Args (from args dict):
        openalex_author_id: OpenAlex author ID, with or without prefix
                           Examples: "A5074091984" or "https://openalex.org/A5074091984"
        per_page: Number of recent works to return (default: 10, max: 200)

    Returns:
        MCP message with JSON containing:
        - author_id: The OpenAlex author ID used
        - count: Number of publications returned
        - publications: List of recent works with id, title, publication_year, doi,
                       type, cited_by_count, is_oa, journal, author_affiliations
    """
    result = await asyncio.to_thread(
        _get_author_recent_works,
        openalex_author_id=args["openalex_author_id"],
        per_page=args.get("per_page", 10)
    )
    return to_message(result)


# ============================================================================
# WORKS SEARCH FUNCTIONS
# ============================================================================

def _search_works_by_title(
    title: str,
    per_page: int = 10
) -> Dict[str, Any]:
    """
    Search for works (papers) by title.

    Args:
        title: Work title to search for
        per_page: Number of results to return (max 200)

    Returns:
        Dictionary with count and results list
    """
    _init_pyalex()

    # Remove commas (interpreted as URL parameters by OpenAlex)
    clean_title = title.replace(',', '')

    # Search works
    works = Works().search_filter(title=clean_title).get(per_page=per_page)

    results = []
    for work in works:
        # Extract authorship information
        authorships = []
        for authorship in work.get("authorships", [])[:10]:  # Limit to first 10 authors
            author = authorship.get("author", {})
            authorships.append({
                "author_id": author.get("id", ""),
                "display_name": author.get("display_name", ""),
                "orcid": author.get("orcid"),
                "author_position": authorship.get("author_position", ""),
            })

        # Get journal/source information
        primary_location = work.get("primary_location", {})
        source = primary_location.get("source", {}) if primary_location else {}

        results.append({
            "openalex_id": work.get("id", ""),
            "doi": work.get("doi"),
            "title": work.get("title", ""),
            "publication_year": work.get("publication_year"),
            "publication_date": work.get("publication_date"),
            "type": work.get("type", ""),
            "cited_by_count": work.get("cited_by_count", 0),
            "is_oa": work.get("open_access", {}).get("is_oa", False),
            "journal": source.get("display_name") if source else None,
            "authorships": authorships,
        })

    return {
        "count": len(results),
        "results": results
    }


def _search_works_by_doi(doi: str) -> Dict[str, Any]:
    """
    Search for a work by DOI.

    Args:
        doi: DOI of the work (with or without prefix)

    Returns:
        Dictionary with found flag and work details
    """
    _init_pyalex()

    # Normalize DOI
    if not doi.startswith("https://doi.org/"):
        doi = f"https://doi.org/{doi}"

    try:
        work = Works()[doi]

        # Extract authorship information
        authorships = []
        for authorship in work.get("authorships", []):
            author = authorship.get("author", {})
            authorships.append({
                "author_id": author.get("id", ""),
                "display_name": author.get("display_name", ""),
                "orcid": author.get("orcid"),
                "author_position": authorship.get("author_position", ""),
            })

        # Get journal/source information
        primary_location = work.get("primary_location", {})
        source = primary_location.get("source", {}) if primary_location else {}

        return {
            "found": True,
            "openalex_id": work.get("id", ""),
            "doi": work.get("doi"),
            "title": work.get("title", ""),
            "publication_year": work.get("publication_year"),
            "publication_date": work.get("publication_date"),
            "type": work.get("type", ""),
            "cited_by_count": work.get("cited_by_count", 0),
            "is_oa": work.get("open_access", {}).get("is_oa", False),
            "journal": source.get("display_name") if source else None,
            "authorships": authorships,
        }
    except Exception as e:
        return {
            "found": False,
            "message": f"No work found with DOI {doi}: {str(e)}"
        }


def _search_works_by_title_and_author(
    title: str,
    author_name: str,
    per_page: int = 10
) -> Dict[str, Any]:
    """
    Search for works by title with author name filter.

    Args:
        title: Work title to search for
        author_name: Author name to filter by
        per_page: Number of results to return (max 200)

    Returns:
        Dictionary with count, author_searched, and results list
    """
    _init_pyalex()

    # Remove commas
    clean_title = title.replace(',', '')

    # Search works by title, then filter by author in Python
    # (OpenAlex API doesn't support combining title search with author filter easily)
    works = Works().search_filter(title=clean_title).get(per_page=per_page * 3)

    # Filter results by author name in Python
    filtered_works = []
    norm_author = author_name.lower()
    for work in works:
        # Check if author is in authorships
        for authorship in work.get("authorships", []):
            author = authorship.get("author", {})
            display_name = author.get("display_name", "").lower()
            if norm_author in display_name:
                filtered_works.append(work)
                break  # Found author, no need to check other authorships
        
        if len(filtered_works) >= per_page:
            break
    
    works = filtered_works

    results = []
    for work in works:
        # Extract authorship information
        authorships = []
        for authorship in work.get("authorships", [])[:10]:
            author = authorship.get("author", {})
            authorships.append({
                "author_id": author.get("id", ""),
                "display_name": author.get("display_name", ""),
                "orcid": author.get("orcid"),
                "author_position": authorship.get("author_position", ""),
            })

        # Get journal/source information
        primary_location = work.get("primary_location", {})
        source = primary_location.get("source", {}) if primary_location else {}

        results.append({
            "openalex_id": work.get("id", ""),
            "doi": work.get("doi"),
            "title": work.get("title", ""),
            "publication_year": work.get("publication_year"),
            "publication_date": work.get("publication_date"),
            "type": work.get("type", ""),
            "cited_by_count": work.get("cited_by_count", 0),
            "is_oa": work.get("open_access", {}).get("is_oa", False),
            "journal": source.get("display_name") if source else None,
            "authorships": authorships,
        })

    return {
        "count": len(results),
        "author_searched": author_name,
        "results": results
    }


def _get_work_details(openalex_work_id: str) -> Dict[str, Any]:
    """
    Get complete details for a specific work.

    Args:
        openalex_work_id: OpenAlex work ID (with or without URL prefix)

    Returns:
        Dictionary with work details
    """
    _init_pyalex()

    # Normalize ID
    if not openalex_work_id.startswith("https://openalex.org/"):
        if not openalex_work_id.startswith("W"):
            openalex_work_id = f"W{openalex_work_id}"
        openalex_work_id = f"https://openalex.org/{openalex_work_id}"

    try:
        work = Works()[openalex_work_id]

        # Extract full authorship information
        authorships = []
        for authorship in work.get("authorships", []):
            author = authorship.get("author", {})
            institutions = []
            for inst in authorship.get("institutions", []):
                institutions.append({
                    "id": inst.get("id", ""),
                    "display_name": inst.get("display_name", ""),
                    "country_code": inst.get("country_code", ""),
                })

            authorships.append({
                "author_id": author.get("id", ""),
                "display_name": author.get("display_name", ""),
                "orcid": author.get("orcid"),
                "author_position": authorship.get("author_position", ""),
                "is_corresponding": authorship.get("is_corresponding", False),
                "institutions": institutions,
            })

        # Get journal/source information
        primary_location = work.get("primary_location", {})
        source = primary_location.get("source", {}) if primary_location else {}

        # Get abstract inverted index (if available)
        abstract_inverted_index = work.get("abstract_inverted_index")

        # Get concepts/topics
        concepts = work.get("concepts", [])[:10]  # Top 10 concepts

        return {
            "found": True,
            "openalex_id": work.get("id", ""),
            "doi": work.get("doi"),
            "title": work.get("title", ""),
            "publication_year": work.get("publication_year"),
            "publication_date": work.get("publication_date"),
            "type": work.get("type", ""),
            "cited_by_count": work.get("cited_by_count", 0),
            "is_oa": work.get("open_access", {}).get("is_oa", False),
            "oa_url": work.get("open_access", {}).get("oa_url"),
            "journal": source.get("display_name") if source else None,
            "journal_issn": source.get("issn") if source else None,
            "has_abstract": abstract_inverted_index is not None,
            "authorships": authorships,
            "concepts": concepts,
            "referenced_works_count": work.get("referenced_works_count", 0),
        }
    except Exception as e:
        return {
            "found": False,
            "message": f"Work not found: {str(e)}"
        }


# ============================================================================
# WORKS SEARCH TOOLS (MCP)
# ============================================================================

@tool(
    "search_works_by_title",
    "Search for academic works (papers) by title. Returns work ID, DOI, title, "
    "publication year, authors, journal, and citation count. Useful for finding "
    "papers when you have the title or partial title.",
    {
        "title": str,
        "per_page": int,
    }
)
async def search_works_by_title(args: dict[str, Any]) -> dict[str, Any]:
    """
    MCP tool for searching works by title.

    Returns:
        MCP message with JSON containing:
        - count: Number of works found
        - results: List of works with openalex_id, doi, title, publication_year,
                  type, cited_by_count, is_oa, journal, and authorships
    """
    result = await asyncio.to_thread(
        _search_works_by_title,
        title=args["title"],
        per_page=args.get("per_page", 10)
    )
    return to_message(result)


@tool(
    "search_works_by_doi",
    "Search for a specific work by DOI (Digital Object Identifier). "
    "Most reliable method to get exact work when DOI is known.",
    {
        "doi": str,
    }
)
async def search_works_by_doi(args: dict[str, Any]) -> dict[str, Any]:
    """
    MCP tool for searching works by DOI.

    Returns:
        MCP message with JSON containing:
        - found: Boolean indicating if work was found
        - openalex_id, doi, title, publication_year, type, cited_by_count,
          is_oa, journal, and authorships (if found)
    """
    result = await asyncio.to_thread(
        _search_works_by_doi,
        doi=args["doi"]
    )
    return to_message(result)


@tool(
    "search_works_by_title_and_author",
    "Search for works by title filtered by author name. Useful when you know "
    "both the paper title and author to narrow down results.",
    {
        "title": str,
        "author_name": str,
        "per_page": int,
    }
)
async def search_works_by_title_and_author(args: dict[str, Any]) -> dict[str, Any]:
    """
    MCP tool for searching works by title and author.

    Returns:
        MCP message with JSON containing:
        - count: Number of works found
        - author_searched: The author name used for filtering
        - results: List of works with openalex_id, doi, title, publication_year,
                  type, cited_by_count, is_oa, journal, and authorships
    """
    result = await asyncio.to_thread(
        _search_works_by_title_and_author,
        title=args["title"],
        author_name=args["author_name"],
        per_page=args.get("per_page", 10)
    )
    return to_message(result)


@tool(
    "get_work_details",
    "Get complete details for a specific work by OpenAlex work ID. "
    "Returns comprehensive information including full author list with affiliations, "
    "abstract availability, concepts, and citation information.",
    {
        "openalex_work_id": str,
    }
)
async def get_work_details(args: dict[str, Any]) -> dict[str, Any]:
    """
    MCP tool for getting complete work details.

    Returns:
        MCP message with JSON containing:
        - found: Boolean indicating if work was found
        - Complete work information including openalex_id, doi, title,
          publication_year, type, cited_by_count, is_oa, oa_url, journal,
          has_abstract, authorships (with institutions), concepts, and more
    """
    result = await asyncio.to_thread(
        _get_work_details,
        openalex_work_id=args["openalex_work_id"]
    )
    return to_message(result)


# Create the MCP server
openalex_mcp_server = create_sdk_mcp_server(
    name="openalex_mcp",
    tools=[
        # Author search tools
        search_authors_by_name,
        search_authors_by_orcid,
        search_authors_by_name_and_institution,
        get_author_details,
        get_author_recent_works,
        # Works search tools
        search_works_by_title,
        search_works_by_doi,
        search_works_by_title_and_author,
        get_work_details,
    ],
)


# Create Claude Agent options with MCP server and allowed tools
OPENALEX_MCP_OPTIONS = ClaudeAgentOptions(
    mcp_servers={"openalex_mcp": openalex_mcp_server},
    allowed_tools=[
        # Author search tools
        "mcp__openalex_mcp__search_authors_by_name",
        "mcp__openalex_mcp__search_authors_by_orcid",
        "mcp__openalex_mcp__search_authors_by_name_and_institution",
        "mcp__openalex_mcp__get_author_details",
        "mcp__openalex_mcp__get_author_recent_works",
        # Works search tools
        "mcp__openalex_mcp__search_works_by_title",
        "mcp__openalex_mcp__search_works_by_doi",
        "mcp__openalex_mcp__search_works_by_title_and_author",
        "mcp__openalex_mcp__get_work_details",
        # Web tools
        "WebSearch",
        "WebFetch",
    ],
)


# Create Claude SDK client
OPENALEX_MCP_CLIENT = ClaudeSDKClient(options=OPENALEX_MCP_OPTIONS)
