Artifex/tools/runtime.py

from __future__ import annotations

import subprocess
import os
import re
import shutil
from dataclasses import dataclass
from pathlib import Path

from tools.capabilities import Capability


@dataclass(frozen=True)
class ToolResult:
    returncode: int
    stdout: str
    stderr: str


@dataclass(frozen=True)
class MutationResult:
    success: bool
    operation: str
    files_changed: list[str]
    hunks_applied: int = 0
    error: str = ""


class CapabilityError(PermissionError):
    pass


class WorktreeTools:
    def __init__(self, worktree_path: Path, capabilities: set[Capability]) -> None:
        self.worktree_path = worktree_path.resolve()
        self.capabilities = capabilities

    def _require(self, capability: Capability) -> None:
        if capability not in self.capabilities:
            raise CapabilityError(f"Missing capability: {capability.value}")

    def _safe_path(self, relative_path: str) -> Path:
        path = (self.worktree_path / relative_path).resolve()
        if self.worktree_path not in path.parents and path != self.worktree_path:
            raise CapabilityError("Path escapes worktree")
        return path

    def read_text(self, relative_path: str) -> str:
        self._require(Capability.READ_REPOSITORY)
        return self._safe_path(relative_path).read_text(encoding="utf-8")

    def list_directory(self, relative_path: str = ".") -> list[str]:
        self._require(Capability.INVESTIGATE_WORKTREE)
        path = self._safe_path(relative_path)
        if not path.is_dir():
            raise CapabilityError("Path is not a directory")
        entries: list[str] = []
        for child in sorted(path.iterdir(), key=lambda item: item.name):
            if child.name == ".git" or child.name == "__pycache__":
                continue
            suffix = "/" if child.is_dir() else ""
            entries.append(f"{child.relative_to(self.worktree_path).as_posix()}{suffix}")
        return entries

    def read_file(self, relative_path: str, max_chars: int = 12000) -> str:
        self._require(Capability.INVESTIGATE_WORKTREE)
        path = self._safe_path(relative_path)
        if not path.is_file():
            raise CapabilityError("Path is not a file")
        return path.read_text(encoding="utf-8")[:max_chars]

    def search_code(self, pattern: str, include: str = "*.py", max_results: int = 50) -> list[dict[str, object]]:
        self._require(Capability.INVESTIGATE_WORKTREE)
        regex = re.compile(pattern)
        results: list[dict[str, object]] = []
        for path in self.worktree_path.rglob(include):
            if ".git" in path.parts or "__pycache__" in path.parts or not path.is_file():
                continue
            relative = path.relative_to(self.worktree_path).as_posix()
            for number, line in enumerate(path.read_text(encoding="utf-8", errors="ignore").splitlines(), start=1):
                if regex.search(line):
                    results.append({"path": relative, "line": number, "text": line[:300]})
                    if len(results) >= max_results:
                        return results
        return results

    def find_symbol(self, symbol: str, max_results: int = 50) -> list[dict[str, object]]:
        escaped = re.escape(symbol)
        return self.search_code(rf"^\s*(class|def)\s+{escaped}\b|\b{escaped}\b", max_results=max_results)

    def write_text(self, relative_path: str, content: str) -> None:
        self._require(Capability.WRITE_WORKTREE)
        path = self._safe_path(relative_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(content, encoding="utf-8")

    def write_file(self, relative_path: str, content: str) -> MutationResult:
        try:
            self.write_text(relative_path, content)
        except Exception as exc:
            return MutationResult(False, "write_file", [], error=str(exc))
        return MutationResult(True, "write_file", [relative_path])

    def create_directory(self, relative_path: str) -> MutationResult:
        self._require(Capability.WRITE_WORKTREE)
        try:
            path = self._safe_path(relative_path)
            if path.exists() and not path.is_dir():
                return MutationResult(False, "create_directory", [], error="Target exists as a file")
            path.mkdir(parents=True, exist_ok=True)
        except Exception as exc:
            return MutationResult(False, "create_directory", [], error=str(exc))
        return MutationResult(True, "create_directory", [relative_path])

    def delete_file(self, relative_path: str) -> MutationResult:
        self._require(Capability.WRITE_WORKTREE)
        try:
            path = self._safe_path(relative_path)
            if not path.exists():
                return MutationResult(False, "delete_file", [], error="Target file does not exist")
            if path.is_dir():
                return MutationResult(False, "delete_file", [], error="Target is a directory")
            path.unlink()
        except Exception as exc:
            return MutationResult(False, "delete_file", [], error=str(exc))
        return MutationResult(True, "delete_file", [relative_path])

    def move_file(self, source: str, destination: str, overwrite: bool = False) -> MutationResult:
        self._require(Capability.WRITE_WORKTREE)
        try:
            source_path = self._safe_path(source)
            destination_path = self._safe_path(destination)
            if not source_path.exists():
                return MutationResult(False, "move_file", [], error="Source file does not exist")
            if source_path.is_dir():
                return MutationResult(False, "move_file", [], error="Source is a directory")
            if destination_path.exists() and not overwrite:
                return MutationResult(False, "move_file", [], error="Destination already exists")
            destination_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(str(source_path), str(destination_path))
        except Exception as exc:
            return MutationResult(False, "move_file", [], error=str(exc))
        return MutationResult(True, "move_file", [source, destination])

    def apply_patch(self, relative_path: str, patch: str) -> MutationResult:
        self._require(Capability.WRITE_WORKTREE)
        try:
            path = self._safe_path(relative_path)
            if not path.is_file():
                return MutationResult(False, "apply_patch", [], error="Target file does not exist")
            original = path.read_text(encoding="utf-8").splitlines(keepends=True)
            updated, hunks = self._apply_unified_patch(original, patch)
            path.write_text("".join(updated), encoding="utf-8")
        except Exception as exc:
            return MutationResult(False, "apply_patch", [], error=str(exc))
        return MutationResult(True, "apply_patch", [relative_path], hunks_applied=hunks)

    def _apply_unified_patch(self, original: list[str], patch: str) -> tuple[list[str], int]:
        patch_lines = patch.splitlines(keepends=True)
        output: list[str] = []
        source_index = 0
        index = 0
        hunks = 0
        while index < len(patch_lines):
            line = patch_lines[index]
            if not line.startswith("@@"):
                index += 1
                continue
            match = re.match(r"@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@", line)
            if not match:
                raise ValueError(f"Invalid hunk header: {line.strip()}")
            old_start = int(match.group(1)) - 1
            index += 1
            hunk_lines: list[str] = []
            while index < len(patch_lines) and not patch_lines[index].startswith("@@"):
                hunk_lines.append(patch_lines[index])
                index += 1
            old_lines = self._patch_old_lines(hunk_lines)
            actual_start = self._locate_hunk(original, old_lines, old_start, source_index)
            if actual_start is None:
                raise ValueError("Patch context does not match target file")
            if actual_start < source_index:
                raise ValueError("Overlapping or out-of-order patch hunk")
            output.extend(original[source_index:actual_start])
            source_index = actual_start
            hunks += 1
            for hunk_line in hunk_lines:
                if hunk_line.startswith(("---", "+++")):
                    continue
                if hunk_line.startswith(" "):
                    expected = hunk_line[1:]
                    if source_index >= len(original) or original[source_index] != expected:
                        raise ValueError("Patch context does not match target file")
                    output.append(original[source_index])
                    source_index += 1
                elif hunk_line.startswith("-"):
                    expected = hunk_line[1:]
                    if source_index >= len(original) or original[source_index] != expected:
                        raise ValueError("Patch removal does not match target file")
                    source_index += 1
                elif hunk_line.startswith("+"):
                    output.append(hunk_line[1:])
                elif hunk_line.startswith("\\ No newline at end of file"):
                    pass
                elif hunk_line.strip() == "":
                    raise ValueError("Blank patch lines must be prefixed with context/add/remove marker")
                else:
                    raise ValueError(f"Invalid hunk line: {hunk_line[:40]}")
        if hunks == 0:
            raise ValueError("Patch contains no hunks")
        output.extend(original[source_index:])
        return output, hunks

    def _patch_old_lines(self, hunk_lines: list[str]) -> list[str]:
        old_lines: list[str] = []
        for hunk_line in hunk_lines:
            if hunk_line.startswith(("---", "+++", "+", "\\ No newline at end of file")):
                continue
            if hunk_line.startswith((" ", "-")):
                old_lines.append(hunk_line[1:])
            elif hunk_line.strip() == "":
                raise ValueError("Blank patch lines must be prefixed with context/add/remove marker")
            else:
                raise ValueError(f"Invalid hunk line: {hunk_line[:40]}")
        return old_lines

    def _locate_hunk(self, original: list[str], old_lines: list[str], old_start: int, source_index: int) -> int | None:
        if not old_lines:
            return max(old_start, source_index) if old_start >= source_index else None
        if self._matches_at(original, old_lines, old_start):
            return old_start
        window_start = max(source_index, old_start - 200)
        window_end = min(len(original), old_start + 200)
        for candidate in range(window_start, window_end + 1):
            if self._matches_at(original, old_lines, candidate):
                return candidate
        for candidate in range(source_index, len(original) + 1):
            if self._matches_at(original, old_lines, candidate):
                return candidate
        return None

    def _matches_at(self, original: list[str], old_lines: list[str], start: int) -> bool:
        if start < 0 or start + len(old_lines) > len(original):
            return False
        return original[start : start + len(old_lines)] == old_lines

    def run(self, command: list[str], timeout: int = 60) -> ToolResult:
        self._require(Capability.RUN_TESTS)
        env = os.environ.copy()
        for key in ["DJANGO_SETTINGS_MODULE", "DATABASE_URL"]:
            env.pop(key, None)
        completed = subprocess.run(
            command,
            cwd=self.worktree_path,
            check=False,
            capture_output=True,
            text=True,
            timeout=timeout,
            env=env,
        )
        return ToolResult(completed.returncode, completed.stdout, completed.stderr)

    def git(self, args: list[str]) -> ToolResult:
        required = Capability.COMMIT_CHANGES if args and args[0] == "commit" else Capability.READ_REPOSITORY
        self._require(required)
        completed = subprocess.run(
            ["git", *args],
            cwd=self.worktree_path,
            check=False,
            capture_output=True,
            text=True,
        )
        return ToolResult(completed.returncode, completed.stdout, completed.stderr)

    def status(self) -> str:
        return self.git(["status", "--short"]).stdout

    def diff(self) -> str:
        return self.git(["diff", "--", "."]).stdout

    def git_status(self) -> str:
        self._require(Capability.INVESTIGATE_WORKTREE)
        return self.status()

    def git_diff(self) -> str:
        self._require(Capability.INVESTIGATE_WORKTREE)
        return self.diff()

    def commit_all(self, message: str) -> str:
        self._require(Capability.COMMIT_CHANGES)
        add = self.git(["add", "."])
        if add.returncode != 0:
            raise RuntimeError(add.stderr)
        commit = self.git(["-c", "user.name=Artifex", "-c", "user.email=artifex@example.invalid", "commit", "-m", message])
        if commit.returncode != 0:
            raise RuntimeError(commit.stderr or commit.stdout)
        rev = self.git(["rev-parse", "HEAD"])
        if rev.returncode != 0:
            raise RuntimeError(rev.stderr)
        return rev.stdout.strip()