#!/usr/bin/env python3
"""Serve a local paper browser with lightweight Ollama-backed actions."""

from __future__ import annotations

import argparse
import hashlib
import html
import json
import os
import re
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from http import HTTPStatus
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
WEB_ROOT = ROOT / "web"
STATIC_ROOT = WEB_ROOT / "static"
INDEX_PATH = ROOT / "data" / "index.json"
CACHE_ROOT = WEB_ROOT / "cache"
ARXIV_CACHE = CACHE_ROOT / "arxiv"
AI_CACHE = CACHE_ROOT / "ai"
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://192.168.1.10:11434").rstrip("/")

DEFAULT_MODELS = {
    "translate": "ChatGPT-5.6:light",
    "summary": "ChatGPT-5.6:fast",
    "deep": "ChatGPT-5.6:large",
}

NS = {
    "atom": "http://www.w3.org/2005/Atom",
    "arxiv": "http://arxiv.org/schemas/atom",
}

INDEX_LOCK = threading.Lock()
OLLAMA_LOCK = threading.Lock()
INDEX_CACHE: dict[str, Any] = {"mtime": 0.0, "papers": [], "by_id": {}}


def ensure_cache_dirs() -> None:
    ARXIV_CACHE.mkdir(parents=True, exist_ok=True)
    AI_CACHE.mkdir(parents=True, exist_ok=True)


def clean_title(value: Any) -> str:
    title = str(value or "").strip()
    if len(title) >= 2 and title[0] == title[-1] == '"':
        return title[1:-1]
    return title


def item_id(path: str) -> str:
    return hashlib.sha1(path.encode("utf-8")).hexdigest()[:14]


def as_list(value: Any) -> list[str]:
    if value is None:
        return []
    if isinstance(value, list):
        return [str(item) for item in value if str(item).strip()]
    text = str(value).strip()
    return [text] if text else []


def score_value(value: Any) -> int:
    try:
        return int(str(value or "0"))
    except ValueError:
        return 0


def normalize_paper(item: dict[str, Any]) -> dict[str, Any]:
    meta = item.get("meta") or {}
    path = str(item.get("path") or "")
    title = clean_title(meta.get("title") or item.get("title") or Path(path).stem)
    topics = as_list(meta.get("topics"))
    datasets = as_list(meta.get("datasets"))
    authors = str(meta.get("authors") or "").strip()
    paper = {
        "id": item_id(path),
        "path": path,
        "title": title,
        "authors": authors,
        "year": str(meta.get("year") or ""),
        "venue": str(meta.get("venue") or ""),
        "url": str(meta.get("url") or ""),
        "code_url": as_list(meta.get("code_url")),
        "source": str(meta.get("source") or ""),
        "published_at": str(meta.get("published_at") or ""),
        "updated_at": str(meta.get("updated_at") or ""),
        "status": str(meta.get("status") or "unknown"),
        "relevance": str(meta.get("relevance") or ""),
        "topics": topics,
        "datasets": datasets,
        "collection_score": score_value(meta.get("collection_score")),
        "collection_queries": str(meta.get("collection_queries") or ""),
        "meta": meta,
    }
    search_text = " ".join(
        [
            title,
            authors,
            paper["year"],
            paper["venue"],
            paper["url"],
            paper["status"],
            paper["relevance"],
            paper["collection_queries"],
            " ".join(topics),
            " ".join(datasets),
        ]
    )
    paper["_search"] = search_text.lower()
    return paper


def load_papers() -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]]]:
    mtime = INDEX_PATH.stat().st_mtime
    with INDEX_LOCK:
        if INDEX_CACHE["mtime"] == mtime:
            return INDEX_CACHE["papers"], INDEX_CACHE["by_id"]
        raw_items = json.loads(INDEX_PATH.read_text(encoding="utf-8"))
        papers = [
            normalize_paper(item)
            for item in raw_items
            if item.get("collection") == "papers" and item.get("path")
        ]
        by_id = {paper["id"]: paper for paper in papers}
        INDEX_CACHE.update({"mtime": mtime, "papers": papers, "by_id": by_id})
        return papers, by_id


