Spaces:

SagarMythos
/

sunee3

Sleeping

File size: 40,795 Bytes

df6cd5e

"""
core/zip_handler.py — Zip Intelligence helpers (v4.5 — Line-Locked Edition).

Features
--------
- safely extract uploaded ZIP archives into an isolated workspace
- inspect file structure and text-file contents for AI context building
- apply AI-suggested changes back onto the extracted project using THREE
  strategies (in strict priority order):

    1. **Line-Locked Patch blocks** (```linepatch``` fences) — surgical edits
       bounded by an explicit ``LINES: <start>-<end>`` header. Any content the
       model produces outside that range is REJECTED. This is the preferred
       format for Phase-1 reliability.
    2. **Unified Diff / Git patch** blocks — applied via ``git apply``.
    3. **# FILE: full-file rewrite** blocks — kept for backwards compatibility
       but now flagged with ``strategy="file_blocks"`` so the pipeline can
       downgrade its trust score.

- recompress the modified project into a downloadable ZIP.
"""

from __future__ import annotations

import os
import re
import shutil
import subprocess
import tempfile
import zipfile
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

IGNORE_DIRS = {
    ".git", ".venv", "venv", "node_modules", "__pycache__",
    ".idea", ".vscode", "dist", "build", ".next", ".cache",
    "__MACOSX",
}

TEXT_EXTENSIONS = {
    ".py", ".js", ".ts", ".tsx", ".jsx", ".java", ".go", ".rs",
    ".c", ".cpp", ".h", ".hpp", ".rb", ".php", ".sh", ".yml",
    ".yaml", ".json", ".md", ".html", ".css", ".sql", ".toml",
    ".cfg", ".ini", ".txt", ".env", ".dockerfile", ".xml",
}

_CODE_BLOCK_RE = re.compile(r"```(?P<lang>[\w.+-]*)\n(?P<body>.*?)```", re.S)
_FILE_HEADER_RE = re.compile(r"^[#/;!\- ]*FILE:\s*([^\n\r]+)", re.I | re.M)
_DIFF_PLUS_RE = re.compile(r"^\+\+\+\s+(?:b/)?(.+)$", re.M)

# ---------------------------------------------------------------------------
# Line-Locked patch format
# ---------------------------------------------------------------------------
# Expected block (fence language: ``linepatch``, ``line-patch``, or ``patch-lines``):
#
#     ```linepatch
#     FILE: relative/path/to/file.py
#     LINES: 42-57
#     ---
#     <verbatim replacement content for lines 42..57 inclusive>
#     ```
#
# Multiple linepatch blocks may appear in one response, even against the same
# file. They are applied bottom-up (highest starting line first) so earlier
# line numbers remain valid.

_LINEPATCH_LANGS = {"linepatch", "line-patch", "patch-lines", "surgical"}
_LP_FILE_RE = re.compile(r"^\s*FILE:\s*(?P<path>[^\n\r]+)\s*$", re.I | re.M)
_LP_LINES_RE = re.compile(
    r"^\s*LINES:\s*(?P<start>\d+)\s*-\s*(?P<end>\d+)\s*$",
    re.I | re.M,
)
_LP_SEPARATOR_RE = re.compile(r"^\s*---+\s*$", re.M)


def _safe_rel_path(rel_path: str) -> str:
    rel = str(rel_path or "").replace("\\", "/").strip().lstrip("/")
    rel = re.sub(r"/+", "/", rel)
    if not rel or rel in {".", ".."}:
        raise ValueError("empty relative path")
    parts = [p for p in rel.split("/") if p not in {"", "."}]
    if any(part == ".." for part in parts):
        raise ValueError(f"unsafe path: {rel_path}")
    return "/".join(parts)


def _collect_candidate_modified_files(text: str) -> List[str]:
    candidates: List[str] = []
    for block in _extract_linepatch_blocks(text):
        path = block.get("path")
        if path and path not in candidates:
            candidates.append(path)
    for block in _extract_diff_blocks(text):
        for path in _DIFF_PLUS_RE.findall(block):
            raw = (path or "").strip()
            if not raw or raw == "/dev/null":
                continue
            try:
                clean = _safe_rel_path(raw)
            except ValueError:
                continue
            if clean not in candidates:
                candidates.append(clean)
    for path, _content in _extract_file_blocks(text):
        if path not in candidates:
            candidates.append(path)
    return candidates


def _derive_allowed_dirs(allowed_files: Sequence[str]) -> List[str]:
    dirs = set()
    for rel in allowed_files:
        clean = _safe_rel_path(rel)
        parent = clean.rsplit("/", 1)[0] if "/" in clean else ""
        if parent:
            dirs.add(parent)
    return sorted(dirs)


