app.py

# import json
# import os
# import queue
# import threading
# import time
# import uuid
# from dataclasses import dataclass, field
# from typing import Any, Dict, Optional

# from flask import Flask, Response, jsonify, render_template, request

# from bibcheck.engine import verify_bibtex

# app = Flask(__name__)
# app.config["MAX_CONTENT_LENGTH"] = 1 * 1024 * 1024  # 1MB


# # @dataclass
# # class Job:
# #     id: str
# #     q: "queue.Queue[Dict[str, Any]]" = field(default_factory=queue.Queue)
# #     done: bool = False
# #     result: Optional[Dict[str, Any]] = None
# #     error: Optional[str] = None
# #     started_at: float = field(default_factory=time.time)

# @dataclass
# class Job:
#     id: str
#     q: "queue.Queue[Dict[str, Any]]" = field(default_factory=queue.Queue)
#     done: bool = False
#     result: Optional[Dict[str, Any]] = None
#     error: Optional[str] = None
#     started_at: float = field(default_factory=time.time)
#     debug: Dict[str, Any] = field(default_factory=dict)  # ✅ 新增



# JOBS: Dict[str, Job] = {}
# JOBS_LOCK = threading.Lock()


# def push(job: Job, event: str, data: Any):
#     job.q.put({"event": event, "data": data})


# def run_job(job: Job, bib_content: str, config: Dict[str, Any]):
#     """
#     ⚠️ 这个函数在后台线程执行，严禁直接使用 flask.request / session 等请求上下文对象
#     """
#     try:
#         def emit(event: str, data: Any):
#             push(job, event, data)

#         result = verify_bibtex(
#             bib_content,
#             emit=emit,
#             config=config,   # ✅ 使用 submit() 里传进来的 config
#             debug_store=job.debug,  # 新增，加入详细信息输出
#         )
#         job.result = result
#         push(job, "result", result)
#         push(job, "done", {"ok": True})
#     except Exception as e:
#         job.error = f"{type(e).__name__}: {e}"
#         push(job, "error", {"message": job.error})
#         push(job, "done", {"ok": False})
#     finally:
#         job.done = True


# @app.route("/", methods=["GET"])
# def index():
#     return render_template("index.html")


# @app.route("/submit", methods=["POST"])
# def submit():
#     bib_text = ""

#     # 1) file
#     if "bibfile" in request.files:
#         f = request.files["bibfile"]
#         if f and f.filename:
#             bib_text = f.read().decode("utf-8", errors="replace")

#     # 2) paste
#     pasted = request.form.get("bibtext", "").strip()
#     if pasted:
#         bib_text = pasted

#     if not bib_text.strip():
#         return jsonify({"ok": False, "error": "Empty input. Please upload a .bib or paste BibTeX content."}), 400

#     job_id = uuid.uuid4().hex
#     job = Job(id=job_id)
#     with JOBS_LOCK:
#         JOBS[job_id] = job

#     # ✅ 在 request context 内把需要的信息取出来
#     # （你也可以直接写死为 None，不影响功能，只是对 Crossref 更“礼貌”）
#     crossref_mailto = request.headers.get("X-User-Email") or None
    
#     s2_api_key = os.getenv("SEMANTIC_SCHOLAR_API_KEY", "").strip() or None

#     config = {
#         "crossref_mailto": crossref_mailto,
#         "timeout_sec": 12,
#         "max_candidates": 5,
#         "per_source_min_interval_sec": 0.25,
#         # "scholar_min_interval_sec": 2.0,   # 太小很容易 429；你可按环境调大
#         # "scholar_fill_top": True,          # 只 fill 最匹配的一个
#         "s2_api_key": s2_api_key,
#         "s2_min_interval_sec": 0.25,   # 你可以按需要调大
#         "s2_max_retries": 2,
#     }

#     # Start worker thread（把 config 传进去）
#     t = threading.Thread(target=run_job, args=(job, bib_text, config), daemon=True)
#     t.start()

#     return jsonify({"ok": True, "job_id": job_id})


# @app.route("/stream/<job_id>", methods=["GET"])
# def stream(job_id: str):
#     with JOBS_LOCK:
#         job = JOBS.get(job_id)

#     if not job:
#         return jsonify({"ok": False, "error": "Job not found"}), 404

#     def gen():
#         yield "event: ping\ndata: {}\n\n"
#         while True:
#             try:
#                 msg = job.q.get(timeout=15)
#             except queue.Empty:
#                 yield "event: ping\ndata: {}\n\n"
#                 if job.done:
#                     break
#                 continue

#             ev = msg.get("event", "log")
#             data = msg.get("data", {})
#             payload = json.dumps(data, ensure_ascii=False)
#             yield f"event: {ev}\ndata: {payload}\n\n"

