from __future__ import annotations

import argparse
import asyncio
import os
import re
from dataclasses import dataclass
from datetime import datetime
from itertools import combinations
from pathlib import Path
from typing import Literal, cast

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["openalex"]


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, ...] = ("openalex",)
DEFAULT_DOWNLOAD_PATH = _first_env("ACADEMIC_MCP_DOWNLOAD_PATH") or "./downloads"

OPENALEX_API_KEY = _first_env("OPENALEX_API_KEY") or ""
OPENALEX_MAILTO = _first_env("OPENALEX_MAILTO") or ""

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

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 _count_or_zero(value: object) -> int:
    count = _safe_int(value)
    return count if count is not None else 0


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 = {"openalex"}
    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()}..."


def _decode_openalex_abstract(inverted_index: object) -> str:
    if not isinstance(inverted_index, dict) or not inverted_index:
        return ""

    pairs: list[tuple[int, str]] = []
    for token, positions in inverted_index.items():
        if not isinstance(token, str) or not isinstance(positions, list):
            continue
        for pos in positions:
            idx = _safe_int(pos)
            if idx is not None and idx >= 0:
                pairs.append((idx, token))

    if not pairs:
        return ""

    pairs.sort(key=lambda item: item[0])
    words = [word for _, word in pairs]
    return _clean_text(" ".join(words))


def _openalex_authorships_to_authors(authorships: object) -> list[str]:
    if not isinstance(authorships, list):
        return []
    result: list[str] = []
    for item in authorships:
        if not isinstance(item, dict):
            continue
        author = item.get("author")
        if not isinstance(author, dict):
            continue
        name = _clean_text(str(author.get("display_name", "")))
        if name:
            result.append(name)
    return result