def _enforce_diff_guard(text: str, allowed_files: Optional[Sequence[str]] = None) -> Optional[Dict[str, Any]]:
    if not allowed_files:
        return None
    normalized_files = []
    for rel in allowed_files:
        try:
            clean = _safe_rel_path(rel)
        except ValueError:
            continue
        if clean not in normalized_files:
            normalized_files.append(clean)
    if not normalized_files:
        return None

    allowed_set = set(normalized_files)
    allowed_dirs = _derive_allowed_dirs(normalized_files)
    touched = _collect_candidate_modified_files(text)
    blocked: List[str] = []
    for rel in touched:
        if rel in allowed_set:
            continue
        if any(rel.startswith(directory + "/") for directory in allowed_dirs):
            continue
        blocked.append(rel)
    if not blocked:
        return None
    return {
        "applied": False,
        "strategy": "diff_guard",
        "modified_files": [],
        "errors": [
            "Diff Guard rejected patch outside GraphMapper scope: " + ", ".join(blocked)
        ],
        "blocked_files": blocked,
        "allowed_scope": normalized_files,
    }


def _looks_like_text(path: Path, sample_size: int = 2048) -> bool:
    ext = path.suffix.lower()
    if ext in TEXT_EXTENSIONS:
        return True
    try:
        with open(path, "rb") as fh:
            sample = fh.read(sample_size)
        if b"\x00" in sample:
            return False
        sample.decode("utf-8")
        return True
    except Exception:
        return False


def _iter_project_files(project_root: str) -> Iterable[Tuple[str, Path]]:
    root = Path(project_root).resolve()
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if d not in IGNORE_DIRS]
        for filename in sorted(filenames):
            abs_path = Path(dirpath) / filename
            rel_path = abs_path.relative_to(root).as_posix()
            yield rel_path, abs_path


def _detect_archive_root(zf: zipfile.ZipFile) -> str:
    top_levels = set()
    for name in zf.namelist():
        clean = name.replace("\\", "/").strip("/")
        if not clean or clean.startswith("__MACOSX/"):
            continue
        parts = clean.split("/")
        if parts and parts[0]:
            top_levels.add(parts[0])
        if len(top_levels) > 1:
            return ""
    return next(iter(top_levels), "") if len(top_levels) == 1 else ""


def extract_uploaded_zip(
    zip_path: str,
    workspace: Optional[str] = None,
    project_id: Optional[str] = None,
) -> Dict[str, Any]:
    """Extract an uploaded ZIP into an isolated temp directory.

    Returns metadata needed for later previewing and re-zipping.
    """
    if not zip_path or not os.path.exists(zip_path):
        raise FileNotFoundError(f"zip not found: {zip_path}")

    base_workspace = workspace or os.path.join(tempfile.gettempdir(), "devai_zipintelligence")
    os.makedirs(base_workspace, exist_ok=True)

    prefix = f"{project_id}_" if project_id else "zipintel_"
    extraction_dir = tempfile.mkdtemp(prefix=prefix, dir=base_workspace)

    with zipfile.ZipFile(zip_path, "r") as zf:
        archive_root_prefix = _detect_archive_root(zf)
        zf.extractall(extraction_dir)

    project_root = os.path.join(extraction_dir, archive_root_prefix) if archive_root_prefix else extraction_dir
    project_root = os.path.abspath(project_root)
    if not os.path.isdir(project_root):
        raise FileNotFoundError("ZIP extracted, but no project root directory was found")

    return {
        "zip_path": os.path.abspath(zip_path),
        "zip_name": os.path.basename(zip_path),
        "extract_dir": os.path.abspath(extraction_dir),
        "project_root": project_root,
        "archive_root_prefix": archive_root_prefix,
    }


def read_project_file(project_root: str, relative_path: str, max_chars: int = 12000) -> str:
    rel = _safe_rel_path(relative_path)
    abs_path = (Path(project_root).resolve() / rel).resolve()
    root = Path(project_root).resolve()
    if root not in abs_path.parents and abs_path != root:
        raise ValueError(f"path escapes project root: {relative_path}")
    if not abs_path.exists() or not abs_path.is_file():
        raise FileNotFoundError(relative_path)
    if not _looks_like_text(abs_path):
        return "[binary or non-text file omitted]"
    with open(abs_path, "r", encoding="utf-8", errors="ignore") as fh:
        text = fh.read(max_chars + 1)
    if len(text) > max_chars:
        return text[:max_chars] + "\n...[truncated]"
    return text


def list_project_files(project_root: str, max_files: int = 1000) -> List[str]:
    files: List[str] = []
    for rel_path, _ in _iter_project_files(project_root):
        files.append(rel_path)
        if len(files) >= max_files:
            break
    return files