#             if ev == "done":
#                 break

#     headers = {
#         "Content-Type": "text/event-stream",
#         "Cache-Control": "no-cache, no-transform",
#         "Connection": "keep-alive",
#         "X-Accel-Buffering": "no",
#     }
#     return Response(gen(), headers=headers)


# @app.route("/result/<job_id>", methods=["GET"])
# def get_result(job_id: str):
#     with JOBS_LOCK:
#         job = JOBS.get(job_id)
#     if not job:
#         return jsonify({"ok": False, "error": "Job not found"}), 404
#     if not job.done:
#         return jsonify({"ok": False, "error": "Job not finished"}), 202
#     return jsonify({"ok": True, "result": job.result, "error": job.error})


# @app.route("/health", methods=["GET"])
# def health():
#     return jsonify({"ok": True})


# @app.route("/debug/<job_id>/<citekey>", methods=["GET"])
# def debug_entry(job_id: str, citekey: str):
#     with JOBS_LOCK:
#         job = JOBS.get(job_id)
#     if not job:
#         return jsonify({"ok": False, "error": "Job not found"}), 404

#     dbg = job.debug.get(citekey)
#     if dbg is None:
#         return jsonify({"ok": False, "error": f"Debug not found for key={citekey}"}), 404

#     return jsonify({"ok": True, "key": citekey, "debug": dbg})



# if __name__ == "__main__":
#     app.run(host="0.0.0.0", port=5001, debug=True, threaded=True, use_reloader=False)

import json
import os
import queue
import sqlite3
import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

from cryptography.fernet import Fernet
from flask import Flask, Response, jsonify, render_template, request, session, redirect, url_for, g, flash
from werkzeug.security import generate_password_hash, check_password_hash

from bibcheck.engine import verify_bibtex

# -------------------- App --------------------
app = Flask(__name__)
app.config["MAX_CONTENT_LENGTH"] = 1 * 1024 * 1024  # 1MB

BASE_DIR = os.path.abspath(os.path.dirname(__file__))
INSTANCE_DIR = os.path.join(BASE_DIR, "instance")
DB_PATH = os.path.join(INSTANCE_DIR, "refcheck.db")
SECRET_PATH = os.path.join(INSTANCE_DIR, "secret.key")
FERNET_PATH = os.path.join(INSTANCE_DIR, "fernet.key")

os.makedirs(INSTANCE_DIR, exist_ok=True)


def _load_or_create_file(path: str, nbytes: int = 32) -> bytes:
    if os.path.exists(path):
        with open(path, "rb") as f:
            return f.read().strip()
    data = os.urandom(nbytes)
    with open(path, "wb") as f:
        f.write(data)
    return data


def _init_secret_keys():
    # Flask session secret
    secret = os.getenv("REFCHECK_SECRET_KEY")
    if secret:
        app.config["SECRET_KEY"] = secret
    else:
        app.config["SECRET_KEY"] = _load_or_create_file(SECRET_PATH, 32)

    # Fernet encryption key for storing API keys at rest
    fernet_key = os.getenv("REFCHECK_FERNET_KEY")
    if fernet_key:
        # must be urlsafe_b64 key
        fk = fernet_key.encode("utf-8")
    else:
        if os.path.exists(FERNET_PATH):
            fk = _load_or_create_file(FERNET_PATH, 44)  # Fernet key len typically 44 chars base64
            # If file was random bytes, regenerate correctly:
            try:
                Fernet(fk)
            except Exception:
                fk = Fernet.generate_key()
                with open(FERNET_PATH, "wb") as f:
                    f.write(fk)
        else:
            fk = Fernet.generate_key()
            with open(FERNET_PATH, "wb") as f:
                f.write(fk)
    return Fernet(fk)


FERNET = _init_secret_keys()


# -------------------- DB --------------------
def get_db() -> sqlite3.Connection:
    if "db" not in g:
        conn = sqlite3.connect(DB_PATH)
        conn.row_factory = sqlite3.Row
        g.db = conn
    return g.db


@app.teardown_appcontext
def close_db(_exc):
    db = g.pop("db", None)
    if db is not None:
        db.close()


def init_db():
    db = get_db()
    db.execute(
        """
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            created_at INTEGER NOT NULL
        )
        """
    )
    db.execute(
        """
        CREATE TABLE IF NOT EXISTS user_keys (
            user_id INTEGER UNIQUE NOT NULL,
            s2_api_key_enc TEXT,
            updated_at INTEGER NOT NULL,
            FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE
        )
        """
    )
    db.commit()


