import os
import json
import re
import time
from datetime import datetime, timedelta
from pathlib import Path

import requests
from flask import Flask, jsonify, render_template, request

app = Flask(__name__)

CACHE_DIR = Path("cache")
CACHE_DIR.mkdir(exist_ok=True)
CACHE_TTL_DAYS = int(os.environ.get("CACHE_TTL_DAYS", "30"))

OA_API = "https://api.openalex.org"

SEARCH_LIMIT = int(os.environ.get("SEARCH_LIMIT", "20"))
REFERENCE_LIMIT = int(os.environ.get("REFERENCE_LIMIT", "100"))

REQUEST_TIMEOUT = 30
API_DELAY_SECONDS = 0.2
MAX_RETRIES = 3
RETRY_BACKOFF = 2  # seconds

OA_EMAIL = os.environ.get("OA_EMAIL", "")
OA_API_KEY = os.environ.get("OA_API_KEY")


def safe_paper_id(paper_id: str) -> str:
    return re.sub(r"[^a-zA-Z0-9_-]", "_", paper_id)[:128]


def cache_path(paper_id: str) -> Path:
    return CACHE_DIR / f"{safe_paper_id(paper_id)}.json"


def is_cache_fresh(path: Path) -> bool:
    if not path.exists():
        return False
    mtime = datetime.fromtimestamp(path.stat().st_mtime)
    return datetime.now() - mtime < timedelta(days=CACHE_TTL_DAYS)


def load_cache(paper_id: str):
    path = cache_path(paper_id)
    if path.exists():
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    return None


def save_cache(paper_id: str, data: dict):
    path = cache_path(paper_id)
    payload = {
        "paper_id": paper_id,
        "created_at": datetime.now().isoformat(),
        "data": data,
    }
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)


def oa_get(url, params=None):
    headers = {
        "User-Agent": "connected-papers-local/1.0 (personal use)",
        "Accept": "application/json",
    }
    params = params or {}
    if OA_EMAIL:
        params["mailto"] = OA_EMAIL
    if OA_API_KEY:
        params["api_key"] = OA_API_KEY

    last_exception = None
    for attempt in range(MAX_RETRIES):
        try:
            r = requests.get(url, params=params, headers=headers, timeout=REQUEST_TIMEOUT)
            r.raise_for_status()
            return r.json()
        except requests.HTTPError as e:
            last_exception = e
            status = getattr(e.response, "status_code", None)
            if status in (429, 503) and attempt < MAX_RETRIES - 1:
                time.sleep(RETRY_BACKOFF * (attempt + 1))
                continue
            raise
        except requests.RequestException as e:
            last_exception = e
            if attempt < MAX_RETRIES - 1:
                time.sleep(RETRY_BACKOFF * (attempt + 1))
                continue
            raise
    raise last_exception


def oa_work_to_paper(work: dict) -> dict:
    authors = [
        {"name": a["author"]["display_name"]}
        for a in work.get("authorships", [])
        if a.get("author")
    ]
    title = work.get("display_name") or work.get("title") or ""
    if not title.strip():
        # fallback to DOI or arXiv ID if title is missing
        ids = work.get("ids", {})
        title = ids.get("arxiv") or ids.get("doi", "").replace("https://doi.org/", "") or "Untitled"
    return {
        "paperId": work.get("id"),
        "title": title,
        "authors": authors,
        "year": work.get("publication_year"),
        "citationCount": work.get("cited_by_count", 0),
        "referenceCount": len(work.get("referenced_works", [])),
    }


def sort_by_citations(papers: list) -> list:
    return sorted(papers, key=lambda p: p.get("citationCount") or 0, reverse=True)


def looks_like_doi(query: str) -> bool:
    return query.startswith("10.") and "/" in query


def looks_like_arxiv(query: str) -> bool:
    q = query.lower()
    if q.startswith("arxiv:") or q.startswith("arXiv:"):
        return True
    import re
    return bool(re.fullmatch(r"\d{4}\.\d{4,5}(v\d+)?", query))


def extract_arxiv_id(query: str) -> str:
    q = query.lower()
    if q.startswith("arxiv:"):
        return query[6:].strip()
    return query.strip()


def resolve_arxiv(arxiv_id: str):
    """Use the arXiv API to get title and authors. Returns dict or None."""
    try:
        url = f"http://export.arxiv.org/api/query?id_list={arxiv_id}"
        r = requests.get(url, timeout=REQUEST_TIMEOUT, allow_redirects=True)
        r.raise_for_status()
        import xml.etree.ElementTree as ET
        root = ET.fromstring(r.text)
        ns = {"atom": "http://www.w3.org/2005/Atom", "arxiv": "http://arxiv.org/schemas/atom"}
        entry = root.find("atom:entry", ns)
        if entry is None:
            return None
        title_elem = entry.find("atom:title", ns)
        title = title_elem.text.strip() if title_elem is not None and title_elem.text else ""
        authors = [
            {"name": author.find("atom:name", ns).text.strip()}
            for author in entry.findall("atom:author", ns)
            if author.find("atom:name", ns) is not None and author.find("atom:name", ns).text
        ]
        return {"title": title, "authors": authors}
    except Exception:
        return None


