AlgoTune/algotune_agent

"""Encode AlgoTune Agent's multi-file workspace as one Galapagos candidate artifact.

Galapagos evaluates one candidate file. AlgoTune agents can create helper modules, build metadata,
and Cython/Pythran/DaCe sources. A small import-time bootstrap materializes those sidecars before
the editable ``solver.py`` body executes, preserving the upstream workspace model without changing
the task evaluator contract.
"""
from __future__ import annotations

import ast
import base64
import hashlib
import json
import re
import warnings
from pathlib import PurePosixPath

from .validator import tampering_error

_MANIFEST_PREFIX = "# ALGOTUNE-WORKSPACE-V1:"
_SOLVER_MARKER = "# ALGOTUNE-SOLVER-PY-V1"
_SAFE_FILE = re.compile(r"^[A-Za-z0-9_./-]+$")


def validate_filename(name: str) -> str:
    name = str(name).strip().replace("\\", "/")
    path = PurePosixPath(name)
    if (
        not name
        or not _SAFE_FILE.fullmatch(name)
        or path.is_absolute()
        or ".." in path.parts
        or any(part.startswith(".algotune_") for part in path.parts)
        or name.endswith("/")
    ):
        raise ValueError(
            f"Invalid filename {name!r}; use a relative path containing letters, numbers, "
            "dots, dashes, underscores, or slashes"
        )
    return path.as_posix()


def unpack_workspace(candidate: str) -> dict[str, str]:
    """Return virtual files; an ordinary candidate is simply ``solver.py``."""
    lines = candidate.splitlines()
    if not lines or not lines[0].startswith(_MANIFEST_PREFIX):
        return {"solver.py": candidate}
    if len(lines) < 2 or not lines[1].startswith("import base64 as _at_b64"):
        # A solver is allowed to contain arbitrary comments, including one that resembles our
        # manifest sentinel. Only generated candidates carry the bootstrap signature too.
        return {"solver.py": candidate}
    try:
        encoded = lines[0][len(_MANIFEST_PREFIX) :].strip()
        payload = json.loads(base64.urlsafe_b64decode(encoded.encode()).decode())
        if not isinstance(payload, dict):
            raise ValueError("workspace manifest is not an object")
        files = {validate_filename(k): str(v) for k, v in payload.items()}
        # Current bundles keep solver.py in the manifest and compile it as an independent module
        # source. This preserves module docstrings and ``from __future__`` placement even though a
        # bootstrap has to run first. Keep reading the earlier development format so checkpoints
        # written while the port was being tested remain resumable.
        if "solver.py" in files and _SOLVER_MARKER not in lines:
            return files
        marker = lines.index(_SOLVER_MARKER)
        solver = "\n".join(lines[marker + 1 :])
        if candidate.endswith("\n"):
            solver += "\n"
        files["solver.py"] = solver
        return files
    except Exception as exc:  # noqa: BLE001 - malformed model-authored artifact
        raise ValueError(f"Malformed AlgoTune workspace bundle: {exc}") from exc