def create_user(username: str, password: str) -> Optional[int]:
    db = get_db()
    ph = generate_password_hash(password)
    try:
        cur = db.execute(
            "INSERT INTO users(username, password_hash, created_at) VALUES(?,?,?)",
            (username, ph, int(time.time())),
        )
        uid = cur.lastrowid
        db.execute("INSERT OR IGNORE INTO user_keys(user_id, s2_api_key_enc, updated_at) VALUES(?,?,?)", (uid, None, int(time.time())))
        db.commit()
        return int(uid)
    except sqlite3.IntegrityError:
        return None


def get_user_by_username(username: str):
    db = get_db()
    return db.execute("SELECT * FROM users WHERE username=?", (username,)).fetchone()


def get_user_by_id(uid: int):
    db = get_db()
    return db.execute("SELECT * FROM users WHERE id=?", (uid,)).fetchone()


def verify_user(username: str, password: str) -> Optional[int]:
    u = get_user_by_username(username)
    if not u:
        return None
    if not check_password_hash(u["password_hash"], password):
        return None
    return int(u["id"])


def set_s2_api_key(uid: int, api_key: Optional[str]):
    db = get_db()
    if api_key is None or api_key.strip() == "":
        enc = None
    else:
        enc = FERNET.encrypt(api_key.strip().encode("utf-8")).decode("utf-8")
    db.execute(
        "INSERT INTO user_keys(user_id, s2_api_key_enc, updated_at) VALUES(?,?,?) "
        "ON CONFLICT(user_id) DO UPDATE SET s2_api_key_enc=excluded.s2_api_key_enc, updated_at=excluded.updated_at",
        (uid, enc, int(time.time())),
    )
    db.commit()


def get_s2_api_key(uid: int) -> Optional[str]:
    db = get_db()
    row = db.execute("SELECT s2_api_key_enc FROM user_keys WHERE user_id=?", (uid,)).fetchone()
    if not row:
        return None
    enc = row["s2_api_key_enc"]
    if not enc:
        return None
    try:
        return FERNET.decrypt(enc.encode("utf-8")).decode("utf-8")
    except Exception:
        return None


# -------------------- Auth helpers --------------------
def login_required(fn):
    def wrapper(*args, **kwargs):
        if not g.user:
            flash("请先登录后再访问该页面。", "warning")
            return redirect(url_for("login", next=request.path))
        return fn(*args, **kwargs)
    wrapper.__name__ = fn.__name__
    return wrapper


@app.before_request
def load_user():
    init_db()
    uid = session.get("user_id")
    g.user = get_user_by_id(uid) if uid else None
    g.has_s2_key = False
    if g.user:
        g.has_s2_key = bool(get_s2_api_key(int(g.user["id"])))


@app.context_processor
def inject_globals():
    return {
        "current_user": g.user,
        "has_s2_key": getattr(g, "has_s2_key", False),
    }


# -------------------- Job system (SSE) --------------------
@dataclass
class Job:
    id: str
    q: "queue.Queue[Dict[str, Any]]" = field(default_factory=queue.Queue)
    done: bool = False
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    started_at: float = field(default_factory=time.time)
    debug: Dict[str, Any] = field(default_factory=dict)


JOBS: Dict[str, Job] = {}
JOBS_LOCK = threading.Lock()


def push(job: Job, event: str, data: Any):
    job.q.put({"event": event, "data": data})


def run_job(job: Job, bib_content: str, config: Dict[str, Any]):
    try:
        def emit(event: str, data: Any):
            push(job, event, data)

        result = verify_bibtex(
            bib_content,
            emit=emit,
            config=config,
            debug_store=job.debug,
        )
        job.result = result
        push(job, "result", result)
        push(job, "done", {"ok": True})
    except Exception as e:
        job.error = f"{type(e).__name__}: {e}"
        push(job, "error", {"message": job.error})
        push(job, "done", {"ok": False})
    finally:
        job.done = True


# -------------------- Routes --------------------
@app.route("/", methods=["GET"])
def index():
    return render_template("index.html")


@app.route("/register", methods=["GET", "POST"])
def register():
    if request.method == "GET":
        return render_template("register.html")

    username = (request.form.get("username") or "").strip()
    password = request.form.get("password") or ""
    password2 = request.form.get("password2") or ""

    if len(username) < 3:
        flash("用户名至少 3 个字符。", "error")
        return redirect(url_for("register"))
    if len(password) < 6:
        flash("密码至少 6 位。", "error")
        return redirect(url_for("register"))
    if password != password2:
        flash("两次输入的密码不一致。", "error")
        return redirect(url_for("register"))

    uid = create_user(username, password)
    if uid is None:
        flash("该用户名已被注册。", "error")
        return redirect(url_for("register"))

    session["user_id"] = uid
    flash("注册成功。你可以在个人页面填写 Semantic Scholar API Key 以启用完整功能。", "success")
    return redirect(url_for("account"))


