conjurer/latex_functions.py

import asyncio
import logging
import re
import shutil
import tempfile
import zipfile
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

# ===== Bezpieczeństwo nazw/rozszerzeń =====
SAFE_TEX_NAME = re.compile(r"^[\w\-. ]+\.tex$", re.IGNORECASE)
SAFE_ASSET_NAME = re.compile(
    r"^[\w\-. /]+?\.(tex|png|jpg|jpeg|pdf|svg|eps|bmp|gif|sty|cls|bib|bst|bbx|cbx|def|cfg)$",
    re.IGNORECASE,
)

def is_safe_tex_name(name: str) -> bool:
    return bool(SAFE_TEX_NAME.match(name or ""))

def is_safe_asset_name(name: str) -> bool:
    return bool(SAFE_ASSET_NAME.match(name or ""))

# ===== Uruchamianie kompilatora =====

async def _run(cmd: List[str], cwd: Path, timeout_sec: int, logger_name: Optional[str]) -> Tuple[int, str]:
    """Uruchamia proces kompilatora w cwd, zwraca (rc, sklejone_stdout_stderr)."""
    logger = logging.getLogger(logger_name) if logger_name else logging.getLogger()
    logger.debug("RUN %s (cwd=%s, timeout=%s)", " ".join(cmd), cwd, timeout_sec)

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
        cwd=str(cwd),
    )
    try:
        out = await asyncio.wait_for(proc.communicate(), timeout=timeout_sec)
    except asyncio.TimeoutError:
        try:
            proc.kill()
        finally:
            pass
        return 124, "TIMEOUT: compiler took too long"
    rc = proc.returncode
    log = (out[0] or b"").decode("utf-8", errors="ignore")
    logger.debug("RC=%s, LOG TAIL:\n%s", rc, log[-1500:])
    return rc, log

def _build_cmd(tex_engine: str, main_tex: str) -> List[str]:
    """Buduje komendę kompilatora; domyślnie tectonic/pdflatex/xelatex."""
    if tex_engine.lower() == "tectonic":
        # Tectonic sam robi wielofazowość, ale i tak uruchamiamy 2x (bezpiecznie dla bib/refs)
        return ["tectonic", "--keep-intermediates", "--synctex", main_tex]
    elif tex_engine.lower() in ("pdflatex", "xelatex", "lualatex"):
        return [tex_engine, "-interaction=nonstopmode", "-halt-on-error", main_tex]
    else:
        # fallback – traktujemy jak pdflatex
        return ["pdflatex", "-interaction=nonstopmode", "-halt-on-error", main_tex]

async def run_latex_two_passes(
    tex_filename: str, workdir: Path, tex_engine: str, timeout_sec: int, logger: Optional[str] = None
) -> Tuple[bool, str, Path]:
    """
    Uruchamia kompilację 2 razy (bezpieczeństwo referencji/bib).
    Zwraca: (ok, log_text, pdf_path).
    """
    logger_obj = logging.getLogger(logger) if logger else logging.getLogger()
    logger_obj.info("Compiling %s with %s engine two passes", tex_filename, tex_engine)
    logs: List[str] = []
    pdf_path = workdir / (Path(tex_filename).stem + ".pdf")

    # PASS 1
    rc1, log1 = await _run(_build_cmd(tex_engine, tex_filename), workdir, timeout_sec, logger)
    logs.append(log1)
    if rc1 != 0:
        return False, "".join(logs), pdf_path

    # PASS 2 (często już nic nie robi, ale nie szkodzi)
    rc2, log2 = await _run(_build_cmd(tex_engine, tex_filename), workdir, timeout_sec, logger)
    logs.append(log2)
    ok = (rc2 == 0) and pdf_path.exists()
    return ok, "".join(logs), pdf_path

# ===== Kompilacje: pojedynczy TEX + ZIP =====