def _bootstrap(files: dict[str, str]) -> str:
    payload = base64.urlsafe_b64encode(
        json.dumps(files, sort_keys=True, separators=(",", ":")).encode()
    ).decode()
    digest = hashlib.sha256(payload.encode()).hexdigest()[:16]
    # This code executes inside the task sandbox, before its stopwatch starts. That is the same
    # place upstream AlgoTune pays setup/Cython compilation cost.
    return f'''{_MANIFEST_PREFIX} {payload}
import base64 as _at_b64, json as _at_json, os as _at_os
import pathlib as _at_path, shutil as _at_shutil, subprocess as _at_subprocess, sys as _at_sys
import zipfile as _at_zipfile
_at_files = _at_json.loads(_at_b64.urlsafe_b64decode({payload!r}).decode())
_at_root = _at_path.Path("/tmp/algotune_workspace_{digest}")
_at_root.mkdir(parents=True, exist_ok=True)
for _at_name, _at_text in _at_files.items():
    _at_target = _at_root / _at_name
    _at_target.parent.mkdir(parents=True, exist_ok=True)
    if not _at_target.exists() or _at_target.read_text() != _at_text:
        _at_target.write_text(_at_text)
_at_pythran = [
    _at_name for _at_name, _at_text in _at_files.items()
    if _at_name.endswith(".py") and "pythran export" in _at_text.lower()
]
_at_dace = [
    _at_name for _at_name, _at_text in _at_files.items()
    if _at_name.endswith(".py") and "@dace.program" in _at_text
]
_at_build = "setup.py" if "setup.py" in _at_files else (
    "pyproject.toml" if "pyproject.toml" in _at_files else ""
)
if (_at_build or _at_pythran or _at_dace) and not (_at_root / ".algotune_build_ok").exists():
    if "setup.py" in _at_files:
        try:
            _at_proc = _at_subprocess.run(
                [_at_sys.executable, "setup.py", "build_ext", "--inplace"], cwd=_at_root,
                capture_output=True, text=True, timeout=1800,
            )
        except _at_subprocess.TimeoutExpired as _at_exc:
            raise RuntimeError("Cython compilation failed: timed out after 1800 seconds") from _at_exc
        if _at_proc.returncode:
            raise RuntimeError("Cython compilation failed: " + (_at_proc.stderr or _at_proc.stdout)[-2000:])
    elif "pyproject.toml" in _at_files:
        _at_wheel_dir = _at_root / ".algotune_wheels"
        _at_wheel_dir.mkdir(exist_ok=True)
        try:
            _at_proc = _at_subprocess.run(
                [_at_sys.executable, "-m", "pip", "wheel", ".", "--no-deps",
                 "--no-build-isolation", "--wheel-dir", str(_at_wheel_dir)],
                cwd=_at_root, capture_output=True, text=True, timeout=1800,
            )
        except _at_subprocess.TimeoutExpired as _at_exc:
            raise RuntimeError("Cython compilation failed: timed out after 1800 seconds") from _at_exc
        if _at_proc.returncode:
            raise RuntimeError("Cython compilation failed: " + (_at_proc.stderr or _at_proc.stdout)[-2000:])
        for _at_wheel in _at_wheel_dir.glob("*.whl"):
            with _at_zipfile.ZipFile(_at_wheel) as _at_archive:
                _at_archive.extractall(_at_root)
    if _at_build:
        _at_extensions = [
            _at_extension
            for _at_extension in list(_at_root.glob("**/*.so")) + list(_at_root.glob("**/*.pyd"))
            if "build" not in _at_extension.relative_to(_at_root).parts
        ]
        if not _at_extensions:
            _at_extensions = list((_at_root / "build").glob("**/*.so")) + list(
                (_at_root / "build").glob("**/*.pyd")
            )
            for _at_extension in _at_extensions:
                _at_shutil.copy2(_at_extension, _at_root / _at_extension.name)
        if not _at_extensions:
            raise RuntimeError("Cython compilation failed: build produced no .so/.pyd extension")
    for _at_name in _at_pythran:
        _at_compile_target = _at_root / _at_name
        try:
            _at_proc = _at_subprocess.run(
                [_at_sys.executable, "-m", "pythran", "-O3", "-march=native",
                 _at_compile_target.name], cwd=_at_compile_target.parent,
                capture_output=True, text=True, timeout=300,
            )
        except _at_subprocess.TimeoutExpired as _at_exc:
            raise RuntimeError("Pythran compilation failed: timed out after 300 seconds") from _at_exc
        if _at_proc.returncode:
            raise RuntimeError("Pythran compilation failed: " + (_at_proc.stderr or _at_proc.stdout)[-2000:])
    for _at_name in _at_dace:
        _at_dace_target = _at_root / _at_name
        try:
            _at_proc = _at_subprocess.run(
                [_at_sys.executable, "-c", "import " + _at_dace_target.stem],
                cwd=_at_dace_target.parent, capture_output=True, text=True, timeout=300,
            )
        except _at_subprocess.TimeoutExpired as _at_exc:
            raise RuntimeError("DaCe compilation failed: timed out after 300 seconds") from _at_exc
        if _at_proc.returncode:
            raise RuntimeError("DaCe compilation failed: " + (_at_proc.stderr or _at_proc.stdout)[-2000:])
    (_at_root / ".algotune_build_ok").write_text("ok")
if str(_at_root) not in _at_sys.path:
    _at_sys.path.insert(0, str(_at_root))
_at_solver_path = _at_root / "solver.py"
_at_solver_source = _at_solver_path.read_text()
del _at_b64, _at_json, _at_os, _at_path, _at_shutil, _at_subprocess, _at_sys, _at_zipfile
del _at_files
del _at_name, _at_text, _at_target, _at_pythran, _at_dace, _at_build
exec(compile(_at_solver_source, str(_at_solver_path), "exec"), globals(), globals())
'''


def pack_workspace(files: dict[str, str]) -> str:
    normalized = {validate_filename(name): str(content) for name, content in files.items()}
    solver = normalized.get("solver.py", "")
    if (
        not any(name != "solver.py" for name in normalized)
        and not solver.startswith(_MANIFEST_PREFIX)
    ):
        return solver
    normalized["solver.py"] = solver
    return _bootstrap(normalized)


def syntax_error(files: dict[str, str]) -> str | None:
    """Apply the upstream edit-time critical Python linter gate."""
    for name, source in files.items():
        if not name.endswith(".py"):
            continue
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", SyntaxWarning)
                tree = ast.parse(source, filename=name)
        except SyntaxError as exc:
            return f"{name}:{exc.lineno}:{exc.offset}: {exc.msg}"
        violation = tampering_error(tree, source)
        if violation:
            return f"{name}: {violation}"
    return None


def numbered(source: str, start_line: int = 1, limit: int = 100) -> str:
    lines = source.splitlines()
    start = max(1, int(start_line))
    selected = lines[start - 1 : start - 1 + limit]
    return "\n".join(f"{index:4d}: {line}" for index, line in enumerate(selected, start))


def replace_lines(source: str, start: int, end: int, content: str) -> str:
    lines = source.splitlines(keepends=True)
    cleaned = content.lstrip(":") if content else ""
    replacement = [line + "\n" for line in cleaned.splitlines()]
    if start < 0 or end < 0 or (start > 0 and end < start):
        raise ValueError(f"invalid line range {start}-{end}")
    if start == 0:
        result = replacement + lines[end:]
    else:
        if start > len(lines) + 1:
            raise ValueError(
                f"start line {start} is greater than the file length ({len(lines)}) + 1"
            )
        result = lines[: start - 1] + replacement + lines[end:]
    return "".join(result)


def delete_lines(source: str, start: int, end: int) -> str:
    if start < 1 or end < start:
        raise ValueError(f"invalid delete range {start}-{end}")
    return replace_lines(source, start, end, "")