import contextlib
import json
import os
import tempfile
import time
import urllib.error
import urllib.request
from typing import Any, Dict, List

import gradio as gr
import uvicorn
from fastapi import FastAPI

from scholarfetch_cli import UnifiedRecord
from scholarfetch_fastmcp import build_server
from scholarfetch_mcp import ScholarFetchService


SERVICE = ScholarFetchService()
ALLOWED_ENGINES = ["elsevier", "openalex", "crossref", "arxiv", "europepmc", "springer", "semanticscholar"]
PAPER_HEADERS = ["idx", "title", "doi", "year", "authors", "venue", "engine", "abstract", "fulltext", "url", "pdf_url"]
REF_HEADERS = ["idx", "reference", "doi"]
AUTHOR_HEADERS = ["idx", "display_name", "author_id", "affiliation", "works_count", "cited_by_count"]


def _public_base_url():
    host = (os.getenv("SPACE_HOST") or "").strip()
    if host:
        return f"https://{host}"
    return f"http://127.0.0.1:{int(os.getenv('PORT', '7860'))}"


def _public_mcp_url():
    return _public_base_url().rstrip("/") + "/mcp"


def _parse_engine_list(raw):
    if isinstance(raw, list):
        selected = [x for x in raw if x in ALLOWED_ENGINES]
    else:
        selected = []
    return selected or ["elsevier", "openalex", "crossref", "arxiv", "europepmc", "springer"]


def _cli(engines):
    return SERVICE._runtime(_parse_engine_list(engines))


def _record_from_payload(payload: Dict[str, Any]) -> UnifiedRecord:
    return UnifiedRecord(
        engine=str(payload.get("engine") or ""),
        title=str(payload.get("title") or ""),
        doi=str(payload.get("doi") or ""),
        year=str(payload.get("year") or ""),
        authors=str(payload.get("authors") or ""),
        venue=str(payload.get("venue") or ""),
        abstract=str(payload.get("abstract") or ""),
        url=str(payload.get("url") or ""),
        pdf_url=str(payload.get("pdf_url") or ""),
        raw_id=str(payload.get("raw_id") or payload.get("doi") or payload.get("title") or ""),
    )


def _paper_to_row(i: int, payload: Dict[str, Any], cli) -> List[Any]:
    rec = _record_from_payload(payload)
    return [
        i,
        rec.title,
        rec.doi,
        rec.year,
        rec.authors,
        rec.venue,
        rec.engine,
        "YES" if (rec.abstract or "").strip() else "NO",
        cli._record_fulltext_status(rec).upper(),
        rec.url,
        rec.pdf_url,
    ]


def _rows_from_papers(results: List[Dict[str, Any]], engines) -> List[List[Any]]:
    cli = _cli(engines)
    return [_paper_to_row(i + 1, payload, cli) for i, payload in enumerate(results)]


def _ensure_index(index, size: int) -> int:
    idx = int(index)
    if idx < 1 or idx > size:
        raise ValueError(f"Index out of range: 1..{size}")
    return idx - 1


def check_mcp_live():
    url = _public_mcp_url().rstrip("/") + "/"
    started = time.time()
    req = urllib.request.Request(url, method="GET")
    try:
        with urllib.request.urlopen(req, timeout=6) as resp:
            ms = int((time.time() - started) * 1000)
            return f"MCP is reachable: HTTP {resp.status} in {ms} ms\n{url}"
    except urllib.error.HTTPError as err:
        ms = int((time.time() - started) * 1000)
        if err.code != 404:
            return f"MCP endpoint reachable (HTTP {err.code}) in {ms} ms\n{url}"
        return f"MCP check failed: HTTP 404 in {ms} ms\n{url}"
    except Exception as exc:
        ms = int((time.time() - started) * 1000)
        return f"MCP check failed: {exc} in {ms} ms\n{url}"


