from __future__ import annotations

import argparse
import asyncio
import os
import re
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Literal, cast
from xml.etree import ElementTree

import httpx
from mcp.server.fastmcp import FastMCP
from starlette.responses import PlainTextResponse

try:
    from dotenv import load_dotenv

    _ = load_dotenv()
except Exception:
    pass


SupportedSource = Literal["arxiv", "crossref", "semantic", "core"]


def _first_env(*keys: str) -> str | None:
    for key in keys:
        value = os.getenv(key)
        if value:
            return value
    return None


DEFAULT_TIMEOUT_SECONDS = int(_first_env("ACADEMIC_HTTP_TIMEOUT", "HTTP_TIMEOUT_SECONDS") or "25")
DEFAULT_SOURCES: tuple[SupportedSource, ...] = ("arxiv", "crossref", "semantic")
DEFAULT_DOWNLOAD_PATH = _first_env("ACADEMIC_MCP_DOWNLOAD_PATH") or "./downloads"

SEMANTIC_SCHOLAR_API_KEY = _first_env("SEMANTIC_SCHOLAR_API_KEY") or ""
CORE_API_KEY = _first_env("CORE_API_KEY") or ""

mcp = FastMCP(
    name="AcademicMCP",
    instructions="Academic paper search and review generation over MCP",
)

app = mcp.streamable_http_app()


@dataclass(frozen=True)
class SearchRequest:
    query: str
    source: SupportedSource
    max_results: int
    year_from: int | None
    year_to: int | None


def _safe_int(value: object) -> int | None:
    try:
        if value is None:
            return None
        if isinstance(value, bool):
            return int(value)
        if isinstance(value, (int, float, str)):
            return int(value)
        return None
    except Exception:
        return None


def _normalize_sources(sources: list[str] | None) -> list[SupportedSource]:
    if not sources:
        return [cast(SupportedSource, value) for value in DEFAULT_SOURCES]

    normalized: list[SupportedSource] = []
    allowed = {"arxiv", "crossref", "semantic", "core"}
    for source in sources:
        item = source.strip().lower()
        if item in allowed:
            normalized.append(cast(SupportedSource, item))

    if not normalized:
        return [cast(SupportedSource, value) for value in DEFAULT_SOURCES]

    return [cast(SupportedSource, value) for value in dict.fromkeys(normalized)]


def _clean_text(value: str) -> str:
    return re.sub(r"\s+", " ", value).strip()


def _paper_key(paper: dict[str, object]) -> str:
    title = str(paper.get("title", "")).lower().strip()
    doi = str(paper.get("doi", "")).lower().strip()
    if doi:
        return f"doi:{doi}"
    return f"title:{title}"


def _snippet(text: str, limit: int = 220) -> str:
    content = _clean_text(text)
    if len(content) <= limit:
        return content
    return f"{content[:limit].rstrip()}..."


async def _search_arxiv(client: httpx.AsyncClient, req: SearchRequest) -> list[dict[str, object]]:
    params = {
        "search_query": f"all:{req.query}",
        "start": 0,
        "max_results": req.max_results,
    }
    response = await client.get("http://export.arxiv.org/api/query", params=params)
    response.raise_for_status()

    root = ElementTree.fromstring(response.text)
    ns = {"atom": "http://www.w3.org/2005/Atom"}
    entries = root.findall("atom:entry", ns)
    papers: list[dict[str, object]] = []

    for entry in entries:
        title = _clean_text(entry.findtext("atom:title", default="", namespaces=ns))
        abstract = _clean_text(entry.findtext("atom:summary", default="", namespaces=ns))
        published = entry.findtext("atom:published", default="", namespaces=ns)
        paper_year = _safe_int((published or "")[:4])
        if req.year_from and paper_year and paper_year < req.year_from:
            continue
        if req.year_to and paper_year and paper_year > req.year_to:
            continue

        authors = [
            _clean_text(node.text or "")
            for node in entry.findall("atom:author/atom:name", ns)
            if node.text
        ]
        paper_id = _clean_text(entry.findtext("atom:id", default="", namespaces=ns))

        papers.append(
            {
                "source": "arxiv",
                "id": paper_id,
                "title": title,
                "abstract": abstract,
                "authors": authors,
                "year": paper_year,
                "published": published,
                "url": paper_id,
                "doi": "",
                "citation_count": None,
            }
        )

    return papers