def build_file_tree(project_root: str, max_entries: int = 250) -> str:
    lines = ["."]
    count = 0
    root = Path(project_root).resolve()
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = sorted([d for d in dirnames if d not in IGNORE_DIRS])
        rel_dir = Path(dirpath).resolve().relative_to(root).as_posix()
        depth = 0 if rel_dir == "." else rel_dir.count("/") + 1
        base_indent = "  " * depth
        if rel_dir != ".":
            lines.append(f"{base_indent}{Path(dirpath).name}/")
            count += 1
            if count >= max_entries:
                lines.append("  ...[tree truncated]")
                break
        for filename in sorted(filenames):
            rel_file = (Path(dirpath) / filename).resolve().relative_to(root).as_posix()
            lines.append(f"{base_indent}  {filename}")
            count += 1
            if count >= max_entries:
                lines.append("  ...[tree truncated]")
                return "\n".join(lines)
    return "\n".join(lines)


def build_zip_ai_context(
    project_root: str,
    focus_files: Optional[Sequence[str]] = None,
    max_files: int = 18,
    per_file_chars: int = 8000,
    max_total_chars: int = 70000,
    include_line_numbers: bool = True,
) -> str:
    """Build the context blob shown to the AI.

    When ``include_line_numbers`` is True (the v4.5 default), each source file
    is rendered with a ``NNN: `` prefix on every line so the model can address
    edits by exact line numbers — the foundation of Line-Locking.
    """
    files = list_project_files(project_root, max_files=1200)
    ordered: List[str] = []
    seen = set()

    for rel in focus_files or []:
        if rel in files and rel not in seen:
            ordered.append(rel)
            seen.add(rel)

    for rel in files:
        if rel not in seen:
            ordered.append(rel)
            seen.add(rel)

    sections = [
        "## ZIP Intelligence — Project file tree",
        "```text",
        build_file_tree(project_root),
        "```",
        "",
        "## ZIP Intelligence — Readable file contents (with line numbers)"
        if include_line_numbers
        else "## ZIP Intelligence — Readable file contents",
    ]

    total = sum(len(s) for s in sections)
    file_count = 0
    for rel in ordered:
        abs_path = Path(project_root) / rel
        if not abs_path.is_file() or not _looks_like_text(abs_path):
            continue
        snippet = read_project_file(project_root, rel, max_chars=per_file_chars)
        if include_line_numbers:
            numbered_lines = []
            for idx, line in enumerate(snippet.splitlines(), start=1):
                numbered_lines.append(f"{idx:>4}: {line}")
            snippet = "\n".join(numbered_lines)
        block = f"### FILE: {rel}\n```\n{snippet}\n```\n"
        if total + len(block) > max_total_chars and file_count > 0:
            sections.append("\n...[additional files omitted to stay within context budget]")
            break
        sections.append(block)
        total += len(block)
        file_count += 1
        if file_count >= max_files:
            sections.append("\n...[file limit reached]\n")
            break

    if file_count == 0:
        sections.append("_(No readable text files found in the uploaded ZIP.)_")

    return "\n".join(sections).strip()


def preview_project_contents(
    project_root: str,
    max_files: int = 30,
    per_file_chars: int = 3000,
    max_total_chars: int = 90000,
) -> str:
    return build_zip_ai_context(
        project_root=project_root,
        focus_files=None,
        max_files=max_files,
        per_file_chars=per_file_chars,
        max_total_chars=max_total_chars,
        include_line_numbers=False,
    )


# ---------------------------------------------------------------------------
# Strategy 1 — Line-Locked patch blocks (surgical edit)
# ---------------------------------------------------------------------------

def _extract_linepatch_blocks(text: str) -> List[Dict[str, Any]]:
    """Return a list of parsed linepatch blocks.

    Each entry: ``{"path": str, "start": int, "end": int, "content": str}``.
    Malformed blocks are silently skipped; callers can compare the number of
    returned entries against the number of ``linepatch`` fences to detect
    invalid submissions.
    """
    blocks: List[Dict[str, Any]] = []
    for match in _CODE_BLOCK_RE.finditer(text or ""):
        lang = (match.group("lang") or "").strip().lower()
        if lang not in _LINEPATCH_LANGS:
            continue
        body = match.group("body") or ""

        file_m = _LP_FILE_RE.search(body)
        lines_m = _LP_LINES_RE.search(body)
        if not file_m or not lines_m:
            continue

        try:
            path = _safe_rel_path(file_m.group("path").strip())
        except ValueError:
            continue

        start = int(lines_m.group("start"))
        end = int(lines_m.group("end"))
        if start < 1 or end < start:
            continue

        # Find where the header ends. Prefer an explicit `---` separator; if
        # missing, take everything after the last header line.
        header_end = max(file_m.end(), lines_m.end())
        sep_m = _LP_SEPARATOR_RE.search(body, header_end)
        content_start = sep_m.end() if sep_m else header_end
        # Skip a single leading newline so the payload starts cleanly.
        content = body[content_start:]
        if content.startswith("\n"):
            content = content[1:]
        # Strip trailing whitespace-only newline; keep meaningful indentation.
        content = content.rstrip("\n")

        blocks.append({
            "path": path,
            "start": start,
            "end": end,
            "content": content,
        })
    return blocks


