"""
动态任务依赖图（DAG）— V2 P1-14 修复。
替代 V1 静态 todo.md 表格，支持 Arc 增量重规划，Eng 增量吸收。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


@dataclass
class TaskNode:
    id: str
    status: str = "TODO"  # TODO | DISPATCHED | DONE | BLOCKED | INVALIDATED
    task: str = ""
    files_dirs: str = ""
    done_when: str = ""
    in_degree: int = 0
    out_edges: list[str] = field(default_factory=list)
    write_set: list[str] = field(default_factory=list)
    meta: dict[str, Any] = field(default_factory=dict)
    test_required: bool = False  # P1-19.1: 边界测试强制标记
    adr_refs: list[str] = field(default_factory=list)  # P1-21: ADR→任务溯源链
    keep_constraints: list[str] = field(default_factory=list)  # P1-24: 保留约束（不可修改的文件/目录）


@dataclass
class Edge:
    source: str
    target: str
    kind: str = "dependency"  # dependency | conflict | sync


@dataclass
class EdgeChange:
    added: list[Edge] = field(default_factory=list)
    removed: list[Edge] = field(default_factory=list)


@dataclass
class CascadeReport:
    """P1-21: ADR 变更级联失效报告。"""
    invalidated_completed: int = 0
    terminated_in_progress: int = 0
    cascaded_downstream: int = 0
    rollback_ref: str = ""
    invalidated_task_ids: list[str] = field(default_factory=list)


@dataclass
class PlanDelta:
    """Arc 重规划产出的增量差异，替代全量覆盖 todo.md。"""
    removed_tasks: list[str] = field(default_factory=list)
    added_tasks: list[TaskNode] = field(default_factory=list)
    modified_tasks: list[TaskNode] = field(default_factory=list)
    edge_changes: EdgeChange = field(default_factory=EdgeChange)
    rollback_ref: str = ""  # P1-21: 回滚快照引用


class TaskGraph:
    """动态任务依赖图，支持增量更新和全量替换。"""

    def __init__(self):
        self.nodes: dict[str, TaskNode] = {}
        self.edges: list[Edge] = []
        self.dispatch_frozen: bool = False  # P1-21: 调度冻结

    def add_node(self, node: TaskNode) -> None:
        self.nodes[node.id] = node

    def add_edge(self, edge: Edge) -> None:
        self.edges.append(edge)
        if edge.target in self.nodes:
            self.nodes[edge.target].in_degree += 1
        if edge.source in self.nodes:
            self.nodes[edge.source].out_edges.append(edge.target)

    def apply_delta(self, delta: PlanDelta) -> None:
        """增量吸收 Arc 的重规划结果，保留已调度任务不受影响。"""
        for task_id in delta.removed_tasks:
            self._remove_node(task_id)
        for node in delta.added_tasks:
            self._add_node(node)
        for node in delta.modified_tasks:
            self._update_node(node)
        for edge in delta.edge_changes.removed:
            self._remove_edge(edge)
        for edge in delta.edge_changes.added:
            self._add_edge(edge)

    def apply_full_replace(self, nodes: list[TaskNode], edges: list[Edge]) -> None:
        """全量替换模式：Arc 产出完整 DAG，保留已完成任务状态。
        INVALIDATED 状态不保留（已被级联失效标记的任务在全量替换时重置）。"""
        done_status = {tid: n.status for tid, n in self.nodes.items()
                       if n.status in ("DONE", "DISPATCHED")}
        self.nodes = {n.id: n for n in nodes}
        self.edges = list(edges)
        for tid, status in done_status.items():
            if tid in self.nodes:
                self.nodes[tid].status = status
        for edge in self.edges:
            if edge.target in self.nodes:
                self.nodes[edge.target].in_degree += 1
            if edge.source in self.nodes:
                self.nodes[edge.source].out_edges.append(edge.target)

    def ready_tasks(self) -> list[str]:
        """返回当前入度为 0 且状态为 TODO 的任务。调度冻结时返回空。"""
        if self.dispatch_frozen:
            return []
        return [nid for nid, n in self.nodes.items() if n.in_degree == 0 and n.status == "TODO"]

    def diff(self, other: TaskGraph) -> PlanDelta:
        """对比自身与 other，产出 PlanDelta（add/remove/modify node + edge changes）。

        self = 新图, other = 旧图（before replan）。
        """
        delta = PlanDelta()
        old_ids = set(other.nodes.keys())
        new_ids = set(self.nodes.keys())

        # 移除
        delta.removed_tasks = list(old_ids - new_ids)

        # 新增
        delta.added_tasks = [self.nodes[tid] for tid in (new_ids - old_ids)]

        # 修改
        for tid in old_ids & new_ids:
            old_n = other.nodes[tid]
            new_n = self.nodes[tid]
            if (old_n.task != new_n.task
                    or old_n.files_dirs != new_n.files_dirs
                    or old_n.done_when != new_n.done_when
                    or old_n.write_set != new_n.write_set):
                delta.modified_tasks.append(new_n)

        # Edge 差异
        old_edges = {(e.source, e.target, e.kind) for e in other.edges}
        new_edges = {(e.source, e.target, e.kind) for e in self.edges}
        for s, t, k in (new_edges - old_edges):
            delta.edge_changes.added.append(Edge(source=s, target=t, kind=k))
        for s, t, k in (old_edges - new_edges):
            delta.edge_changes.removed.append(Edge(source=s, target=t, kind=k))

        return delta

    @classmethod
    def load(cls, path) -> TaskGraph:
        """从 _export_task_graph_json 写的格式还原 TaskGraph。"""
        from pathlib import Path
        from air_runtime.io import safe_json_load
        p = Path(path)
        data = safe_json_load(p)
        graph = cls()
        if not data or not isinstance(data, dict):
            return graph
        graph.dispatch_frozen = data.get("dispatchFrozen", False)
        for nid, nd in data.get("nodes", {}).items():
            graph.nodes[nid] = TaskNode(
                id=nd.get("id", nid),
                status=nd.get("status", "TODO"),
                task=nd.get("task", ""),
                files_dirs=nd.get("filesDirs", ""),
                done_when=nd.get("doneWhen", ""),
                in_degree=nd.get("inDegree", 0),
                out_edges=list(nd.get("outEdges", [])),
                write_set=list(nd.get("writeSet", [])),
                test_required=nd.get("testRequired", False),
                adr_refs=list(nd.get("adrRefs", [])),
                keep_constraints=list(nd.get("keepConstraints", [])),
            )
        for ed in data.get("edges", []):
            graph.edges.append(Edge(
                source=ed["source"], target=ed["target"],
                kind=ed.get("kind", "dependency"),
            ))
        return graph

    def task_ids_by_status(self, status: str) -> list[str]:
        return [nid for nid, n in self.nodes.items() if n.status == status]

    def find_cycles(self) -> list[list[str]]:
        """检测依赖环（DFS）。"""
        visited: set[str] = set()
        rec_stack: set[str] = set()
        cycles: list[list[str]] = []

        def dfs(node_id: str, path: list[str]) -> None:
            visited.add(node_id)
            rec_stack.add(node_id)
            path.append(node_id)
            for target in self.nodes.get(node_id, TaskNode(id=node_id)).out_edges:
                if target not in visited:
                    dfs(target, path.copy())
                elif target in rec_stack:
                    cycle_start = path.index(target)
                    cycles.append(path[cycle_start:])
            rec_stack.discard(node_id)

        for nid in self.nodes:
            if nid not in visited:
                dfs(nid, [])

        return cycles

    def export_todo_md(self) -> str:
        """导出为人可读的 todo.md 表格，保留 V1 的可见性优势。"""
        lines = ["| Task | Status | Files/Dirs | Done When | Validation | ADR |",
                 "|------|--------|------------|-----------|------------|-----|"]
        for nid, node in self.nodes.items():
            lines.append(f"| {node.task} | {node.status} | {node.files_dirs} | "
                        f"{node.done_when} | | |")
        return "\n".join(lines) + "\n"

    def _remove_node(self, task_id: str) -> None:
        if task_id in self.nodes:
            del self.nodes[task_id]
        self.edges = [e for e in self.edges if e.source != task_id and e.target != task_id]

    def _add_node(self, node: TaskNode) -> None:
        self.nodes[node.id] = node

    def _update_node(self, node: TaskNode) -> None:
        if node.id in self.nodes:
            existing_status = self.nodes[node.id].status
            self.nodes[node.id] = node
            # INVALIDATED 可覆盖 DONE/DISPATCHED（P1-21: ADR 级联失效）
            if existing_status in ("DISPATCHED", "DONE") and node.status != "INVALIDATED":
                self.nodes[node.id].status = existing_status

    def _remove_edge(self, edge: Edge) -> None:
        self.edges = [e for e in self.edges
                      if not (e.source == edge.source and e.target == edge.target)]
        if edge.target in self.nodes:
            self.nodes[edge.target].in_degree = max(0, self.nodes[edge.target].in_degree - 1)

    def _add_edge(self, edge: Edge) -> None:
        self.edges.append(edge)
        if edge.target in self.nodes:
            self.nodes[edge.target].in_degree += 1
        if edge.source in self.nodes:
            self.nodes[edge.source].out_edges.append(edge.target)

    # P1-21: ADR 级联失效

    def tasks_by_adr(self, adr_id: str) -> list[TaskNode]:
        """查找所有引用指定 ADR 的任务（含已完成）。"""
        return [n for n in self.nodes.values() if adr_id in n.adr_refs]

    def _find_downstream(self, task_ids: list[str]) -> list[str]:
        """BFS 遍历下游依赖任务。"""
        visited: set[str] = set()
        queue = list(task_ids)
        while queue:
            current = queue.pop(0)
            if current in visited:
                continue
            visited.add(current)
            node = self.nodes.get(current)
            if node:
                for target in node.out_edges:
                    if target not in visited:
                        queue.append(target)
        # 排除起点自身
        return [tid for tid in visited if tid not in set(task_ids)]

    def invalidate_by_adr(self, adr_id: str, delta: PlanDelta) -> CascadeReport:
        """P1-21: ADR 变更时级联失效所有相关任务。"""
        affected = self.tasks_by_adr(adr_id)
        completed = [t for t in affected if t.status == "DONE"]
        in_progress = [t for t in affected if t.status == "DISPATCHED"]
        pending = [t for t in affected if t.status == "TODO"]

        # 1. 冻结调度
        self.dispatch_frozen = True

        # 2. 标记已完成任务为 INVALIDATED
        for t in completed:
            t.status = "INVALIDATED"
            delta.removed_tasks.append(t.id)

        # 3. 标记进行中任务为 INVALIDATED（调用方负责中止 Worker）
        for t in in_progress:
            t.status = "INVALIDATED"
            delta.removed_tasks.append(t.id)

        # 4. 标记 ADR 直接关联的 TODO 任务为 INVALIDATED
        for t in pending:
            t.status = "INVALIDATED"
            delta.removed_tasks.append(t.id)

        # 5. 级联失效下游
        downstream_ids = self._find_downstream([t.id for t in completed + in_progress + pending])
        cascaded = []
        for tid in downstream_ids:
            node = self.nodes.get(tid)
            if node and node.status in ("TODO", "DISPATCHED"):
                node.status = "INVALIDATED"
                delta.removed_tasks.append(tid)
                cascaded.append(tid)

        # 6. 回滚快照引用（由调用方在 git revert 后填入）
        all_invalidated = [t.id for t in completed + in_progress + pending] + cascaded

        return CascadeReport(
            invalidated_completed=len(completed),
            terminated_in_progress=len(in_progress),
            cascaded_downstream=len(cascaded),
            rollback_ref=delta.rollback_ref,
            invalidated_task_ids=all_invalidated,
        )

    def unfreeze_dispatch(self) -> None:
        """P1-21: 解冻调度，在 Arc 重新生成受影响任务后调用。"""
        self.dispatch_frozen = False