conjurer/latex_functions.py

# latex_functions.py
# -*- coding: utf-8 -*-
from __future__ import annotations

import asyncio
import re
import shlex
import tempfile
import zipfile
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from ai_functions import handle_response

SAFE_TEX_NAME = re.compile(r"^[\w\-. ]+\.tex$", re.IGNORECASE)


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


def is_tex_attachment(att, max_mb: int) -> bool:
    """Filtruje discord.Attachment: tylko .tex + limit MB + nazwa bezpieczna."""
    return bool(
        getattr(att, "filename", "")
        and is_safe_tex_name(att.filename)
        and getattr(att, "size", 0) <= max_mb * 1024 * 1024
    )


async def run_latex_two_passes(
    tex_filename: str, workdir: Path, tex_engine: str, timeout_sec: int, logger=None
) -> Tuple[bool, str, Path]:
    cmd = f"{tex_engine} -interaction=nonstopmode -halt-on-error -file-line-error -no-shell-escape {shlex.quote(tex_filename)}"
    if logger:
        logger.debug("LaTeX cmd: %s (cwd=%s)", cmd, workdir)

    async def _run_once() -> Tuple[int, str]:
        proc = await asyncio.create_subprocess_shell(
            cmd,
            cwd=str(workdir),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
        try:
            out_b, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout_sec)
            out_t = (out_b or b"").decode("utf-8", errors="replace")
            return proc.returncode or 0, out_t
        except asyncio.TimeoutError:
            try:
                proc.kill()
            except Exception:
                pass
            return 124, f"[Timeout] Compilation exceeded {timeout_sec}s.\n"

    logs: List[str] = []
    rc1, out1 = await _run_once()
    logs.append(out1)
    if rc1 != 0:
        log_text = "".join(logs)
        if logger:
            logger.debug("LaTeX pass1 FAIL %s\n%s", tex_filename, log_text[-2000:])
        return False, log_text, workdir / (Path(tex_filename).stem + ".pdf")

    rc2, out2 = await _run_once()
    logs.append(out2)
    log_text = "".join(logs)
    pdf = workdir / (Path(tex_filename).stem + ".pdf")
    ok = (rc2 == 0) and pdf.exists()
    if logger:
        logger.debug("LaTeX pass2 %s ok=%s\n%s", tex_filename, ok, log_text[-2000:])
    return ok, log_text, pdf


async def compile_single_tex_bytes(
    tex_bytes: bytes, filename: str, tex_engine: str, timeout_sec: int, logger=None
) -> Dict[str, Optional[object]]:
    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)
        tex_path = wd / filename
        tex_path.write_bytes(tex_bytes)
        ok, log_text, pdf_path = await run_latex_two_passes(
            filename, wd, tex_engine, timeout_sec, logger=logger
        )
        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=None
) -> Tuple[bytes, List[str]]:
    failed: List[str] = []
    out_pdf_paths: List[Path] = []
    with tempfile.TemporaryDirectory(prefix="latex_zip_") as td:
        td_path = Path(td)
        in_zip = td_path / "in.zip"
        in_zip.write_bytes(zip_bytes)
        extract_dir = td_path / "in"
        extract_dir.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(in_zip, "r") as zf:
            zf.extractall(extract_dir)

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

        for tex in tex_files:
            try:
                work = td_path / f"build_{tex.stem}"
                work.mkdir(parents=True, exist_ok=True)
                target_tex = work / tex.name
                target_tex.write_bytes(tex.read_bytes())
                ok, log_text, pdf_path = await run_latex_two_passes(
                    target_tex.name, work, tex_engine, timeout_sec, logger=logger
                )
                if ok and pdf_path.exists():
                    out_pdf_paths.append(pdf_path)
                else:
                    failed.append(tex.name)
                    if logger:
                        logger.debug("BATCH FAIL %s\n%s", tex.name, log_text[-1500:])
            except Exception as e:
                failed.append(f"{tex.name} (exception: {e})")
                if logger:
                    logger.debug("BATCH EXC %s: %s", tex.name, e)

        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


async def ask_openai_diagnosis(
    log_text: str, tex_excerpt: str, api_key: str, model: str, logger=None
) -> str:
    sys = (
        "You are a LaTeX build diagnostician. Analyze pdflatex log and return:\n"
        "1) Root causes (bullets)\n2) Minimal working fix (code block)\n3) Alternative fix"
    )
    usr = f"---LOG---\n{log_text[-9000:]}\n---END LOG---"
    if tex_excerpt:
        usr += f"\n---TEX EXCERPT---\n{tex_excerpt[:4000]}\n---END EXCERPT---"
        prompt = [{"role": sys, "content": "none"}, {"role": "user", "content": usr}]

        response, _ = await handle_response(
            prompt,
            None,
            None,
            prompt,
            "latex_diagnosis",
            "NONE",
            algorithm=model,
            none_request=prompt,
        )
    return response