async def _search_crossref(client: httpx.AsyncClient, req: SearchRequest) -> list[dict[str, object]]:
    params = {
        "query": req.query,
        "rows": req.max_results,
        "select": "DOI,title,author,issued,URL,abstract,is-referenced-by-count,published-print,published-online",
    }
    response = await client.get("https://api.crossref.org/works", params=params)
    response.raise_for_status()
    payload = response.json()
    items = payload.get("message", {}).get("items", [])
    papers: list[dict[str, object]] = []

    for item in items:
        title_list = item.get("title", [])
        title = _clean_text(title_list[0]) if title_list else ""

        date_parts = (
            item.get("issued", {}).get("date-parts", [[]])[0]
            or item.get("published-print", {}).get("date-parts", [[]])[0]
            or item.get("published-online", {}).get("date-parts", [[]])[0]
        )
        paper_year = _safe_int(date_parts[0] if date_parts else None)
        if req.year_from and paper_year and paper_year < req.year_from:
            continue
        if req.year_to and paper_year and paper_year > req.year_to:
            continue

        authors: list[str] = []
        for author in item.get("author", [])[:15]:
            given = _clean_text(author.get("given", ""))
            family = _clean_text(author.get("family", ""))
            full_name = _clean_text(f"{given} {family}")
            if full_name:
                authors.append(full_name)

        abstract_raw = item.get("abstract") or ""
        abstract = _clean_text(re.sub(r"<[^>]+>", " ", abstract_raw)) if abstract_raw else ""

        papers.append(
            {
                "source": "crossref",
                "id": _clean_text(item.get("DOI", "")),
                "title": title,
                "abstract": abstract,
                "authors": authors,
                "year": paper_year,
                "published": "",
                "url": _clean_text(item.get("URL", "")),
                "doi": _clean_text(item.get("DOI", "")),
                "citation_count": _safe_int(item.get("is-referenced-by-count")),
            }
        )

    return papers


async def _search_semantic(client: httpx.AsyncClient, req: SearchRequest) -> list[dict[str, object]]:
    headers: dict[str, str] = {}
    if SEMANTIC_SCHOLAR_API_KEY:
        headers["x-api-key"] = SEMANTIC_SCHOLAR_API_KEY

    params = {
        "query": req.query,
        "limit": req.max_results,
        "fields": "title,abstract,authors,year,url,externalIds,citationCount",
    }
    response = await client.get(
        "https://api.semanticscholar.org/graph/v1/paper/search",
        params=params,
        headers=headers,
    )
    response.raise_for_status()
    payload = response.json()
    data = payload.get("data", [])
    papers: list[dict[str, object]] = []

    for item in data:
        paper_year = _safe_int(item.get("year"))
        if req.year_from and paper_year and paper_year < req.year_from:
            continue
        if req.year_to and paper_year and paper_year > req.year_to:
            continue

        doi = ""
        external_ids = item.get("externalIds") or {}
        if isinstance(external_ids, dict):
            doi = _clean_text(external_ids.get("DOI", ""))

        papers.append(
            {
                "source": "semantic",
                "id": _clean_text(item.get("paperId", "")),
                "title": _clean_text(item.get("title", "")),
                "abstract": _clean_text(item.get("abstract", "")),
                "authors": [
                    _clean_text(author.get("name", ""))
                    for author in item.get("authors", [])
                    if author.get("name")
                ],
                "year": paper_year,
                "published": "",
                "url": _clean_text(item.get("url", "")),
                "doi": doi,
                "citation_count": _safe_int(item.get("citationCount")),
            }
        )

    return papers


async def _search_core(client: httpx.AsyncClient, req: SearchRequest) -> list[dict[str, object]]:
    headers: dict[str, str] = {}
    if CORE_API_KEY:
        headers["Authorization"] = f"Bearer {CORE_API_KEY}"

    params = {
        "q": req.query,
        "limit": req.max_results,
    }
    response = await client.get("https://api.core.ac.uk/v3/search/works", params=params, headers=headers)
    response.raise_for_status()
    payload = response.json()
    results = payload.get("results", [])
    papers: list[dict[str, object]] = []

    for item in results:
        paper_year = _safe_int(item.get("yearPublished"))
        if req.year_from and paper_year and paper_year < req.year_from:
            continue
        if req.year_to and paper_year and paper_year > req.year_to:
            continue

        papers.append(
            {
                "source": "core",
                "id": _clean_text(str(item.get("id", ""))),
                "title": _clean_text(item.get("title", "")),
                "abstract": _clean_text(item.get("abstract", "")),
                "authors": [_clean_text(name) for name in item.get("authors", []) if name],
                "year": paper_year,
                "published": "",
                "url": _clean_text(item.get("downloadUrl", "") or item.get("sourceFulltextUrls", [""])[0]),
                "doi": _clean_text(item.get("doi", "")),
                "citation_count": _safe_int(item.get("citationCount")),
            }
        )

    return papers