def read_json_body(handler: SimpleHTTPRequestHandler) -> dict[str, Any]:
    length = int(handler.headers.get("content-length") or "0")
    if length <= 0:
        return {}
    raw = handler.rfile.read(length)
    return json.loads(raw.decode("utf-8"))


def write_json(handler: SimpleHTTPRequestHandler, payload: Any, status: int = 200) -> None:
    data = json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8")
    handler.send_response(status)
    handler.send_header("content-type", "application/json; charset=utf-8")
    handler.send_header("cache-control", "no-store")
    handler.send_header("content-length", str(len(data)))
    handler.end_headers()
    handler.wfile.write(data)


def write_error(handler: SimpleHTTPRequestHandler, status: int, message: str, **extra: Any) -> None:
    payload = {"error": message, **extra}
    write_json(handler, payload, status)


def paper_public(paper: dict[str, Any]) -> dict[str, Any]:
    return {key: value for key, value in paper.items() if not key.startswith("_") and key != "meta"}


def topic_counts(papers: list[dict[str, Any]]) -> list[dict[str, Any]]:
    counts: dict[str, int] = {}
    for paper in papers:
        for topic in paper["topics"]:
            counts[topic] = counts.get(topic, 0) + 1
    return [
        {"topic": topic, "count": count}
        for topic, count in sorted(counts.items(), key=lambda item: (-item[1], item[0]))
    ]


def arxiv_id_from_paper(paper: dict[str, Any]) -> str:
    url = str(paper.get("url") or "")
    match = re.search(r"arxiv\.org/(?:abs|html|pdf)/([0-9]{4}\.[0-9]+)", url)
    return match.group(1) if match else ""


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


def fetch_arxiv(arxiv_id: str) -> dict[str, Any]:
    ensure_cache_dirs()
    cache_path = ARXIV_CACHE / f"{arxiv_id}.json"
    if cache_path.exists():
        return json.loads(cache_path.read_text(encoding="utf-8"))

    query = urllib.parse.urlencode({"id_list": arxiv_id})
    request = urllib.request.Request(
        f"https://export.arxiv.org/api/query?{query}",
        headers={"User-Agent": "agent-kb-paper-browser/0.1"},
    )
    with urllib.request.urlopen(request, timeout=20) as response:
        xml_text = response.read().decode("utf-8", "replace")
    root = ET.fromstring(xml_text)
    entry = root.find("atom:entry", NS)
    if entry is None:
        raise ValueError(f"arXiv entry not found: {arxiv_id}")

    authors = [
        normalize_space(author.findtext("atom:name", default="", namespaces=NS))
        for author in entry.findall("atom:author", NS)
    ]
    categories = [
        category.attrib.get("term", "")
        for category in entry.findall("atom:category", NS)
        if category.attrib.get("term")
    ]
    payload = {
        "arxiv_id": arxiv_id,
        "title": normalize_space(entry.findtext("atom:title", default="", namespaces=NS)),
        "abstract": normalize_space(entry.findtext("atom:summary", default="", namespaces=NS)),
        "authors": authors,
        "published": entry.findtext("atom:published", default="", namespaces=NS)[:10],
        "updated": entry.findtext("atom:updated", default="", namespaces=NS)[:10],
        "categories": categories,
        "url": f"https://arxiv.org/abs/{arxiv_id}",
    }
    cache_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return payload