def _apply_linepatch_blocks(project_root: str, text: str) -> Dict[str, Any]:
    """Apply Line-Locked patches to real files on disk.

    * Refuses to touch lines outside the declared ``LINES:`` range.
    * If the target file has fewer lines than ``end``, the patch is REJECTED
      for that block (returned in ``errors``).
    * Multiple blocks against the same file are applied bottom-up so earlier
      line numbers stay valid after each splice.
    """
    parsed = _extract_linepatch_blocks(text)
    if not parsed:
        return {
            "applied": False,
            "strategy": "line_locked",
            "modified_files": [],
            "patch_details": [],
            "errors": ["no linepatch blocks found"],
        }

    root = Path(project_root).resolve()

    # Group by file so we can sort bottom-up per file.
    per_file: Dict[str, List[Dict[str, Any]]] = {}
    for blk in parsed:
        per_file.setdefault(blk["path"], []).append(blk)

    modified_files: List[str] = []
    patch_details: List[Dict[str, Any]] = []
    errors: List[str] = []

    for rel_path, blocks in per_file.items():
        try:
            abs_path = (root / rel_path).resolve()
        except Exception as exc:
            errors.append(f"{rel_path}: bad path ({exc})")
            continue
        if root not in abs_path.parents and abs_path != root:
            errors.append(f"{rel_path}: escapes project root")
            continue
        if not abs_path.exists() or not abs_path.is_file():
            errors.append(f"{rel_path}: file not found — line-locked patches "
                          f"can only edit existing files")
            continue

        try:
            original_text = abs_path.read_text(encoding="utf-8", errors="ignore")
        except Exception as exc:
            errors.append(f"{rel_path}: read failed ({exc})")
            continue

        original_lines = original_text.splitlines(keepends=False)
        total_lines = len(original_lines)

        # Validate every block against the CURRENT (unmodified) line count.
        valid_blocks: List[Dict[str, Any]] = []
        for blk in blocks:
            if blk["start"] > total_lines or blk["end"] > total_lines:
                errors.append(
                    f"{rel_path}: LINES {blk['start']}-{blk['end']} out of "
                    f"range (file has {total_lines} lines) — patch rejected"
                )
                continue
            valid_blocks.append(blk)

        if not valid_blocks:
            continue

        # Detect overlapping ranges within the same file — reject the file
        # entirely to preserve the surgical-editing guarantee.
        sorted_by_start = sorted(valid_blocks, key=lambda b: b["start"])
        for i in range(1, len(sorted_by_start)):
            prev = sorted_by_start[i - 1]
            cur = sorted_by_start[i]
            if cur["start"] <= prev["end"]:
                errors.append(
                    f"{rel_path}: overlapping linepatch ranges "
                    f"({prev['start']}-{prev['end']} vs "
                    f"{cur['start']}-{cur['end']}) — file skipped"
                )
                valid_blocks = []
                break

        if not valid_blocks:
            continue

        # Apply bottom-up so line indices for earlier blocks stay valid.
        new_lines = list(original_lines)
        applied_ranges: List[Tuple[int, int, int]] = []  # (start, end, new_len)
        for blk in sorted(valid_blocks, key=lambda b: b["start"], reverse=True):
            start_idx = blk["start"] - 1        # inclusive, 0-based
            end_idx = blk["end"]                # exclusive slice bound
            replacement = blk["content"].splitlines(keepends=False)
            new_lines[start_idx:end_idx] = replacement
            applied_ranges.append((blk["start"], blk["end"], len(replacement)))

        # Preserve trailing newline behaviour of the original file.
        new_text = "\n".join(new_lines)
        if original_text.endswith("\n") and not new_text.endswith("\n"):
            new_text += "\n"

        try:
            abs_path.write_text(new_text, encoding="utf-8")
        except Exception as exc:
            errors.append(f"{rel_path}: write failed ({exc})")
            continue

        modified_files.append(rel_path)
        for orig_start, orig_end, new_len in sorted(applied_ranges):
            patch_details.append({
                "file": rel_path,
                "original_range": [orig_start, orig_end],
                "original_line_count": orig_end - orig_start + 1,
                "replacement_line_count": new_len,
            })

    return {
        "applied": bool(modified_files),
        "strategy": "line_locked",
        "modified_files": modified_files,
        "patch_details": patch_details,
        "errors": errors,
    }