async def _search_openalex(client: httpx.AsyncClient, req: SearchRequest) -> list[dict[str, object]]:
    if not OPENALEX_API_KEY:
        raise ValueError("OPENALEX_API_KEY is required.")

    params: dict[str, str | int] = {
        "search": req.query,
        "per-page": max(1, min(req.max_results, 200)),
        "sort": "cited_by_count:desc",
        "api_key": OPENALEX_API_KEY,
    }

    filters: list[str] = []
    if req.year_from is not None and req.year_to is not None:
        filters.append(f"publication_year:{req.year_from}-{req.year_to}")
    elif req.year_from is not None:
        filters.append(f"publication_year:>{req.year_from - 1}")
    elif req.year_to is not None:
        filters.append(f"publication_year:<{req.year_to + 1}")

    if filters:
        params["filter"] = ",".join(filters)
    if OPENALEX_MAILTO:
        params["mailto"] = OPENALEX_MAILTO

    query_params = {key: str(value) for key, value in params.items()}
    response = await client.get("https://api.openalex.org/works", params=query_params)
    response.raise_for_status()
    payload = response.json()
    results = payload.get("results", [])
    papers: list[dict[str, object]] = []

    for item in results:
        if not isinstance(item, dict):
            continue

        paper_year = _safe_int(item.get("publication_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 = _clean_text(str(item.get("doi", "") or ""))
        openalex_id = _clean_text(str(item.get("id", "") or ""))

        primary_location = item.get("primary_location")
        landing_page = ""
        if isinstance(primary_location, dict):
            landing_page = _clean_text(str(primary_location.get("landing_page_url", "") or ""))

        url = doi or landing_page or openalex_id
        abstract = _decode_openalex_abstract(item.get("abstract_inverted_index"))

        papers.append(
            {
                "source": "openalex",
                "id": openalex_id,
                "title": _clean_text(str(item.get("title", "") or item.get("display_name", ""))),
                "abstract": abstract,
                "authors": _openalex_authorships_to_authors(item.get("authorships")),
                "year": paper_year,
                "published": _clean_text(str(item.get("publication_date", "") or "")),
                "url": url,
                "doi": doi,
                "citation_count": _safe_int(item.get("cited_by_count")),
            }
        )

    return papers


async def _search_one(client: httpx.AsyncClient, req: SearchRequest) -> tuple[str, list[dict[str, object]], str | None]:
    try:
        if req.source == "openalex":
            return req.source, await _search_openalex(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 _paper_brief(paper: dict[str, object], rank: int) -> dict[str, object]:
    return {
        "rank": rank,
        "title": str(paper.get("title", "")),
        "year": _safe_int(paper.get("year")),
        "url": str(paper.get("url", "")),
    }


def _unique_authors(paper: dict[str, object]) -> list[str]:
    raw = paper.get("authors", [])
    if not isinstance(raw, list):
        return []

    names: list[str] = []
    seen: set[str] = set()
    for item in raw:
        if not isinstance(item, str):
            continue
        name = _clean_text(item)
        if not name or name in seen:
            continue
        seen.add(name)
        names.append(name)
    return names


def _summarize_papers(papers: list[dict[str, object]]) -> list[dict[str, object]]:
    summarized: list[dict[str, object]] = []
    for index, paper in enumerate(papers, start=1):
        authors = _unique_authors(paper)
        summarized.append(
            {
                "rank": index,
                "title": str(paper.get("title", "")),
                "year": _safe_int(paper.get("year")),
                "citation_count": _safe_int(paper.get("citation_count")),
                "authors_total": len(authors),
                "url": str(paper.get("url", "")),
            }
        )
    return summarized


def _analyze_author_network(
    papers: list[dict[str, object]],
    min_author_paper_count: int,
    max_authors_per_paper: int,
    top_coauthor_pairs: int,
) -> dict[str, object]:
    author_to_papers: dict[str, list[dict[str, object]]] = {}
    coauthor_counts: dict[tuple[str, str], int] = {}
    coauthor_samples: dict[tuple[str, str], list[dict[str, object]]] = {}

    papers_used_for_pairs = 0
    papers_skipped_for_pairs = 0

    for index, paper in enumerate(papers, start=1):
        authors = _unique_authors(paper)
        brief = _paper_brief(paper, index)

        for author in authors:
            author_to_papers.setdefault(author, []).append(brief)

        if len(authors) < 2:
            continue
        if len(authors) > max_authors_per_paper:
            papers_skipped_for_pairs += 1
            continue

        papers_used_for_pairs += 1

        for author_a, author_b in combinations(authors, 2):
            pair = (author_a, author_b) if author_a <= author_b else (author_b, author_a)
            coauthor_counts[pair] = coauthor_counts.get(pair, 0) + 1

            sample = coauthor_samples.setdefault(pair, [])
            if len(sample) < 5:
                sample.append(brief)

    recurring_authors: list[dict[str, object]] = []
    for author, authored_papers in author_to_papers.items():
        if len(authored_papers) < min_author_paper_count:
            continue
        recurring_authors.append(
            {
                "author": author,
                "paper_count": len(authored_papers),
                "papers": authored_papers,
            }
        )

    recurring_authors.sort(
        key=lambda item: (-_count_or_zero(item.get("paper_count")), str(item.get("author", "")))
    )

    top_authors = [
        {
            "author": author,
            "paper_count": len(authored_papers),
        }
        for author, authored_papers in sorted(
            author_to_papers.items(),
            key=lambda item: (-len(item[1]), item[0]),
        )
    ]

    coauthor_pairs: list[dict[str, object]] = []
    ranked_pairs = sorted(
        coauthor_counts.items(),
        key=lambda item: (-item[1], item[0][0], item[0][1]),
    )
    for (author_a, author_b), count in ranked_pairs[:top_coauthor_pairs]:
        coauthor_pairs.append(
            {
                "author_a": author_a,
                "author_b": author_b,
                "coauthored_paper_count": count,
                "sample_papers": coauthor_samples.get((author_a, author_b), []),
            }
        )

    return {
        "recurring_authors": recurring_authors,
        "coauthor_pairs": coauthor_pairs,
        "top_authors": top_authors,
        "papers_used_for_pairs": papers_used_for_pairs,
        "papers_skipped_for_pairs": papers_skipped_for_pairs,
    }


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("- 数据源：OpenAlex")
    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("- 本综述基于 OpenAlex 检索结果自动生成，正式引用前请核验原文。")
    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 {
        "openalex_api_key_configured": bool(OPENALEX_API_KEY),
        "openalex_mailto": OPENALEX_MAILTO,
        "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,
    }


@mcp.tool()
async def analyze_author_network(
    query: str,
    sources: list[str] | None = None,
    max_results_per_source: int = 20,
    year_from: int | None = None,
    year_to: int | None = None,
    min_author_paper_count: int = 2,
    max_authors_per_paper: int = 20,
    top_coauthor_pairs: int = 30,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
) -> dict[str, object]:
    min_paper_count = max(2, min(min_author_paper_count, 50))
    max_authors = max(2, min(max_authors_per_paper, 500))
    top_pairs = max(1, min(top_coauthor_pairs, 500))

    search_result = await paper_search(
        query=query,
        sources=sources,
        max_results_per_source=max_results_per_source,
        year_from=year_from,
        year_to=year_to,
        timeout_seconds=timeout_seconds,
    )

    papers_raw = search_result.get("papers", [])
    papers: list[dict[str, object]] = []
    if isinstance(papers_raw, list):
        papers = [paper for paper in papers_raw if isinstance(paper, dict)]

    network = _analyze_author_network(
        papers=papers,
        min_author_paper_count=min_paper_count,
        max_authors_per_paper=max_authors,
        top_coauthor_pairs=top_pairs,
    )

    errors = search_result.get("errors", [])
    if not isinstance(errors, list):
        errors = []

    return {
        "success": True,
        "query": query,
        "sources": _normalize_sources(sources),
        "year_from": year_from,
        "year_to": year_to,
        "total_papers": len(papers),
        "paper_list": _summarize_papers(papers),
        "min_author_paper_count": min_paper_count,
        "max_authors_per_paper": max_authors,
        "top_coauthor_pairs": top_pairs,
        "recurring_authors": network["recurring_authors"],
        "coauthor_pairs": network["coauthor_pairs"],
        "top_authors": network["top_authors"],
        "papers_used_for_pairs": network["papers_used_for_pairs"],
        "papers_skipped_for_pairs": network["papers_skipped_for_pairs"],
        "errors": errors,
    }


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())