@app.route("/login", methods=["GET", "POST"])
def login():
    if request.method == "GET":
        return render_template("login.html", next=request.args.get("next", ""))

    username = (request.form.get("username") or "").strip()
    password = request.form.get("password") or ""
    uid = verify_user(username, password)
    if uid is None:
        flash("用户名或密码错误。", "error")
        return redirect(url_for("login"))

    session["user_id"] = uid
    flash("登录成功。", "success")
    nxt = request.form.get("next") or ""
    return redirect(nxt or url_for("index"))


@app.route("/logout", methods=["POST"])
def logout():
    session.pop("user_id", None)
    flash("已退出登录。", "success")
    return redirect(url_for("index"))


@app.route("/account", methods=["GET", "POST"])
@login_required
def account():
    uid = int(g.user["id"])
    if request.method == "POST":
        action = request.form.get("action") or ""
        if action == "save":
            api_key = (request.form.get("s2_api_key") or "").strip()
            if not api_key:
                flash("API Key 不能为空。", "error")
                return redirect(url_for("account"))
            set_s2_api_key(uid, api_key)
            flash("API Key 已保存。", "success")
            return redirect(url_for("account"))

        if action == "delete":
            set_s2_api_key(uid, None)
            flash("API Key 已删除。", "success")
            return redirect(url_for("account"))

        flash("未知操作。", "error")
        return redirect(url_for("account"))

    key = get_s2_api_key(uid)
    masked = None
    if key:
        masked = (key[:4] + "..." + key[-4:]) if len(key) >= 10 else "***"
    return render_template("account.html", masked_key=masked)


@app.route("/submit", methods=["POST"])
def submit():
    bib_text = ""

    if "bibfile" in request.files:
        f = request.files["bibfile"]
        if f and f.filename:
            bib_text = f.read().decode("utf-8", errors="replace")

    pasted = (request.form.get("bibtext") or "").strip()
    if pasted:
        bib_text = pasted

    if not bib_text.strip():
        return jsonify({"ok": False, "error": "Empty input. Please upload a .bib or paste BibTeX content."}), 400

    job_id = uuid.uuid4().hex
    job = Job(id=job_id)

    with JOBS_LOCK:
        JOBS[job_id] = job

    # Use user's S2 key if logged in and provided
    s2_key = None
    if g.user:
        s2_key = get_s2_api_key(int(g.user["id"]))

    crossref_mailto = None  # optional

    config = {
        "timeout_sec": 12,
        "max_candidates": 5,
        "per_source_min_interval_sec": 0.25,
        "crossref_mailto": crossref_mailto,

        # Semantic Scholar user key
        "s2_api_key": s2_key,
        "enable_semanticscholar": bool(s2_key),
        "s2_min_interval_sec": 0.25,
        "s2_max_retries": 2,
    }

    t = threading.Thread(target=run_job, args=(job, bib_text, config), daemon=True)
    t.start()
    return jsonify({"ok": True, "job_id": job_id})


@app.route("/stream/<job_id>", methods=["GET"])
def stream(job_id: str):
    with JOBS_LOCK:
        job = JOBS.get(job_id)

    if not job:
        return jsonify({"ok": False, "error": "Job not found"}), 404

    def gen():
        yield "event: ping\ndata: {}\n\n"
        while True:
            try:
                msg = job.q.get(timeout=15)
            except queue.Empty:
                yield "event: ping\ndata: {}\n\n"
                if job.done:
                    break
                continue

            ev = msg.get("event", "log")
            data = msg.get("data", {})
            payload = json.dumps(data, ensure_ascii=False)
            yield f"event: {ev}\ndata: {payload}\n\n"

            if ev == "done":
                break

    headers = {
        "Content-Type": "text/event-stream",
        "Cache-Control": "no-cache, no-transform",
        "Connection": "keep-alive",
        "X-Accel-Buffering": "no",
    }
    return Response(gen(), headers=headers)


@app.route("/debug/<job_id>/<citekey>", methods=["GET"])
def debug_entry(job_id: str, citekey: str):
    with JOBS_LOCK:
        job = JOBS.get(job_id)
    if not job:
        return jsonify({"ok": False, "error": "Job not found"}), 404
    dbg = job.debug.get(citekey)
    if dbg is None:
        return jsonify({"ok": False, "error": f"Debug not found for key={citekey}"}), 404
    return jsonify({"ok": True, "key": citekey, "debug": dbg})


@app.route("/health", methods=["GET"])
def health():
    return jsonify({"ok": True})


if __name__ == "__main__":
    # dev only
    app.run(host="0.0.0.0", port=5001, debug=True, threaded=True, use_reloader=False)