# ---------------------------------------------------------------------------
# Strategy 2 — Unified diff / git patch
# ---------------------------------------------------------------------------

def _extract_diff_blocks(text: str) -> List[str]:
    blocks: List[str] = []
    for match in _CODE_BLOCK_RE.finditer(text or ""):
        lang = (match.group("lang") or "").strip().lower()
        body = match.group("body") or ""
        if lang in {"diff", "patch"} and body.strip():
            blocks.append(body)
    if not blocks and ("diff --git" in (text or "") or ("@@" in (text or "") and "+++" in (text or ""))):
        blocks.append(text)
    return blocks


def _apply_diff_blocks(project_root: str, text: str) -> Dict[str, Any]:
    diff_blocks = _extract_diff_blocks(text)
    if not diff_blocks:
        return {
            "applied": False,
            "strategy": "unified_diff",
            "modified_files": [],
            "errors": ["no unified diff found"],
        }

    if not shutil.which("git"):
        return {
            "applied": False,
            "strategy": "unified_diff",
            "modified_files": [],
            "errors": ["`git` binary unavailable — cannot apply unified diff"],
        }

    patch_text = "\n\n".join(block.strip() for block in diff_blocks if block.strip()).strip() + "\n"
    patch_file = tempfile.NamedTemporaryFile("w", suffix=".patch", delete=False, encoding="utf-8")
    try:
        patch_file.write(patch_text)
        patch_file.close()
        proc = subprocess.run(
            ["git", "-C", project_root, "apply", "--whitespace=nowarn", "--reject", patch_file.name],
            capture_output=True,
            text=True,
        )
        if proc.returncode != 0:
            return {
                "applied": False,
                "strategy": "unified_diff",
                "modified_files": [],
                "errors": [proc.stderr.strip() or proc.stdout.strip() or "git apply failed"],
            }
        modified = [
            _safe_rel_path(path.strip())
            for path in _DIFF_PLUS_RE.findall(patch_text)
            if path.strip() and path.strip() != "/dev/null"
        ]
        deduped = list(dict.fromkeys(modified))
        return {
            "applied": True,
            "strategy": "unified_diff",
            "modified_files": deduped,
            "errors": [],
        }
    finally:
        try:
            os.unlink(patch_file.name)
        except OSError:
            pass


# ---------------------------------------------------------------------------
# Strategy 3 — # FILE: full-file rewrite (fallback / new-file support)
# ---------------------------------------------------------------------------

def _extract_file_blocks(text: str) -> List[Tuple[str, str]]:
    files: List[Tuple[str, str]] = []
    for match in _CODE_BLOCK_RE.finditer(text or ""):
        lang = (match.group("lang") or "").strip().lower()
        # Skip specialised fences so a linepatch/diff never doubles as a file block.
        if lang in _LINEPATCH_LANGS or lang in {"diff", "patch"}:
            continue
        body = match.group("body") or ""
        header = _FILE_HEADER_RE.search(body)
        if not header:
            continue
        try:
            path = _safe_rel_path(header.group(1).strip())
        except ValueError:
            continue
        content = body[header.end():].lstrip("\n")
        files.append((path, content))
    return files


def _apply_file_blocks(project_root: str, text: str) -> Dict[str, Any]:
    parsed = _extract_file_blocks(text)
    modified_files: List[str] = []
    if not parsed:
        return {
            "applied": False,
            "strategy": "file_blocks",
            "modified_files": modified_files,
            "errors": ["no # FILE blocks found"],
        }

    root = Path(project_root).resolve()
    for rel_path, content in parsed:
        abs_path = (root / rel_path).resolve()
        if root not in abs_path.parents and abs_path != root:
            raise ValueError(f"unsafe write target: {rel_path}")
        abs_path.parent.mkdir(parents=True, exist_ok=True)
        with open(abs_path, "w", encoding="utf-8", errors="ignore") as fh:
            fh.write(content)
        modified_files.append(rel_path)

    return {
        "applied": True,
        "strategy": "file_blocks",
        "modified_files": modified_files,
        "errors": [],
    }


# ---------------------------------------------------------------------------
# Public entry point: apply_ai_changes
# ---------------------------------------------------------------------------

