pdf2zh/webapp/runner.py

"""Framework-agnostic translation runner.

Runs the OCR -> translate -> render pipeline in a worker thread (so its internal
``asyncio.run`` works) and streams per-phase progress through a queue. Knows
nothing about Gradio, so it can be unit-tested in isolation.
"""

from __future__ import annotations

import logging
import queue
import tempfile
import threading
import traceback
import uuid
from collections.abc import Callable, Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class TranslationRequest:
    """One translation job, with the provider already resolved to its key."""

    pdf_path: str | None
    provider: str  # Phase-2 provider key (e.g. "openrouter")
    api_key: str
    model: str | None
    src_lang: str
    tgt_lang: str
    font: str
    pages: list[int] | None


@dataclass(frozen=True)
class Progress:
    """A per-phase progress update streamed while the pipeline runs."""

    frac: float
    msg: str


@dataclass(frozen=True)
class Result:
    """The terminal outcome of a run.

    ``data`` carries the step's payload for the stepped flow (e.g. the parsed
    dict, or ``{"translated": ..., "out_path": ...}``). ``out_path`` is kept for
    the legacy one-shot ``stream_translation`` path.
    """

    status: str  # "ok" | "invalid" | "error"
    out_path: str | None = None
    detail: str = ""
    data: Any = None


def validate(req: TranslationRequest) -> str | None:
    """Return a user-facing error message if the request can't run, else None."""
    if not req.pdf_path:
        return "Vui lòng tải lên một file PDF."
    if not req.api_key or not req.api_key.strip():
        return "Thiếu API key — nhập API key của provider ở thanh bên."
    if not req.src_lang or not req.tgt_lang:
        return "Chọn ngôn ngữ nguồn và ngôn ngữ đích."
    return None


def list_models(provider: str, api_key: str) -> list[str]:
    """Fetch model ids from a provider's OpenAI-compatible ``GET /models`` endpoint.

    ``provider`` is the resolved key (e.g. "deepseek"). Returns a sorted list of
    model ids, or [] on any failure (bad key, no endpoint, non-OpenAI response) —
    the UI then falls back to free-text entry.
    """
    import httpx

    from pdf2zh.translation.config import provider_base_url

    if not api_key or not api_key.strip():
        return []
    try:
        base = provider_base_url(provider).rstrip("/")
        resp = httpx.get(
            f"{base}/models",
            headers={"Authorization": f"Bearer {api_key.strip()}"},
            timeout=15,
            verify=False,
        )
        resp.raise_for_status()
        data = resp.json().get("data", [])
        return sorted({m["id"] for m in data if isinstance(m, dict) and m.get("id")})
    except Exception:  # noqa: BLE001 — listing is best-effort; fall back to manual
        logger.warning("list_models failed for provider %s", provider, exc_info=True)
        return []


def stream_translation(req: TranslationRequest) -> Iterator[Progress | Result]:
    """Yield ``Progress`` updates while translating, then one terminal ``Result``."""
    # Imported lazily so the lightweight bits above (dataclasses, validate) stay
    # importable without the heavy ML stack (torch, surya, ...) that e2e pulls in.
    from pdf2zh.e2e import run_pipeline

    q: queue.Queue = queue.Queue()
    work_dir = Path(tempfile.gettempdir()) / f"pdf2zh_{uuid.uuid4().hex}"

    def on_progress(frac: float, msg: str) -> None:
        q.put(Progress(frac, msg))

    def worker() -> None:
        try:
            out = run_pipeline(
                pdf_path=req.pdf_path,
                src_lang=req.src_lang,
                tgt_lang=req.tgt_lang,
                provider=req.provider,
                api_key=req.api_key,
                model=req.model,
                pages=req.pages,
                font=req.font,
                work_dir=work_dir,
                progress=on_progress,
            )
            q.put(Result("ok", out_path=out))
        except ValueError as exc:  # user-facing input error
            q.put(Result("invalid", detail=str(exc)))
        except Exception as exc:  # noqa: BLE001 — surface anything else to the UI
            logger.exception("pipeline failed")
            tail = "".join(traceback.format_exc().splitlines(keepends=True)[-6:])
            q.put(
                Result("error", detail=f"{type(exc).__name__}: {exc}\n```\n{tail}\n```")
            )

    threading.Thread(target=worker, daemon=True).start()
    while True:
        item = q.get()
        yield item
        if isinstance(item, Result):
            return


# --------------------------------------------------------------------------- #
# Stepped flow — run one phase in a worker thread and stream its progress.
# --------------------------------------------------------------------------- #
def _stream(
    fn: Callable[[Callable[[float, str], None]], Any],
) -> Iterator[Progress | Result]:
    """Run ``fn(progress_cb)`` in a worker thread; stream Progress then a Result.

    ``fn`` receives a ``progress(frac, msg)`` callback and returns the payload
    placed on ``Result.data``. A ValueError becomes an ``invalid`` result
    (user-facing input error); anything else becomes an ``error`` result.
    """
    q: queue.Queue = queue.Queue()

    def on_progress(frac: float, msg: str) -> None:
        q.put(Progress(frac, msg))

    def worker() -> None:
        try:
            data = fn(on_progress)
            q.put(Result("ok", data=data))
        except ValueError as exc:
            q.put(Result("invalid", detail=str(exc)))
        except Exception as exc:  # noqa: BLE001 — surface anything else to the UI
            logger.exception("step failed")
            tail = "".join(traceback.format_exc().splitlines(keepends=True)[-6:])
            q.put(
                Result("error", detail=f"{type(exc).__name__}: {exc}\n```\n{tail}\n```")
            )

    threading.Thread(target=worker, daemon=True).start()
    while True:
        item = q.get()
        yield item
        if isinstance(item, Result):
            return


def stream_parse(
    pdf_path: str, pages: list[int] | None, work_dir: str | Path
) -> Iterator[Progress | Result]:
    """Phase 1 — parse. ``Result.data`` is the parsed doc dict."""
    from pdf2zh.e2e import run_parse

    return _stream(lambda p: run_parse(pdf_path, pages, work_dir, p))


def stream_translate_render(
    pdf_path: str,
    parsed: dict,
    src_lang: str,
    tgt_lang: str,
    provider: str,
    api_key: str,
    model: str | None,
    pages: list[int] | None,
    font: str,
    work_dir: str | Path,
) -> Iterator[Progress | Result]:
    """Phase 2 + 3 — translate the (edited) parsed doc, then render.

    ``Result.data`` is ``{"translated": dict, "out_path": str}``.
    """
    from pdf2zh.e2e import run_render, run_translate

    def fn(p: Callable[[float, str], None]) -> dict:
        translated = run_translate(
            parsed, src_lang, tgt_lang, provider, api_key, model, work_dir, p
        )
        out_path = run_render(pdf_path, translated, pages, font, work_dir, p)
        return {"translated": translated, "out_path": out_path}

    return _stream(fn)


def stream_render(
    pdf_path: str,
    translated: dict,
    pages: list[int] | None,
    font: str,
    work_dir: str | Path,
) -> Iterator[Progress | Result]:
    """Phase 3 only — re-render the (edited) translated doc. ``Result.data`` is
    ``{"out_path": str}``."""
    from pdf2zh.e2e import run_render

    def fn(p: Callable[[float, str], None]) -> dict:
        return {"out_path": run_render(pdf_path, translated, pages, font, work_dir, p)}

    return _stream(fn)