async def compile_single_tex_bytes(
    tex_bytes: bytes,
    filename: str,
    tex_engine: str,
    timeout_sec: int,
    logger: Optional[str] = None,
    support_files: Optional[Iterable[Tuple[str, bytes]]] = None,
) -> Dict[str, Optional[object]]:
    """
    Kompiluje pojedynczy .tex + dowolne pliki pomocnicze (obrazy, sty/cls/bib, inne .tex).
    Każdy plik zapisujemy w katalogu roboczym widocznym dla kompilatora.
    """
    log = logging.getLogger(logger) if logger else logging.getLogger()
    log.info("Compiling %s with %s engine single tex", filename, tex_engine)
    if not is_safe_tex_name(filename):
        return {"ok": False, "pdf_bytes": None, "log_text": "Invalid .tex filename", "pdf_name": None}

    with tempfile.TemporaryDirectory(prefix="latex_one_") as td:
        wd = Path(td)
        (wd / Path(filename).name).write_bytes(tex_bytes)

        # Zapisz wszystkie supporty obok głównego pliku
        if support_files:
            for fname, data in support_files:
                if not fname or not is_safe_asset_name(fname):
                    continue
                # spłaszczamy do nazwy bazowej (bez podkatalogów z zewn. źródeł)
                (wd / Path(fname).name).write_bytes(data)

        ok, log_text, pdf_path = await run_latex_two_passes(Path(filename).name, wd, tex_engine, timeout_sec, logger)
        log.info("LaTeX compile finished: %s", "OK" if ok else "FAIL")
        log.info("Log text:\n%s", log_text[-2000:])  # Log last 2000 chars
        log.info("PDF path: %s", pdf_path)
        if ok and pdf_path.exists():
            return {
                "ok": True,
                "pdf_bytes": pdf_path.read_bytes(),
                "log_text": log_text,
                "pdf_name": pdf_path.name,
            }
        return {"ok": False, "pdf_bytes": None, "log_text": log_text, "pdf_name": pdf_path.name}

async def compile_zip_to_zip(
    zip_bytes: bytes, tex_engine: str, timeout_sec: int, logger: Optional[str] = None
) -> Tuple[bytes, List[str]]:
    """
    Rozpakowuje ZIP, znajduję wszystkie .tex i dla KAŻDEGO tworzy osobny katalog roboczy
    ze skopiowanym CAŁYM drzewem, by zachować ścieżki względne obrazków/plików.
    Zwraca (zip_pdf_bytes, lista_rel_sciezek_które_padły).
    """
    log = logging.getLogger(logger) if logger else logging.getLogger()
    failed: List[str] = []
    out_pdf_paths: List[Path] = []

    with tempfile.TemporaryDirectory(prefix="latex_zip_") as td:
        td_path = Path(td)
        src_root = td_path / "src"
        src_root.mkdir(parents=True, exist_ok=True)

        # 1) Rozpakuj
        zpath = td_path / "in.zip"
        zpath.write_bytes(zip_bytes)
        with zipfile.ZipFile(zpath, "r") as zf:
            zf.extractall(src_root)

        # 2) Zbierz .tex
        tex_files = [p for p in src_root.rglob("*.tex") if is_safe_tex_name(p.name)]
        if not tex_files:
            return b"", ["No .tex files found in archive."]

        # 3) Kompiluj każdy .tex w izolacji (kopiujemy całe drzewo do osobnego workdiru)
        for tex_path in tex_files:
            rel_tex = tex_path.relative_to(src_root).as_posix()
            work = td_path / f"build_{tex_path.stem}"
            try:
                shutil.copytree(src_root, work / "src")
                ok, log_text, pdf = await run_latex_two_passes(rel_tex, work / "src", tex_engine, timeout_sec, logger)
                if ok and pdf.exists():
                    out_pdf_paths.append(pdf)
                else:
                    failed.append(rel_tex)
                    log.info("LaTeX FAIL %s\n%s", rel_tex, (log_text or "")[-2000:])
            except Exception:
                # let it crash – log + re-raise (to zobaczy wyższa warstwa/discord.py)
                log.exception("Exception while compiling %s", rel_tex)
                raise

        # 4) Spakuj wyjściowe PDF-y
        out_zip = td_path / "compiled_pdfs.zip"
        with zipfile.ZipFile(out_zip, "w", compression=zipfile.ZIP_DEFLATED) as zf:
            for pdf in out_pdf_paths:
                zf.write(pdf, arcname=pdf.name)

        return out_zip.read_bytes(), failed