def apply_ai_changes(
    project_root: str,
    generated_text: str,
    allowed_files: Optional[Sequence[str]] = None,
) -> Dict[str, Any]:
    """Apply AI-generated project modifications back onto the extracted ZIP.

    Priority (v4.5):
        1. ```linepatch``` blocks — surgical, line-range-locked edits.
        2. ```diff``` / ```patch``` blocks — unified diff via ``git apply``.
        3. ``# FILE:`` fenced blocks — full-file rewrite (least trusted;
           should only be used for brand-new files).

    When ``allowed_files`` is provided, Diff Guard blocks any patch that tries
    to touch files outside the GraphMapper-derived edit scope.
    """
    text = generated_text or ""
    if not text.strip():
        return {
            "applied": False,
            "strategy": "none",
            "modified_files": [],
            "errors": ["empty generated output"],
        }

    guard = _enforce_diff_guard(text, allowed_files=allowed_files)
    if guard is not None:
        return guard

    # 1. Line-locked surgical patches — the preferred v4.5 path.
    linepatch_result = _apply_linepatch_blocks(project_root, text)
    if linepatch_result.get("applied"):
        return linepatch_result

    # 2. Unified diff / git patches.
    diff_result = _apply_diff_blocks(project_root, text)
    if diff_result.get("applied"):
        return diff_result

    # 3. Full-file rewrite blocks (compatibility fallback).
    file_block_result = _apply_file_blocks(project_root, text)
    if file_block_result.get("applied"):
        return file_block_result

    errors = []
    errors.extend(linepatch_result.get("errors") or [])
    errors.extend(diff_result.get("errors") or [])
    errors.extend(file_block_result.get("errors") or [])
    return {
        "applied": False,
        "strategy": "none",
        "modified_files": [],
        "errors": errors or ["no applicable project modifications found"],
    }


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def describe_bundle(bundle: Dict[str, Any], project_root: Optional[str] = None) -> Dict[str, Any]:
    """Return a lightweight, JSON-safe summary of an extracted ZIP bundle.

    Used by the Gradio UI (Code Graph tab) to show file counts, language mix,
    and top-level entries without dumping the whole preview blob a second time.
    """
    root = project_root or bundle.get("project_root") or ""
    summary: Dict[str, Any] = {
        "zip_name": bundle.get("zip_name", ""),
        "archive_root_prefix": bundle.get("archive_root_prefix", ""),
        "project_root": root,
        "file_count": 0,
        "text_file_count": 0,
        "languages": {},
        "top_level": [],
    }
    if not root or not os.path.isdir(root):
        return summary

    languages: Dict[str, int] = {}
    file_count = 0
    text_count = 0
    for rel_path, abs_path in _iter_project_files(root):
        file_count += 1
        ext = abs_path.suffix.lower() or "(none)"
        languages[ext] = languages.get(ext, 0) + 1
        if _looks_like_text(abs_path):
            text_count += 1

    top_level: List[str] = []
    try:
        for entry in sorted(os.listdir(root)):
            if entry in IGNORE_DIRS:
                continue
            top_level.append(entry + ("/" if os.path.isdir(os.path.join(root, entry)) else ""))
    except OSError:
        pass

    summary.update(
        {
            "file_count": file_count,
            "text_file_count": text_count,
            "languages": dict(sorted(languages.items(), key=lambda kv: -kv[1])[:12]),
            "top_level": top_level[:40],
        }
    )
    return summary


def cleanup_extraction(bundle_or_path: Any) -> bool:
    """Safely remove a previously extracted ZIP workspace.

    Accepts either the bundle dict returned by :func:`extract_uploaded_zip`
    or a raw ``extract_dir`` path string. Returns ``True`` when a directory
    was actually removed. Only touches directories under the OS temp dir or
    the DevAI workspace to avoid accidental recursive deletes.
    """
    if not bundle_or_path:
        return False
    if isinstance(bundle_or_path, dict):
        path = bundle_or_path.get("extract_dir") or bundle_or_path.get("project_root") or ""
    else:
        path = str(bundle_or_path)
    if not path:
        return False
    path = os.path.abspath(path)

    safe_prefixes = [
        os.path.abspath(tempfile.gettempdir()),
        os.path.abspath(os.path.join(tempfile.gettempdir(), "devai_zipintelligence")),
        os.path.abspath(os.environ.get("WORKSPACE_DIR", "/tmp/devai_workspace")),
    ]
    if not any(path.startswith(prefix + os.sep) or path == prefix for prefix in safe_prefixes):
        # Refuse to touch anything outside our known sandboxes.
        return False
    if not os.path.isdir(path):
        return False
    try:
        shutil.rmtree(path, ignore_errors=True)
        return not os.path.isdir(path)
    except Exception:
        return False


