zerolang_editing/zero_tools.py

"""Path-based Zerolang compiler tools for the editing environment."""

from __future__ import annotations

import hashlib
import json
import os
import platform
import re
import shutil
import subprocess
import tempfile
import threading
import urllib.request
from pathlib import Path
from typing import Any


_ZERO_INSTALL_LOCK = threading.Lock()


def _download(url: str, timeout: int = 60) -> bytes:
    with urllib.request.urlopen(url, timeout=timeout) as response:
        return response.read()


def _zero_asset_candidates() -> list[str]:
    system = platform.system()
    machine = platform.machine().lower()
    if machine in {"arm64", "aarch64"}:
        cpu = "arm64"
    elif machine in {"x86_64", "amd64"}:
        cpu = "x64"
    else:
        return []

    if system == "Darwin":
        return [f"zero-darwin-{cpu}"]
    if system == "Linux":
        return [f"zero-linux-musl-{cpu}", f"zero-linux-{cpu}"]
    return []


def _install_zero_binary() -> str | None:
    install_dir = Path(
        os.environ.get("ZERO_INSTALL_DIR")
        or Path(tempfile.gettempdir()) / "zerolang-editing-zero" / "bin"
    ).expanduser()
    binary = install_dir / "zero"
    if binary.exists():
        return str(binary)

    with _ZERO_INSTALL_LOCK:
        if binary.exists():
            return str(binary)

        base_url = os.environ.get(
            "ZERO_DOWNLOAD_BASE_URL",
            "https://github.com/vercel-labs/zero/releases/latest/download",
        ).rstrip("/")
        checksums_text = _download(f"{base_url}/CHECKSUMS.txt").decode()
        checksums = {}
        for line in checksums_text.splitlines():
            parts = line.split()
            if len(parts) >= 2:
                checksums[parts[1]] = parts[0]

        install_dir.mkdir(parents=True, exist_ok=True)
        last_error: Exception | None = None
        for asset in _zero_asset_candidates():
            try:
                data = _download(f"{base_url}/{asset}")
                expected = checksums.get(asset)
                actual = hashlib.sha256(data).hexdigest()
                if expected and actual != expected:
                    raise RuntimeError(f"checksum mismatch for {asset}")
                binary.write_bytes(data)
                os.chmod(binary, 0o755)
                check = subprocess.run(
                    [str(binary), "--version"],
                    text=True,
                    capture_output=True,
                    timeout=10,
                )
                if check.returncode == 0:
                    return str(binary)
            except Exception as exc:
                last_error = exc
                if binary.exists():
                    binary.unlink()
                continue
        if last_error is not None:
            raise RuntimeError(f"failed to install zero binary: {last_error}") from last_error
    return None


def _zero_binary(zero_path: str | None = None) -> str | None:
    candidates = [
        zero_path,
        shutil.which("zero"),
        str(Path.home() / ".zero" / "bin" / "zero"),
    ]
    for candidate in candidates:
        if candidate and Path(candidate).exists():
            return candidate
    return _install_zero_binary()


def _json_tool_result(result: dict[str, Any]) -> str:
    return json.dumps(result, indent=2, sort_keys=True)


def read_source(path: str | Path) -> str:
    return Path(path).read_text()


