conjurer/latex_functions.py

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

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

from ai_functions import handle_response
import shutil
from typing import Iterable, Tuple

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_asset_name(name: str) -> bool:
    return bool(SAFE_ASSET_NAME.match(name or ""))


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
    )


def _build_cmd(engine: str, tex_filename: str) -> list[str]:
    base = (engine or "").strip().lower()
    if base in ("pdflatex", "xelatex", "lualatex"):
        return [
            base,
            "-interaction=nonstopmode",
            "-halt-on-error",
            "-file-line-error",
            "-no-shell-escape",
            tex_filename,
        ]
    if base == "latexmk":
        return [
            "latexmk",
            "-pdf",
            "-interaction=nonstopmode",
            "-halt-on-error",
            "-file-line-error",
            tex_filename,
        ]
    if base == "tectonic":
        return [
            "tectonic",
            "-X",
            "compile",
            "--keep-logs",
            "--outdir",
            ".",
            tex_filename,
        ]
    return [
        "pdflatex",
        "-interaction=nonstopmode",
        "-halt-on-error",
        "-file-line-error",
        "-no-shell-escape",
        tex_filename,
    ]


async def run_latex_two_passes(
    tex_filename: str,
    workdir: Path,
    tex_engine: str,
    timeout_sec: int,
    logger: str = None,
):
    logger = logging.getLogger(logger) if logger else None

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

    logs = []
    cmd = _build_cmd(tex_engine, tex_filename)
    if logger:
        logger.debug("LaTeX cmd: %s (cwd=%s)", " ".join(cmd), workdir)

    rc1, out1 = await run_once(cmd)
    logs.append(out1)
    second_pass = not ((tex_engine or "").strip().lower() == "tectonic")

    if rc1 != 0:
        log_text = "".join(logs)
        pdf = workdir / (Path(tex_filename).stem + ".pdf")
        return False, log_text, pdf

    if second_pass:
        rc2, out2 = await run_once(cmd)
        logs.append(out2)
        ok = rc2 == 0
    else:
        ok = rc1 == 0

    log_text = "".join(logs)
    pdf = workdir / (Path(tex_filename).stem + ".pdf")
    ok = ok and pdf.exists()
    return ok, log_text, pdf

async def ask_openai_diagnosis(
    log_text: str, tex_excerpt: str, model: str, logger: str = None
) -> str:
    logger = logging.getLogger(logger) if logger else None
    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"
    )
    logger.debug("LaTeX DIAG (%s)\n%s", model, log_text[-9000:])
    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