AirCoding/AirPlan/docs/spec/AirPlan-ParaV2/lib/air_runtime/review.py

from __future__ import annotations

from itertools import combinations
from pathlib import Path
from typing import Dict, Iterable, List, Set

from .contracts import (
    ACTIVE_TASK_STATUSES,
    ParallelGroup,
    ParallelReview,
    ReviewConflict,
    TaskRecord,
    now_iso,
)
from .todo_parser import parse_tasks


def _overlap_paths(left: Iterable[str], right: Iterable[str]) -> List[str]:
    overlap: Set[str] = set()
    left_items = list(left)
    right_items = list(right)

    for left_path in left_items:
        for right_path in right_items:
            left_clean = left_path.rstrip("/\\")
            right_clean = right_path.rstrip("/\\")
            if left_clean == right_clean:
                overlap.add(left_path)
                continue
            if left_clean.startswith(right_clean):
                overlap.add(right_path)
            elif right_clean.startswith(left_clean):
                overlap.add(left_path)

    return sorted(overlap)


def _task_conflict(left: TaskRecord, right: TaskRecord) -> ReviewConflict | None:
    overlap = _overlap_paths(left.normalized_write_set(), right.normalized_write_set())
    if overlap:
        return ReviewConflict(
            task_ids=[left.task_id, right.task_id],
            reason="shared write set overlap",
            overlap_paths=overlap,
        )

    if left.touches_global_docs() and right.touches_global_docs():
        return ReviewConflict(
            task_ids=[left.task_id, right.task_id],
            reason="both tasks touch global planning or architecture documents",
            overlap_paths=sorted(set(left.global_doc_paths + right.global_doc_paths)),
        )

    return None


def _greedy_parallel_groups(tasks: List[TaskRecord]) -> List[ParallelGroup]:
    groups: List[List[TaskRecord]] = []
    for task in tasks:
        placed = False
        for group in groups:
            if all(_task_conflict(task, existing) is None for existing in group):
                group.append(task)
                placed = True
                break
        if not placed:
            groups.append([task])

    output: List[ParallelGroup] = []
    for index, group in enumerate(groups, start=1):
        output.append(
            ParallelGroup(
                name=f"group-{index}",
                task_ids=[task.task_id for task in group],
                reason="no detected write-set conflict inside this group",
            )
        )
    return output


def _build_dependency_edges(tasks: List[TaskRecord]) -> List[Dict[str, str]]:
    task_ids = {task.task_id for task in tasks}
    edges: List[Dict[str, str]] = []
    for task in tasks:
        for dependency in task.dependencies:
            if dependency in task_ids:
                edges.append({"from": dependency, "to": task.task_id})
    return edges