def lookup_paper(query: str):
    """Try DOI or arXiv direct lookup. Return paper dict or None."""
    if looks_like_doi(query):
        try:
            work = oa_get(f"{OA_API}/works/doi:{query}")
            return oa_work_to_paper(work)
        except requests.RequestException:
            return None

    if looks_like_arxiv(query):
        arxiv_id = extract_arxiv_id(query)
        arxiv_info = resolve_arxiv(arxiv_id)
        if arxiv_info and arxiv_info.get("title"):
            # Search OpenAlex by the exact arXiv title and return candidates
            url = f"{OA_API}/works"
            params = {"search": arxiv_info["title"], "per-page": 10}
            data = oa_get(url, params)
            papers = [oa_work_to_paper(w) for w in data.get("results", [])]
            return sort_by_citations(papers)

    return None


def search_papers(query: str):
    direct = lookup_paper(query)
    if direct:
        if isinstance(direct, list):
            return direct[:SEARCH_LIMIT]
        return [direct]

    url = f"{OA_API}/works"
    params = {"search": query, "per-page": SEARCH_LIMIT}
    data = oa_get(url, params)
    papers = [oa_work_to_paper(w) for w in data.get("results", [])]
    return sort_by_citations(papers)


def get_paper(paper_id: str):
    url = f"{OA_API}/works/{paper_id}"
    return oa_work_to_paper(oa_get(url))


def get_references(paper_id: str, limit: int = REFERENCE_LIMIT):
    time.sleep(API_DELAY_SECONDS)
    work = oa_get(f"{OA_API}/works/{paper_id}")
    ref_ids = work.get("referenced_works", [])
    if not ref_ids:
        return []

    short_ids = [oid.split("/")[-1] for oid in ref_ids]
    refs = []
    chunk_size = 50
    fetch_limit = min(len(short_ids), limit * 3)
    for i in range(0, fetch_limit, chunk_size):
        chunk = short_ids[i : i + chunk_size]
        ids_filter = "|".join(chunk)
        url = f"{OA_API}/works"
        params = {"filter": f"openalex:{ids_filter}", "per-page": chunk_size}
        data = oa_get(url, params)
        refs.extend([oa_work_to_paper(w) for w in data.get("results", [])])

    return sort_by_citations(refs)[:limit]


def format_authors(authors: list, max_len: int = 3) -> str:
    if not authors:
        return "Unknown"
    names = [a.get("name", "Unknown") for a in authors]
    if len(names) > max_len:
        return ", ".join(names[:max_len]) + " et al."
    return ", ".join(names)


def paper_to_node(paper: dict, group: str) -> dict:
    title = paper.get("title") or "Untitled"
    citation_count = paper.get("citationCount") or 0
    return {
        "id": paper.get("paperId"),
        "label": title,
        "title": f"{title}<br>{format_authors(paper.get('authors', []))} ({paper.get('year', 'n.d.')})<br>Citations: {citation_count}",
        "value": citation_count,
        "group": group,
        "year": paper.get("year"),
        "authors": format_authors(paper.get("authors", [])),
        "citationCount": citation_count,
    }


def build_graph(paper_id: str):
    paper = get_paper(paper_id)
    if not paper or not paper.get("paperId"):
        raise ValueError("Paper not found")

    seed_id = paper["paperId"]
    nodes = [paper_to_node(paper, "seed")]
    edges = []

    references = get_references(seed_id, REFERENCE_LIMIT)

    for p in references:
        nodes.append(paper_to_node(p, "reference"))
        edges.append({"from": seed_id, "to": p["paperId"]})

    return {
        "seed_id": seed_id,
        "seed_title": paper.get("title"),
        "source": "openalex",
        "nodes": nodes,
        "edges": edges,
        "cached_at": datetime.now().isoformat(),
    }


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/api/search")
def api_search():
    q = request.args.get("q", "").strip()
    if not q:
        return jsonify({"error": "Missing query"}), 400
    try:
        papers = search_papers(q)
        return jsonify({"papers": papers})
    except requests.RequestException as e:
        return jsonify({"error": str(e)}), 502


@app.route("/api/graph/<path:paper_id>")
def api_graph(paper_id: str):
    paper_id = paper_id.strip()
    if not paper_id:
        return jsonify({"error": "Missing paper id"}), 400

    cached = load_cache(paper_id)
    path = cache_path(paper_id)

    if cached and is_cache_fresh(path):
        return jsonify(cached["data"])

    try:
        graph = build_graph(paper_id)
        save_cache(paper_id, graph)
        return jsonify(graph)
    except requests.RequestException:
        if cached:
            return jsonify(cached["data"])
        return jsonify({"error": "OpenAlex is unreachable. No cached graph available."}), 502
    except Exception as e:
        if cached:
            return jsonify(cached["data"])
        return jsonify({"error": f"Failed to build graph: {e}"}), 500


@app.route("/api/paper/<path:paper_id>")
def api_paper(paper_id: str):
    try:
        paper = get_paper(paper_id)
        return jsonify(paper)
    except requests.RequestException as e:
        return jsonify({"error": str(e)}), 502


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=False)