def run_search(query, limit, engines):
    q = (query or "").strip()
    if not q:
        return [], "Insert a query.", []
    out = SERVICE.search({"query": q, "limit": int(limit), "engines": _parse_engine_list(engines)})
    rows = _rows_from_papers(out.get("results", []), engines)
    msg = f"Found {out.get('count', 0)} results using: {', '.join(out.get('engines_used', _parse_engine_list(engines)))}"
    return rows, msg, out.get("results", [])


def run_author_candidates(name, limit, engines):
    author_name = (name or "").strip()
    if not author_name:
        return [], "Insert an author name.", []
    out = SERVICE.author_candidates({"name": author_name, "limit": int(limit), "engines": _parse_engine_list(engines)})
    candidates = out.get("candidates", [])
    rows = []
    for i, cand in enumerate(candidates, start=1):
        rows.append([
            i,
            cand.get("display_name", ""),
            cand.get("author_id", ""),
            cand.get("affiliation", ""),
            cand.get("works_count", 0),
            cand.get("cited_by_count", 0),
        ])
    return rows, f"Found {len(candidates)} author candidates.", candidates


def run_author_papers(name, candidate_index, limit, filters, engines):
    author_name = (name or "").strip()
    if not author_name:
        return [], "Insert an author name first.", []
    out = SERVICE.author_papers(
        {
            "author_name": author_name,
            "candidate_index": int(candidate_index),
            "limit": int(limit),
            "filters": filters or "",
            "engines": _parse_engine_list(engines),
        }
    )
    rows = _rows_from_papers(out.get("results", []), engines)
    author_meta = out.get("author", {})
    label = author_meta.get("display_name") or author_meta.get("author_id") or author_name
    return rows, f"Loaded {out.get('count', 0)} papers for {label}.", out.get("results", [])


def get_abstract_from_state(index, paper_state, engines):
    if not paper_state:
        return "No active paper list. Search or load author papers first."
    try:
        payload = paper_state[_ensure_index(index, len(paper_state))]
    except Exception as exc:
        return str(exc)
    out = SERVICE.abstract({"doi": payload.get("doi"), "engines": _parse_engine_list(engines)})
    if out.get("abstract"):
        return out.get("abstract")
    return "Abstract not available for the selected paper."


def get_fulltext_from_state(index, paper_state, engines):
    if not paper_state:
        return "No active paper list. Search or load author papers first."
    try:
        payload = paper_state[_ensure_index(index, len(paper_state))]
    except Exception as exc:
        return str(exc)
    rec = _record_from_payload(payload)
    if rec.doi:
        out = SERVICE.article_text({"doi": rec.doi, "engines": _parse_engine_list(engines)})
    elif rec.engine == "arxiv":
        out = _cli(engines)._resolve_fulltext("", seed_record=rec)
        if out.get("found"):
            return out.get("text") or ""
        return "Full text not available for the selected paper."
    else:
        return "Full text not available for the selected paper."
    if out.get("article_text"):
        return out.get("article_text")
    return "Full text not available for the selected paper."


def load_references_from_state(index, paper_state, engines):
    if not paper_state:
        return [], "No active paper list. Search or load author papers first.", []
    try:
        payload = paper_state[_ensure_index(index, len(paper_state))]
    except Exception as exc:
        return [], str(exc), []
    doi = str(payload.get("doi") or "").strip()
    if not doi:
        return [], "Selected paper has no DOI. References are unavailable.", []
    out = SERVICE.references({"doi": doi, "engines": _parse_engine_list(engines)})
    refs = out.get("references", [])
    rows = [[i + 1, ref.get("text", ""), ref.get("doi", "")] for i, ref in enumerate(refs)]
    return rows, f"Loaded {len(refs)} references from {out.get('source') or 'available sources'}.", refs