def recompress_project_zip(
    project_root: str,
    output_dir: str,
    original_zip_name: Optional[str] = None,
    suffix: str = "zip_intelligence",
    archive_root_prefix: str = "",
) -> str:
    os.makedirs(output_dir, exist_ok=True)
    base_name = os.path.splitext(original_zip_name or "project.zip")[0]
    out_path = os.path.join(output_dir, f"{base_name}_{suffix}.zip")

    root = Path(project_root).resolve()
    with zipfile.ZipFile(out_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for rel_path, abs_path in _iter_project_files(str(root)):
            if archive_root_prefix:
                arcname = f"{archive_root_prefix.rstrip('/')}/{rel_path}"
            else:
                arcname = rel_path
            zf.write(abs_path, arcname)
    return out_path


# ---------------------------------------------------------------------------
# v4.6 — Atomic Multi-file Patcher
# ---------------------------------------------------------------------------
#
# Contract
# --------
# If a task requires changes across N files, we MUST either:
#
#   * apply every patch successfully   → commit (all-or-nothing SUCCESS), OR
#   * revert the entire project to its pre-patch state on ANY failure → full
#     rollback (all-or-nothing FAILURE).
#
# The atomic wrapper works by snapshotting every file that appears in the
# generated patch (plus their parent directory state so newly-created files
# can be cleanly removed) BEFORE dispatching to the standard strategies
# implemented in :func:`apply_ai_changes`. On failure, the snapshot is
# restored byte-for-byte and any brand-new files are deleted.
#
# The snapshot is stored on disk under a per-transaction ``.devai_snapshot``
# directory inside ``project_root`` so we survive re-entrant calls. The
# directory is always removed at the end (success or failure).


class AtomicPatchError(RuntimeError):
    """Raised when an atomic multi-file patch could not be fully applied.

    The exception is intentionally lightweight: the recommended failure
    channel is the ``result`` dict returned by
    :func:`apply_ai_changes_atomic`, which already contains ``rolled_back``
    and ``errors`` fields.
    """


def _collect_touched_files(text: str) -> List[str]:
    """Return every relative path referenced by ANY patch strategy.

    Used as the snapshot scope for atomic transactions.
    """
    return _collect_candidate_modified_files(text)


def _snapshot_files(
    project_root: str,
    files: Sequence[str],
) -> Dict[str, Any]:
    """Take a byte-level snapshot of *files* under *project_root*.

    Returns a manifest describing which files existed and where their
    pre-patch contents were stored on disk. The manifest is consumed by
    :func:`_rollback_from_snapshot`.
    """
    root = Path(project_root).resolve()
    snap_dir = Path(tempfile.mkdtemp(prefix="devai_atomic_", dir=str(root)))
    manifest: Dict[str, Any] = {
        "snapshot_dir": str(snap_dir),
        "root": str(root),
        "entries": [],  # list of {"rel": str, "existed": bool, "backup": str|None}
    }

    seen: set = set()
    for raw in files:
        try:
            rel = _safe_rel_path(raw)
        except ValueError:
            continue
        if rel in seen:
            continue
        seen.add(rel)

        abs_path = (root / rel).resolve()
        if root not in abs_path.parents and abs_path != root:
            # Refuse to snapshot anything outside the project.
            continue

        entry: Dict[str, Any] = {"rel": rel, "existed": abs_path.is_file(), "backup": None}
        if entry["existed"]:
            backup_path = snap_dir / rel
            backup_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(abs_path, backup_path)
            entry["backup"] = str(backup_path)
        manifest["entries"].append(entry)
    return manifest


def _rollback_from_snapshot(manifest: Dict[str, Any]) -> List[str]:
    """Restore every file in *manifest* to its pre-patch state.

    Returns the list of relative paths that were rolled back (either
    restored or deleted).
    """
    root = Path(manifest.get("root") or "").resolve()
    entries = manifest.get("entries") or []
    rolled_back: List[str] = []
    for entry in entries:
        rel = entry.get("rel")
        if not rel:
            continue
        abs_path = (root / rel).resolve()
        if root not in abs_path.parents and abs_path != root:
            continue

        if entry.get("existed"):
            backup = entry.get("backup")
            if backup and os.path.isfile(backup):
                try:
                    abs_path.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(backup, abs_path)
                    rolled_back.append(rel)
                except Exception:
                    # Best-effort rollback — surface the failure via the
                    # caller's ``errors`` channel rather than raising here.
                    pass
        else:
            # The file was CREATED by this transaction — delete it if it
            # now exists on disk.
            if abs_path.exists():
                try:
                    abs_path.unlink()
                    rolled_back.append(rel)
                    # Prune empty parent directories that we implicitly
                    # created (do NOT touch pre-existing folders).
                    parent = abs_path.parent
                    while parent != root and parent.is_dir() and not any(parent.iterdir()):
                        parent.rmdir()
                        parent = parent.parent
                except Exception:
                    pass
    return rolled_back


def _discard_snapshot(manifest: Dict[str, Any]) -> None:
    snap_dir = manifest.get("snapshot_dir") if isinstance(manifest, dict) else None
    if snap_dir and os.path.isdir(snap_dir):
        shutil.rmtree(snap_dir, ignore_errors=True)


def _verify_all_touched_files_landed(
    project_root: str,
    result: Dict[str, Any],
    expected: Sequence[str],
) -> Optional[str]:
    """Confirm every file the AI *claimed* to modify actually exists.

    A patch may report ``applied=True`` yet skip individual files (e.g. when
    a linepatch range is out-of-bounds). For atomic guarantees we require
    every declared file to end up on disk after the transaction.
    """
    reported = set(result.get("modified_files") or [])
    missing: List[str] = []
    for rel in expected:
        try:
            clean = _safe_rel_path(rel)
        except ValueError:
            missing.append(rel)
            continue
        if clean not in reported:
            missing.append(clean)
            continue
        abs_path = Path(project_root) / clean
        if not abs_path.is_file():
            missing.append(clean)
    if missing:
        return (
            "atomic guard: the following files were declared in the patch "
            "but did not land on disk — " + ", ".join(sorted(set(missing)))
        )
    return None


def apply_ai_changes_atomic(
    project_root: str,
    generated_text: str,
    allowed_files: Optional[Sequence[str]] = None,
) -> Dict[str, Any]:
    """Atomic multi-file variant of :func:`apply_ai_changes`.

    Behaviour
    ---------
    1. Enumerate every file referenced by the generated patch (across all
       three strategies — linepatch, diff, and ``# FILE:`` blocks).
    2. Snapshot each of those files (and record whether they existed).
    3. Delegate to :func:`apply_ai_changes`.
    4. If ``applied`` is False, or any declared file did not land on disk,
       or an unexpected exception occurred, restore the snapshot and mark
       ``rolled_back=True``.
    5. Otherwise, discard the snapshot and return the successful result
       augmented with ``atomic=True`` and ``rolled_back=False``.

    Returned dict extends the standard :func:`apply_ai_changes` shape with:

        ``atomic``:      ``True`` — this response came from the atomic path.
        ``rolled_back``: whether a rollback was executed.
        ``touched_files``: files referenced by the patch (snapshot scope).
        ``rolled_back_files``: files whose state was actually reverted.
    """
    text = generated_text or ""
    if not text.strip():
        return {
            "applied": False,
            "strategy": "none",
            "modified_files": [],
            "errors": ["empty generated output"],
            "atomic": True,
            "rolled_back": False,
            "touched_files": [],
            "rolled_back_files": [],
        }

    # Diff Guard runs first — no snapshot needed if the patch is out of scope.
    guard = _enforce_diff_guard(text, allowed_files=allowed_files)
    if guard is not None:
        guard.update({
            "atomic": True,
            "rolled_back": False,
            "touched_files": _collect_touched_files(text),
            "rolled_back_files": [],
        })
        return guard

    touched = _collect_touched_files(text)
    manifest = _snapshot_files(project_root, touched)

    try:
        result = apply_ai_changes(project_root, text, allowed_files=allowed_files)
    except Exception as exc:  # pragma: no cover — defensive
        rolled = _rollback_from_snapshot(manifest)
        _discard_snapshot(manifest)
        return {
            "applied": False,
            "strategy": "atomic_rollback",
            "modified_files": [],
            "errors": [f"atomic dispatch raised: {exc}"],
            "atomic": True,
            "rolled_back": True,
            "touched_files": touched,
            "rolled_back_files": rolled,
        }

    if not result.get("applied"):
        rolled = _rollback_from_snapshot(manifest)
        _discard_snapshot(manifest)
        result_out = dict(result)
        result_out.update({
            "atomic": True,
            "rolled_back": True,
            "touched_files": touched,
            "rolled_back_files": rolled,
        })
        return result_out

    # Multi-file completeness check: every declared file must have landed.
    completeness_error = _verify_all_touched_files_landed(
        project_root, result, touched
    )
    if completeness_error is not None:
        rolled = _rollback_from_snapshot(manifest)
        _discard_snapshot(manifest)
        result_out = dict(result)
        errors = list(result_out.get("errors") or [])
        errors.append(completeness_error)
        result_out.update({
            "applied": False,
            "errors": errors,
            "atomic": True,
            "rolled_back": True,
            "touched_files": touched,
            "rolled_back_files": rolled,
        })
        return result_out

    # Success — discard the snapshot.
    _discard_snapshot(manifest)
    result_out = dict(result)
    result_out.update({
        "atomic": True,
        "rolled_back": False,
        "touched_files": touched,
        "rolled_back_files": [],
    })
    return result_out


__all__ = [
    "extract_uploaded_zip",
    "read_project_file",
    "list_project_files",
    "build_file_tree",
    "build_zip_ai_context",
    "preview_project_contents",
    "describe_bundle",
    "cleanup_extraction",
    "apply_ai_changes",
    "apply_ai_changes_atomic",
    "AtomicPatchError",
    "recompress_project_zip",
]