def paper_context(paper: dict[str, Any], abstract: dict[str, Any] | None) -> str:
    meta = {
        "title": paper["title"],
        "authors": paper["authors"],
        "year": paper["year"],
        "venue": paper["venue"],
        "url": paper["url"],
        "topics": paper["topics"],
        "collection_score": paper["collection_score"],
        "collection_queries": paper["collection_queries"],
    }
    parts = [f"Metadata:\n{json.dumps(meta, ensure_ascii=False, indent=2)}"]
    if abstract and abstract.get("abstract"):
        parts.append(f"Abstract:\n{abstract['abstract']}")
    else:
        content = (ROOT / paper["path"]).read_text(encoding="utf-8")
        parts.append(f"Local note:\n{content[:4000]}")
    return "\n\n".join(parts)[:7000]


def prompt_for(mode: str, context: str) -> str:
    no_think = "不要输出思考过程，不要输出 <think> 标签，只输出最终结果。"
    if mode == "translate":
        return (
            "你是论文摘要翻译助手。把下面论文摘要准确翻译成中文，保留技术术语，"
            f"不要扩写，不要编造。如果没有摘要，只翻译已有材料。{no_think}\n\n"
            f"{context}"
        )
    if mode == "deep":
        return (
            "你是 Agent 研究员。请基于给定论文材料做中文深度研读笔记，输出：\n"
            "1. 一句话定位\n2. 研究问题\n3. 方法/系统设计\n4. 评估方式\n"
            "5. 对 Agent 工程实践的启发\n6. 可疑点或待验证问题\n7. 推荐标签。\n"
            f"只根据材料回答，信息不足就明确写“不足”。{no_think}\n\n"
            f"{context}"
        )
    return (
        "你是 Agent 论文知识库助手。请基于给定论文材料输出中文摘要，结构如下：\n"
        "1. 一句话结论\n2. 解决的问题\n3. 核心方法\n4. 评估/实验\n"
        "5. 和 Agent 知识库的关系\n6. 推荐优先级：P0/P1/P2。\n"
        f"保持精炼，只根据材料回答，禁止编造。{no_think}\n\n"
        f"{context}"
    )


def clean_model_response(value: str) -> str:
    text = re.sub(r"<think>[\s\S]*?</think>", "", value, flags=re.IGNORECASE).strip()
    return text or value.strip()


def call_ollama(model: str, prompt: str, mode: str) -> dict[str, Any]:
    options = {
        "temperature": 0.2,
        "num_ctx": 4096,
        "num_predict": 520,
    }
    if mode == "translate":
        options.update({"num_ctx": 3072, "num_predict": 700, "temperature": 0.1})
    elif mode == "deep":
        options.update({"num_ctx": 8192, "num_predict": 900, "temperature": 0.25})

    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "keep_alive": "2m",
        "options": options,
    }
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    request = urllib.request.Request(
        f"{OLLAMA_URL}/api/generate",
        data=data,
        headers={"content-type": "application/json"},
        method="POST",
    )
    started = time.time()
    with urllib.request.urlopen(request, timeout=240) as response:
        result = json.loads(response.read().decode("utf-8"))
    return {
        "model": model,
        "mode": mode,
        "response": clean_model_response(str(result.get("response") or "")),
        "duration_seconds": round(time.time() - started, 2),
        "eval_count": result.get("eval_count"),
        "prompt_eval_count": result.get("prompt_eval_count"),
    }


def ai_cache_path(paper_id: str, mode: str, model: str) -> Path:
    digest = hashlib.sha1(f"{paper_id}:{mode}:{model}".encode("utf-8")).hexdigest()[:18]
    return AI_CACHE / f"{digest}.json"