def resolve_reference(ref_index, ref_state, engines):
    if not ref_state:
        return [], "Load references first.", []
    try:
        ref = ref_state[_ensure_index(ref_index, len(ref_state))]
    except Exception as exc:
        return [], str(exc), []
    cli = _cli(engines)
    payload = None
    doi = str(ref.get("doi") or "").strip()
    if doi:
        payload = SERVICE._select_best_record(cli, cli._parallel_doi_lookup(doi))
    if not payload:
        payload = SERVICE._select_best_record(cli, cli._parallel_search(str(ref.get("text") or ""), limit_per_engine=3))
    if not payload:
        return [], "Could not resolve this reference into a paper record.", []
    rows = _rows_from_papers([payload], engines)
    return rows, "Resolved reference into a paper record.", [payload]


def add_paper_to_saved(index, collection, paper_state):
    if not paper_state:
        return "No active paper list to save from."
    try:
        payload = paper_state[_ensure_index(index, len(paper_state))]
    except Exception as exc:
        return str(exc)
    out = SERVICE.saved_add({"collection": collection or "default", "paper_json": json.dumps(payload, ensure_ascii=False)})
    if out.get("added"):
        return f"Saved paper to collection '{out['collection']}'. Count: {out['count']}"
    return f"Paper already present in collection '{out['collection']}'. Count: {out['count']}"


def list_saved(collection, engines):
    out = SERVICE.saved_list({"collection": collection or "default"})
    rows = _rows_from_papers(out.get("results", []), engines)
    return rows, f"Collection '{out['collection']}' contains {out['count']} saved papers.", out.get("results", [])


def remove_saved(index, collection, saved_state):
    if not saved_state:
        return [], "Saved list is empty.", []
    try:
        payload = saved_state[_ensure_index(index, len(saved_state))]
    except Exception as exc:
        return [], str(exc), saved_state
    out = SERVICE.saved_remove({"collection": collection or "default", "doi": payload.get("doi"), "title": payload.get("title")})
    refreshed = SERVICE.saved_list({"collection": collection or "default"})
    rows = _rows_from_papers(refreshed.get("results", []), ALLOWED_ENGINES)
    msg = f"Removed from '{out['collection']}'. Count: {out['count']}" if out.get("removed") else f"Paper not found in '{out['collection']}'."
    return rows, msg, refreshed.get("results", [])


def clear_saved(collection):
    out = SERVICE.saved_clear({"collection": collection or "default"})
    return [], f"Cleared collection '{out['collection']}'. Removed {out['removed_count']} papers.", []


def export_saved(collection, export_format, citation_style, include_references, engines):
    out = SERVICE.saved_export(
        {
            "collection": collection or "default",
            "format": export_format,
            "style": citation_style,
            "include_references": bool(include_references),
            "engines": _parse_engine_list(engines),
        }
    )
    suffix = ".bib" if export_format == "bib" else ".txt"
    with tempfile.NamedTemporaryFile("w", suffix=suffix, delete=False, encoding="utf-8") as fh:
        fh.write(out.get("content") or "")
        path = fh.name
    note = f"Exported {out.get('count', 0)} saved papers from '{out.get('collection')}' as {out.get('format')}."
    return path, note