async def _search_one(client: httpx.AsyncClient, req: SearchRequest) -> tuple[str, list[dict[str, object]], str | None]:
    try:
        if req.source == "arxiv":
            return req.source, await _search_arxiv(client, req), None
        if req.source == "crossref":
            return req.source, await _search_crossref(client, req), None
        if req.source == "semantic":
            return req.source, await _search_semantic(client, req), None
        if req.source == "core":
            return req.source, await _search_core(client, req), None
        return req.source, [], f"Unsupported source: {req.source}"
    except Exception as exc:
        return req.source, [], str(exc)


def _dedupe_papers(papers: list[dict[str, object]]) -> list[dict[str, object]]:
    seen: set[str] = set()
    result: list[dict[str, object]] = []
    for paper in papers:
        key = _paper_key(paper)
        if key in seen:
            continue
        seen.add(key)
        result.append(paper)
    return result


def _keyword_frequencies(papers: list[dict[str, object]], top_k: int = 12) -> list[dict[str, object]]:
    stop = {
        "the", "and", "for", "with", "from", "that", "this", "into", "their", "using", "based", "study",
        "paper", "approach", "method", "analysis", "research", "results", "a", "an", "of", "to", "in", "on",
    }
    counts: dict[str, int] = {}
    for paper in papers:
        text = f"{paper.get('title', '')} {paper.get('abstract', '')}".lower()
        words = re.findall(r"[a-z]{4,}", text)
        for word in words:
            if word in stop:
                continue
            counts[word] = counts.get(word, 0) + 1

    ranked = sorted(counts.items(), key=lambda item: (-item[1], item[0]))[:top_k]
    return [{"keyword": key, "count": value} for key, value in ranked]


def _sort_for_review(papers: list[dict[str, object]]) -> list[dict[str, object]]:
    def score(paper: dict[str, object]) -> tuple[int, int]:
        citations = _safe_int(paper.get("citation_count")) or 0
        year = _safe_int(paper.get("year")) or 0
        return citations, year

    return sorted(papers, key=score, reverse=True)


def _render_review_markdown(topic: str, papers: list[dict[str, object]], queries: list[str]) -> str:
    now = datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC")
    sorted_papers = _sort_for_review(papers)
    keywords = _keyword_frequencies(sorted_papers)

    lines: list[str] = []
    lines.append(f"# 论文综述草稿：{topic}")
    lines.append("")
    lines.append(f"生成时间：{now}")
    lines.append("")
    lines.append("## 检索策略")
    for query in queries:
        lines.append(f"- 查询词：`{query}`")
    lines.append(f"- 样本论文数：{len(sorted_papers)}")
    lines.append("")

    lines.append("## 论文样本概览")
    lines.append("| 标题 | 来源 | 年份 | 引用 |")
    lines.append("|---|---|---:|---:|")
    for paper in sorted_papers[:20]:
        title = str(paper.get("title", "")).replace("|", " ")
        source = str(paper.get("source", ""))
        year = str(paper.get("year", "") or "-")
        citation = str(paper.get("citation_count", "") or "-")
        lines.append(f"| {title} | {source} | {year} | {citation} |")
    lines.append("")

    lines.append("## 高频主题词")
    for item in keywords:
        lines.append(f"- {item['keyword']}: {item['count']}")
    lines.append("")

    lines.append("## 关键文献证据")
    for idx, paper in enumerate(sorted_papers[:12], start=1):
        title = str(paper.get("title", ""))
        source = str(paper.get("source", ""))
        paper_id = str(paper.get("id", ""))
        abstract = _snippet(str(paper.get("abstract", "")) or "无摘要")
        url = str(paper.get("url", ""))
        lines.append(f"### [{idx}] {title}")
        lines.append(f"- 引用：`[{source}:{paper_id}]`")
        lines.append(f"- 摘要证据：{abstract}")
        if url:
            lines.append(f"- 链接：{url}")
        lines.append("")

    lines.append("## 初步综述结论（基于以上样本）")
    if keywords:
        major = ", ".join(str(item["keyword"]) for item in keywords[:5])
        lines.append(f"- 当前样本中反复出现的研究焦点包括：{major}。")
    lines.append("- 从可得样本看，研究分布受数据源覆盖影响，建议在正式写作前补充更多数据库交叉验证。")
    lines.append("- 对于关键结论，建议回到原文全文核验后再用于正式论文。")
    lines.append("")

    return "\n".join(lines)


@mcp.custom_route("/health", methods=["GET"])
async def health_check(_request: object | None = None):
    return PlainTextResponse("OK")