def _source_fingerprint(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {"path": str(path), "exists": False}
    data = path.read_bytes()
    return {
        "path": str(path),
        "exists": True,
        "bytes": len(data),
        "source_sha256": hashlib.sha256(data).hexdigest(),
    }


def _summarize_graph_dump(graph_dump: str) -> dict[str, Any]:
    summary: dict[str, Any] = {
        "hash": None,
        "literals": [],
        "functions": [],
        "calls": [],
        "identifiers": [],
    }
    for line in graph_dump.splitlines():
        if line.startswith("hash "):
            summary["hash"] = line.split('"', 2)[1]
        elif " Literal " in line:
            match = re.match(r'node (#[0-9a-f]+) Literal type:"([^"]+)" value:"(.*)"', line)
            if match:
                summary["literals"].append(
                    {"node": match.group(1), "type": match.group(2), "value": match.group(3)}
                )
        elif " Function " in line:
            match = re.match(r'node (#[0-9a-f]+) Function name:"([^"]+)" type:"([^"]+)"', line)
            if match:
                summary["functions"].append(
                    {"node": match.group(1), "name": match.group(2), "type": match.group(3)}
                )
        elif " MethodCall " in line:
            match = re.match(r'node (#[0-9a-f]+) MethodCall name:"([^"]+)" type:"([^"]+)"', line)
            if match:
                summary["calls"].append(
                    {"node": match.group(1), "name": match.group(2), "type": match.group(3)}
                )
        elif " Identifier " in line:
            match = re.match(r'node (#[0-9a-f]+) Identifier name:"([^"]+)"', line)
            if match:
                summary["identifiers"].append({"node": match.group(1), "name": match.group(2)})
    return summary


def run_zero_path(args: list[str], path: str | Path, zero_path: str | None = None) -> dict[str, Any]:
    binary = _zero_binary(zero_path)
    source_path = Path(path)
    if binary is None:
        return {
            "ok": False,
            "tool_error": "zero binary not found; install with https://zerolang.ai/install.sh",
            **_source_fingerprint(source_path),
        }
    if not source_path.exists():
        return {
            "ok": False,
            "tool_error": f"source file does not exist: {source_path}",
            **_source_fingerprint(source_path),
        }

    proc = subprocess.run(
        [binary, *args, str(source_path)],
        text=True,
        capture_output=True,
        timeout=10,
        env={**os.environ, "PATH": f"{Path(binary).parent}:{os.environ.get('PATH', '')}"},
    )
    return {
        "ok": proc.returncode == 0,
        "returncode": proc.returncode,
        "stdout": proc.stdout[-12000:],
        "stderr": proc.stderr[-4000:],
        **_source_fingerprint(source_path),
    }


def run_zero_source(
    args: list[str], source: str, zero_path: str | None = None
) -> dict[str, Any]:
    with tempfile.TemporaryDirectory(prefix="zerolang-editing-score-") as tmp:
        source_path = Path(tmp) / "program.0"
        source_path.write_text(source)
        return run_zero_path(args, source_path, zero_path)


def make_zero_tools(zero_path: str | None = None) -> list[Any]:
    def zero_check(path: str) -> str:
        """Run `zero check --json` on a `.0` file path on disk."""
        return _json_tool_result(run_zero_path(["check", "--json"], path, zero_path))

    def zero_graph_summary(path: str) -> str:
        """Return compact graph hash and patchable node facts for a `.0` file path."""
        result = run_zero_path(["graph", "dump"], path, zero_path)
        if result.get("ok"):
            result["summary"] = _summarize_graph_dump(result.get("stdout", ""))
        return _json_tool_result(result)

    def zero_graph_dump(path: str) -> str:
        """Run `zero graph dump` on a `.0` file path on disk."""
        return _json_tool_result(run_zero_path(["graph", "dump"], path, zero_path))

    def zero_graph_json(path: str) -> str:
        """Run `zero graph --json` on a `.0` file path on disk."""
        return _json_tool_result(run_zero_path(["graph", "--json"], path, zero_path))

    def zero_fix_plan(path: str) -> str:
        """Run `zero fix --plan --json` on a `.0` file path on disk."""
        return _json_tool_result(run_zero_path(["fix", "--plan", "--json"], path, zero_path))

    def zero_graph_patch(path: str, expect_graph_hash: str, op: str) -> str:
        """Apply one checked `zero graph patch` operation to a `.0` file path on disk."""
        binary = _zero_binary(zero_path)
        source_path = Path(path)
        if binary is None:
            return _json_tool_result(
                {
                    "ok": False,
                    "tool_error": "zero binary not found; install with https://zerolang.ai/install.sh",
                    **_source_fingerprint(source_path),
                }
            )
        if not source_path.exists():
            return _json_tool_result(
                {
                    "ok": False,
                    "tool_error": f"source file does not exist: {source_path}",
                    **_source_fingerprint(source_path),
                }
            )
        proc = subprocess.run(
            [
                binary,
                "graph",
                "patch",
                str(source_path),
                "--expect-graph-hash",
                expect_graph_hash,
                "--op",
                op,
            ],
            text=True,
            capture_output=True,
            timeout=10,
            env={**os.environ, "PATH": f"{Path(binary).parent}:{os.environ.get('PATH', '')}"},
        )
        return _json_tool_result(
            {
                "ok": proc.returncode == 0,
                "returncode": proc.returncode,
                "stdout": proc.stdout[-12000:],
                "stderr": proc.stderr[-4000:],
                **_source_fingerprint(source_path),
            }
        )

    def zero_skills_get(skill: str) -> str:
        """Return version-matched Zerolang guidance for `language`, `diagnostics`, `stdlib`, or `zero`."""
        binary = _zero_binary(zero_path)
        if binary is None:
            return _json_tool_result(
                {
                    "ok": False,
                    "tool_error": "zero binary not found; install with https://zerolang.ai/install.sh",
                }
            )
        proc = subprocess.run(
            [binary, "skills", "get", skill],
            text=True,
            capture_output=True,
            timeout=10,
            env={**os.environ, "PATH": f"{Path(binary).parent}:{os.environ.get('PATH', '')}"},
        )
        return _json_tool_result(
            {
                "ok": proc.returncode == 0,
                "returncode": proc.returncode,
                "stdout": proc.stdout[-12000:],
                "stderr": proc.stderr[-4000:],
            }
        )

    return [
        zero_check,
        zero_graph_summary,
        zero_graph_dump,
        zero_graph_json,
        zero_fix_plan,
        zero_graph_patch,
        zero_skills_get,
    ]