def build_parallel_review(todo_path: Path) -> ParallelReview:
    all_tasks = parse_tasks(todo_path)
    active_tasks = [task for task in all_tasks if task.status in ACTIVE_TASK_STATUSES]
    active_task_ids = {task.task_id for task in active_tasks}
    task_map = {task.task_id: task for task in active_tasks}
    edges = _build_dependency_edges(active_tasks)

    in_degree = {task.task_id: 0 for task in active_tasks}
    children: Dict[str, List[str]] = {task.task_id: [] for task in active_tasks}
    for edge in edges:
        parent = edge["from"]
        child = edge["to"]
        if parent not in active_task_ids or child not in active_task_ids:
            continue
        in_degree[child] += 1
        children[parent].append(child)

    ready = sorted(
        [task.task_id for task in active_tasks if in_degree[task.task_id] == 0],
        key=lambda task_id: task_map[task_id].line_number,
    )
    scheduled = set()
    parallel_groups: List[ParallelGroup] = []
    wave_index = 1

    while ready:
        current_wave_ids = ready
        ready = []
        current_tasks = [task_map[task_id] for task_id in current_wave_ids]
        wave_groups = _greedy_parallel_groups(current_tasks)
        for group in wave_groups:
            group.name = f"wave-{wave_index}-{group.name}"
        parallel_groups.extend(wave_groups)
        wave_index += 1

        for task_id in current_wave_ids:
            scheduled.add(task_id)
            for child in children.get(task_id, []):
                in_degree[child] -= 1
                if in_degree[child] == 0:
                    ready.append(child)
        ready.sort(key=lambda task_id: task_map[task_id].line_number)

    notes: List[str] = []
    unscheduled = sorted(active_task_ids - scheduled)
    if unscheduled:
        notes.append(
            "Some active tasks could not be layered. Check for cyclic or missing dependencies: "
            + ", ".join(unscheduled)
        )

    if all(not task.dependencies for task in active_tasks) and len(active_tasks) > 1:
        notes.append(
            "No explicit dependency hints were found. Parallel grouping relies on write-set isolation and global-doc serialization rules."
        )

    blocked_tasks = [task.task_id for task in all_tasks if task.status == "BLOCKED"]
    if blocked_tasks:
        notes.append("Blocked tasks were excluded from ready groups: " + ", ".join(blocked_tasks))

    conflicts: List[ReviewConflict] = []
    for left, right in combinations(active_tasks, 2):
        conflict = _task_conflict(left, right)
        if conflict is not None:
            conflicts.append(conflict)

    serialization_points: List[Dict[str, object]] = []
    for task in active_tasks:
        reasons: List[str] = []
        if task.touches_global_docs():
            reasons.append("touches global planning or architecture documents")
        if any(task.task_id in conflict.task_ids for conflict in conflicts):
            reasons.append("has shared write-set conflicts that need engine-level scheduling")
        if reasons:
            serialization_points.append(
                {
                    "taskId": task.task_id,
                    "reasons": reasons,
                    "paths": task.global_doc_paths or task.normalized_write_set(),
                }
            )

    return ParallelReview(
        source_todo=str(todo_path),
        generated_at=now_iso(),
        tasks_considered=active_tasks,
        dependency_edges=edges,
        parallel_groups=parallel_groups,
        conflicts=conflicts,
        serialization_points=serialization_points,
        notes=notes,
    )


def render_review_markdown(review: ParallelReview) -> str:
    lines = [
        "# AirArc Parallel Review",
        "",
        f"- Source TODO: `{review.source_todo}`",
        f"- Generated At: `{review.generated_at}`",
        f"- Active Tasks: `{', '.join(task.task_id for task in review.tasks_considered) or 'none'}`",
        "",
        "## Parallel Groups",
    ]

    if review.parallel_groups:
        for group in review.parallel_groups:
            lines.append(f"- `{group.name}`: {', '.join(group.task_ids)}")
            lines.append(f"  Reason: {group.reason}")
    else:
        lines.append("- No ready parallel groups were detected.")

    lines.extend(["", "## Dependency Edges"])
    if review.dependency_edges:
        for edge in review.dependency_edges:
            lines.append(f"- `{edge['from']}` -> `{edge['to']}`")
    else:
        lines.append("- No explicit dependency edges were detected.")

    lines.extend(["", "## Shared-Write Conflicts"])
    if review.conflicts:
        for conflict in review.conflicts:
            lines.append(
                f"- `{conflict.task_ids[0]}` <-> `{conflict.task_ids[1]}`: {conflict.reason}"
            )
            if conflict.overlap_paths:
                lines.append("  Paths: " + ", ".join(f"`{path}`" for path in conflict.overlap_paths))
    else:
        lines.append("- No shared-write conflicts were detected.")

    lines.extend(["", "## Serialization Points"])
    if review.serialization_points:
        for point in review.serialization_points:
            lines.append(f"- `{point['taskId']}`: {'; '.join(point['reasons'])}")
            paths = point.get("paths", [])
            if paths:
                lines.append("  Paths: " + ", ".join(f"`{path}`" for path in paths))
    else:
        lines.append("- No serialization points were detected.")

    lines.extend(["", "## Notes"])
    if review.notes:
        for note in review.notes:
            lines.append(f"- {note}")
    else:
        lines.append("- No extra notes.")

    return "\n".join(lines) + "\n"