@mcp.tool()
def get_academic_config() -> dict[str, object]:
    return {
        "semantic_scholar_api_key_configured": bool(SEMANTIC_SCHOLAR_API_KEY),
        "core_api_key_configured": bool(CORE_API_KEY),
        "default_sources": list(DEFAULT_SOURCES),
        "download_path": DEFAULT_DOWNLOAD_PATH,
    }


@mcp.tool()
async def paper_search(
    query: str,
    sources: list[str] | None = None,
    max_results_per_source: int = 5,
    year_from: int | None = None,
    year_to: int | None = None,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
) -> dict[str, object]:
    selected = _normalize_sources(sources)
    requests = [
        SearchRequest(
            query=query,
            source=source,
            max_results=max(1, min(max_results_per_source, 25)),
            year_from=year_from,
            year_to=year_to,
        )
        for source in selected
    ]

    async with httpx.AsyncClient(timeout=float(timeout_seconds), follow_redirects=True) as client:
        results = await asyncio.gather(*[_search_one(client, req) for req in requests])

    papers: list[dict[str, object]] = []
    errors: list[dict[str, object]] = []
    for source, source_papers, error in results:
        papers.extend(source_papers)
        if error:
            errors.append({"source": source, "error": error})

    deduped = _dedupe_papers(papers)
    return {
        "success": True,
        "query": query,
        "sources": selected,
        "total": len(deduped),
        "papers": deduped,
        "errors": errors,
    }


@mcp.tool()
async def generate_literature_review(
    topic: str,
    queries: list[str] | None = None,
    sources: list[str] | None = None,
    max_results_per_query: int = 4,
    year_from: int | None = None,
    year_to: int | None = None,
    save_to_docs: bool = True,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
) -> dict[str, object]:
    query_list = queries or [topic]
    selected_sources = _normalize_sources(sources)

    merged: list[dict[str, object]] = []
    errors: list[dict[str, object]] = []

    for query in query_list:
        result = await paper_search(
            query=query,
            sources=selected_sources,
            max_results_per_source=max_results_per_query,
            year_from=year_from,
            year_to=year_to,
            timeout_seconds=timeout_seconds,
        )
        papers = result.get("papers", [])
        if isinstance(papers, list):
            merged.extend([paper for paper in papers if isinstance(paper, dict)])
        tool_errors = result.get("errors", [])
        if isinstance(tool_errors, list):
            errors.extend([err for err in tool_errors if isinstance(err, dict)])

    deduped = _dedupe_papers(merged)
    markdown = _render_review_markdown(topic, deduped, query_list)

    output_path = ""
    if save_to_docs:
        docs_dir = Path("docs")
        docs_dir.mkdir(parents=True, exist_ok=True)
        stamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        safe_topic = re.sub(r"[^a-zA-Z0-9\u4e00-\u9fa5_-]+", "_", topic)[:80]
        output = docs_dir / f"review_{safe_topic}_{stamp}.md"
        output.write_text(markdown, encoding="utf-8")
        output_path = str(output.resolve())

    return {
        "success": True,
        "topic": topic,
        "queries": query_list,
        "sources": selected_sources,
        "papers_used": len(deduped),
        "errors": errors,
        "review_markdown": markdown,
        "output_path": output_path,
    }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Academic MCP server launcher")
    _ = parser.add_argument(
        "--transport",
        choices=["stdio", "http", "sse"],
        default=os.getenv("ACADEMIC_MCP_TRANSPORT", os.getenv("MCP_TRANSPORT", "stdio")),
    )
    _ = parser.add_argument(
        "--host",
        default=os.getenv("ACADEMIC_MCP_HOST", os.getenv("MCP_SERVER_HOST", "127.0.0.1")),
    )
    _ = parser.add_argument(
        "--port",
        type=int,
        default=int(os.getenv("ACADEMIC_MCP_PORT", os.getenv("MCP_SERVER_PORT", "8005"))),
    )
    return parser


@dataclass(frozen=True)
class LaunchArgs:
    transport: str
    host: str
    port: int


def _parse_launch_args(argv: list[str] | None = None) -> LaunchArgs:
    parsed = _build_parser().parse_args(argv)
    return LaunchArgs(
        transport=str(getattr(parsed, "transport", "stdio")),
        host=str(getattr(parsed, "host", "127.0.0.1")),
        port=int(getattr(parsed, "port", "8005")),
    )


def main(argv: list[str] | None = None) -> int:
    args = _parse_launch_args(argv)
    if args.transport.lower() in {"http", "sse"}:
        import uvicorn

        uvicorn.run(app, host=args.host, port=args.port, log_level="info")
        return 0

    mcp.run()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