class PaperBrowserHandler(SimpleHTTPRequestHandler):
    server_version = "AgentPaperBrowser/0.1"

    def log_message(self, format: str, *args: Any) -> None:
        print(f"[paper-web] {self.address_string()} - {format % args}")

    def do_GET(self) -> None:
        parsed = urllib.parse.urlparse(self.path)
        path = parsed.path
        query = urllib.parse.parse_qs(parsed.query)

        if path == "/":
            self.serve_file(STATIC_ROOT / "index.html", "text/html; charset=utf-8")
            return
        if path.startswith("/static/"):
            target = STATIC_ROOT / path.removeprefix("/static/")
            content_type = "text/plain; charset=utf-8"
            if target.suffix == ".css":
                content_type = "text/css; charset=utf-8"
            elif target.suffix == ".js":
                content_type = "application/javascript; charset=utf-8"
            self.serve_file(target, content_type)
            return
        if path == "/api/health":
            self.handle_health()
            return
        if path == "/api/models":
            self.handle_models()
            return
        if path == "/api/topics":
            papers, _ = load_papers()
            write_json(self, {"topics": topic_counts(papers)})
            return
        if path == "/api/papers":
            self.handle_paper_search(query)
            return
        if path == "/api/paper":
            self.handle_paper_detail(query)
            return
        if path == "/api/paper/abstract":
            self.handle_abstract(query)
            return

        write_error(self, HTTPStatus.NOT_FOUND, "not found")

    def do_POST(self) -> None:
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == "/api/ai":
            self.handle_ai()
            return
        write_error(self, HTTPStatus.NOT_FOUND, "not found")

    def serve_file(self, target: Path, content_type: str) -> None:
        try:
            resolved = target.resolve()
            if not str(resolved).startswith(str(STATIC_ROOT.resolve())):
                write_error(self, HTTPStatus.FORBIDDEN, "forbidden")
                return
            data = resolved.read_bytes()
        except FileNotFoundError:
            write_error(self, HTTPStatus.NOT_FOUND, "file not found")
            return
        self.send_response(HTTPStatus.OK)
        self.send_header("content-type", content_type)
        self.send_header("content-length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def handle_health(self) -> None:
        papers, _ = load_papers()
        ollama_ok = False
        try:
            request = urllib.request.Request(f"{OLLAMA_URL}/api/tags")
            with urllib.request.urlopen(request, timeout=3) as response:
                ollama_ok = response.status == 200
        except (urllib.error.URLError, TimeoutError):
            ollama_ok = False
        write_json(
            self,
            {
                "ok": True,
                "paper_count": len(papers),
                "ollama_url": OLLAMA_URL,
                "ollama_ok": ollama_ok,
                "default_models": DEFAULT_MODELS,
            },
        )

    def handle_models(self) -> None:
        try:
            request = urllib.request.Request(f"{OLLAMA_URL}/api/tags")
            with urllib.request.urlopen(request, timeout=5) as response:
                payload = json.loads(response.read().decode("utf-8"))
            models = [
                {
                    "name": item.get("name"),
                    "parameter_size": (item.get("details") or {}).get("parameter_size"),
                    "context_length": (item.get("details") or {}).get("context_length"),
                }
                for item in payload.get("models", [])
            ]
            write_json(self, {"models": models, "default_models": DEFAULT_MODELS})
        except Exception as exc:  # noqa: BLE001
            write_error(self, HTTPStatus.BAD_GATEWAY, f"ollama unavailable: {exc}")

    def handle_paper_search(self, query: dict[str, list[str]]) -> None:
        papers, _ = load_papers()
        q = (query.get("q") or [""])[0].strip().lower()
        topic = (query.get("topic") or [""])[0].strip()
        year = (query.get("year") or [""])[0].strip()
        status = (query.get("status") or [""])[0].strip()
        relevance = (query.get("relevance") or [""])[0].strip()
        sort = (query.get("sort") or ["score"])[0].strip()
        limit = max(1, min(100, int((query.get("limit") or ["40"])[0])))
        offset = max(0, int((query.get("offset") or ["0"])[0]))

        filtered = []
        for paper in papers:
            if q and q not in paper["_search"]:
                continue
            if topic and topic not in paper["topics"]:
                continue
            if year and paper["year"] != year:
                continue
            if status and paper["status"] != status:
                continue
            if relevance and paper["relevance"] != relevance:
                continue
            filtered.append(paper)

        if sort == "date":
            filtered.sort(key=lambda item: (item["published_at"], item["collection_score"]), reverse=True)
        elif sort == "title":
            filtered.sort(key=lambda item: item["title"].lower())
        else:
            filtered.sort(key=lambda item: (item["collection_score"], item["published_at"]), reverse=True)

        page = filtered[offset : offset + limit]
        write_json(
            self,
            {
                "total": len(filtered),
                "offset": offset,
                "limit": limit,
                "items": [paper_public(paper) for paper in page],
            },
        )

    def handle_paper_detail(self, query: dict[str, list[str]]) -> None:
        paper_id = (query.get("id") or [""])[0]
        _, by_id = load_papers()
        paper = by_id.get(paper_id)
        if not paper:
            write_error(self, HTTPStatus.NOT_FOUND, "paper not found")
            return
        text = (ROOT / paper["path"]).read_text(encoding="utf-8")
        write_json(self, {"paper": paper_public(paper), "markdown": text})

    def handle_abstract(self, query: dict[str, list[str]]) -> None:
        paper_id = (query.get("id") or [""])[0]
        _, by_id = load_papers()
        paper = by_id.get(paper_id)
        if not paper:
            write_error(self, HTTPStatus.NOT_FOUND, "paper not found")
            return
        arxiv_id = arxiv_id_from_paper(paper)
        if not arxiv_id:
            write_json(self, {"abstract": None, "message": "not an arXiv paper"})
            return
        try:
            write_json(self, {"abstract": fetch_arxiv(arxiv_id)})
        except Exception as exc:  # noqa: BLE001
            write_error(self, HTTPStatus.BAD_GATEWAY, f"arXiv fetch failed: {exc}")

    def handle_ai(self) -> None:
        try:
            payload = read_json_body(self)
            paper_id = str(payload.get("id") or "")
            mode = str(payload.get("mode") or "summary")
            refresh = bool(payload.get("refresh"))
            model = str(payload.get("model") or DEFAULT_MODELS.get(mode) or DEFAULT_MODELS["summary"])
            if mode not in {"summary", "translate", "deep"}:
                write_error(self, HTTPStatus.BAD_REQUEST, "invalid ai mode")
                return

            _, by_id = load_papers()
            paper = by_id.get(paper_id)
            if not paper:
                write_error(self, HTTPStatus.NOT_FOUND, "paper not found")
                return

            ensure_cache_dirs()
            cache_path = ai_cache_path(paper_id, mode, model)
            if cache_path.exists() and not refresh:
                result = json.loads(cache_path.read_text(encoding="utf-8"))
                result["cached"] = True
                write_json(self, result)
                return

            if not OLLAMA_LOCK.acquire(blocking=False):
                write_error(self, HTTPStatus.TOO_MANY_REQUESTS, "ollama is busy; try again later")
                return
            try:
                abstract = None
                arxiv_id = arxiv_id_from_paper(paper)
                if arxiv_id:
                    abstract = fetch_arxiv(arxiv_id)
                context = paper_context(paper, abstract)
                result = call_ollama(model, prompt_for(mode, context), mode)
                result.update(
                    {
                        "id": paper_id,
                        "title": paper["title"],
                        "cached": False,
                        "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
                    }
                )
                cache_path.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
                write_json(self, result)
            finally:
                OLLAMA_LOCK.release()
        except Exception as exc:  # noqa: BLE001
            write_error(self, HTTPStatus.INTERNAL_SERVER_ERROR, str(exc))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default=os.environ.get("HOST", "0.0.0.0"))
    parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", "18080")))
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    ensure_cache_dirs()
    load_papers()
    server = ThreadingHTTPServer((args.host, args.port), PaperBrowserHandler)
    print(f"Paper browser: http://{args.host}:{args.port}")
    print(f"Ollama: {OLLAMA_URL}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopping paper browser")
    finally:
        server.server_close()
    return 0


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