with gr.Blocks(title="ScholarFetch Web") as demo:
    gr.Markdown("# ScholarFetch Web")
    gr.Markdown(
        "Search, traverse, read, save, and export scholarly papers across multiple engines. "
        "This UI now exposes the same stateful reading-list workflow used by the MCP server."
    )
    gr.Markdown(f"### Public MCP endpoint\n`{_public_mcp_url()}`")

    with gr.Row():
        engines = gr.CheckboxGroup(
            choices=ALLOWED_ENGINES,
            value=["elsevier", "openalex", "crossref", "arxiv", "europepmc", "springer"],
            label="Engines",
        )
        collection = gr.Textbox(label="Saved collection", value="default")

    mcp_check_btn = gr.Button("Check MCP Live")
    mcp_check_out = gr.Textbox(label="MCP status", lines=3, interactive=False)
    mcp_check_btn.click(fn=check_mcp_live, inputs=None, outputs=mcp_check_out, api_name=False)

    search_state = gr.State([])
    author_candidates_state = gr.State([])
    author_papers_state = gr.State([])
    refs_state = gr.State([])
    resolved_ref_paper_state = gr.State([])
    saved_state = gr.State([])

    with gr.Tabs():
        with gr.Tab("Search"):
            query = gr.Textbox(label="Query", placeholder="keywords, DOI, or author name")
            limit = gr.Slider(minimum=1, maximum=100, value=25, step=1, label="Max results")
            search_btn = gr.Button("Search")
            search_status = gr.Markdown()
            search_table = gr.Dataframe(headers=PAPER_HEADERS, datatype=["number"] + ["str"] * 10, interactive=False, label="Search results")
            with gr.Row():
                search_index = gr.Number(label="Selected result index", value=1, precision=0)
                add_search_saved_btn = gr.Button("Save selected paper")
            with gr.Row():
                search_abstract_btn = gr.Button("Show Abstract")
                search_text_btn = gr.Button("Show Full Text")
                search_refs_btn = gr.Button("Load References")
            search_abstract = gr.Textbox(label="Abstract", lines=12)
            search_fulltext = gr.Textbox(label="Full text", lines=20)
            search_save_status = gr.Markdown()
            search_refs_table = gr.Dataframe(headers=REF_HEADERS, datatype=["number", "str", "str"], interactive=False, label="References")
            search_refs_status = gr.Markdown()
            with gr.Row():
                ref_index = gr.Number(label="Reference index", value=1, precision=0)
                resolve_ref_btn = gr.Button("Resolve reference")
                save_ref_btn = gr.Button("Save resolved reference")
            resolved_ref_table = gr.Dataframe(headers=PAPER_HEADERS, datatype=["number"] + ["str"] * 10, interactive=False, label="Resolved reference paper")
            resolved_ref_status = gr.Markdown()

            search_btn.click(fn=run_search, inputs=[query, limit, engines], outputs=[search_table, search_status, search_state], api_name=False)
            add_search_saved_btn.click(fn=add_paper_to_saved, inputs=[search_index, collection, search_state], outputs=search_save_status, api_name=False)
            search_abstract_btn.click(fn=get_abstract_from_state, inputs=[search_index, search_state, engines], outputs=search_abstract, api_name=False)
            search_text_btn.click(fn=get_fulltext_from_state, inputs=[search_index, search_state, engines], outputs=search_fulltext, api_name=False)
            search_refs_btn.click(fn=load_references_from_state, inputs=[search_index, search_state, engines], outputs=[search_refs_table, search_refs_status, refs_state], api_name=False)
            resolve_ref_btn.click(fn=resolve_reference, inputs=[ref_index, refs_state, engines], outputs=[resolved_ref_table, resolved_ref_status, resolved_ref_paper_state], api_name=False)
            save_ref_btn.click(fn=add_paper_to_saved, inputs=[gr.Number(value=1, visible=False), collection, resolved_ref_paper_state], outputs=search_save_status, api_name=False)

        with gr.Tab("Authors"):
            author_name = gr.Textbox(label="Author name", placeholder="Albert Einstein")
            candidate_limit = gr.Slider(minimum=1, maximum=20, value=8, step=1, label="Max author candidates")
            author_candidates_btn = gr.Button("Find author candidates")
            author_candidates_status = gr.Markdown()
            author_candidates_table = gr.Dataframe(headers=AUTHOR_HEADERS, datatype=["number", "str", "str", "str", "number", "number"], interactive=False, label="Author candidates")
            candidate_index = gr.Number(label="Candidate index", value=1, precision=0)
            author_papers_limit = gr.Slider(minimum=1, maximum=200, value=50, step=1, label="Max papers")
            author_filters = gr.Textbox(label="Paper filters", placeholder="year>=2020,has:abstract")
            author_papers_btn = gr.Button("Load papers for selected candidate")
            author_papers_status = gr.Markdown()
            author_papers_table = gr.Dataframe(headers=PAPER_HEADERS, datatype=["number"] + ["str"] * 10, interactive=False, label="Author papers")
            with gr.Row():
                author_paper_index = gr.Number(label="Selected paper index", value=1, precision=0)
                author_save_btn = gr.Button("Save selected paper")
            with gr.Row():
                author_abstract_btn = gr.Button("Show Abstract")
                author_text_btn = gr.Button("Show Full Text")
                author_refs_btn = gr.Button("Load References")
            author_abstract = gr.Textbox(label="Abstract", lines=12)
            author_fulltext = gr.Textbox(label="Full text", lines=20)
            author_refs_table = gr.Dataframe(headers=REF_HEADERS, datatype=["number", "str", "str"], interactive=False, label="References")
            author_save_status = gr.Markdown()
            author_refs_status = gr.Markdown()

            author_candidates_btn.click(fn=run_author_candidates, inputs=[author_name, candidate_limit, engines], outputs=[author_candidates_table, author_candidates_status, author_candidates_state], api_name=False)
            author_papers_btn.click(fn=run_author_papers, inputs=[author_name, candidate_index, author_papers_limit, author_filters, engines], outputs=[author_papers_table, author_papers_status, author_papers_state], api_name=False)
            author_save_btn.click(fn=add_paper_to_saved, inputs=[author_paper_index, collection, author_papers_state], outputs=author_save_status, api_name=False)
            author_abstract_btn.click(fn=get_abstract_from_state, inputs=[author_paper_index, author_papers_state, engines], outputs=author_abstract, api_name=False)
            author_text_btn.click(fn=get_fulltext_from_state, inputs=[author_paper_index, author_papers_state, engines], outputs=author_fulltext, api_name=False)
            author_refs_btn.click(fn=load_references_from_state, inputs=[author_paper_index, author_papers_state, engines], outputs=[author_refs_table, author_refs_status, refs_state], api_name=False)

        with gr.Tab("Saved"):
            saved_refresh_btn = gr.Button("Refresh saved papers")
            saved_status = gr.Markdown()
            saved_table = gr.Dataframe(headers=PAPER_HEADERS, datatype=["number"] + ["str"] * 10, interactive=False, label="Saved papers")
            with gr.Row():
                saved_index = gr.Number(label="Saved paper index", value=1, precision=0)
                saved_remove_btn = gr.Button("Remove selected")
                saved_clear_btn = gr.Button("Clear collection")
            with gr.Row():
                export_format = gr.Dropdown(choices=["citations", "abstracts", "bib", "fulltext"], value="citations", label="Export format")
                citation_style = gr.Dropdown(choices=["harvard", "apa", "ieee"], value="harvard", label="Citation style")
                include_references = gr.Checkbox(label="Include references in fulltext export", value=False)
            saved_export_btn = gr.Button("Export saved collection")
            exported_file = gr.File(label="Export file")
            export_status = gr.Markdown()

            saved_refresh_btn.click(fn=list_saved, inputs=[collection, engines], outputs=[saved_table, saved_status, saved_state], api_name=False)
            saved_remove_btn.click(fn=remove_saved, inputs=[saved_index, collection, saved_state], outputs=[saved_table, saved_status, saved_state], api_name=False)
            saved_clear_btn.click(fn=clear_saved, inputs=[collection], outputs=[saved_table, saved_status, saved_state], api_name=False)
            saved_export_btn.click(fn=export_saved, inputs=[collection, export_format, citation_style, include_references, engines], outputs=[exported_file, export_status], api_name=False)


fastmcp_server = build_server(host="0.0.0.0", port=int(os.getenv("PORT", "7860")), streamable_http_path="/")
fastmcp_app = fastmcp_server.streamable_http_app()


@contextlib.asynccontextmanager
async def lifespan(_app):
    async with fastmcp_server.session_manager.run():
        yield


app = FastAPI(title="ScholarFetch Space", docs_url=None, redoc_url=None, lifespan=lifespan)
app.mount("/mcp", fastmcp_app)
app = gr.mount_gradio_app(app, demo, path="/")


@app.middleware("http")
async def normalize_mcp_path(request, call_next):
    if request.scope.get("path") == "/mcp":
        request.scope["path"] = "/mcp/"
    return await call_next(request)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))