tools/annotate_dmhy_dag_agent.py

"""Run a small DAG-aware LLM annotation pass for DMHY filename suffixes.

This runner is intentionally separate from the prefix-tree annotation scripts.
It selects shared DAG suffix units, asks an OpenAI-compatible Responses API for
strict JSON annotations, falls back to deterministic heuristics on failure, and
emits both terminal patch rows and dmhy_weak-compatible sample records.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

from anifilebert.tokenizer import AnimeTokenizer
from tools.annotate_dmhy_prefix_graph import (
    heuristic_patch,
    normalize_generated_tokens,
    string_list,
    unique_keep_order,
    write_jsonl,
)
from tools.dmhy_dataset import weak_label_filename


DEFAULT_DAG = Path("datasets/AnimeName/dmhy_prefix_dag.json")
DEFAULT_OUTPUT = Path("datasets/AnimeName/dmhy_weak.dag_sample.jsonl")
DEFAULT_PATCH_OUTPUT = Path("datasets/AnimeName/dmhy_prefix_dag.annotations.dag_sample.jsonl")
DEFAULT_UNITS_OUTPUT = Path("datasets/AnimeName/dmhy_prefix_dag.annotation_units.dag_sample.jsonl")
DEFAULT_BASE_URL = "http://10.137.32.209/v1"
DEFAULT_MODEL = "gpt-5.4-mini"
LLM_SOURCE = "responses-dag-v1"
HEURISTIC_SOURCE = "heuristic-dag-v1"
COVERAGE_SHARED_ONLY = "shared-only"
COVERAGE_ALL_TERMINALS = "all-terminals"


@dataclass(frozen=True)
class Args:
    dag: Path
    output: Path
    patch_output: Path
    units_output: Path | None
    limit: int
    max_requests: int
    units_per_request: int
    max_context_units: int
    workers: int
    coverage_mode: str
    min_incoming_count: int
    min_reachable_terminals: int
    max_reachable_terminals: int | None
    example_count: int
    max_records: int
    records_per_terminal: int
    terminal_sink_threshold: int
    terminal_cluster_size: int
    llm: bool
    base_url: str
    api_key: str | None
    model: str
    http_timeout: int
    retries: int
    preserve_i_labels: bool
    resume: bool


def parse_args() -> Args:
    parser = argparse.ArgumentParser(
        description="Annotate a small sample of shared DMHY prefix-DAG suffix units"
    )
    parser.add_argument("--dag", type=Path, default=DEFAULT_DAG)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--patch-output", type=Path, default=DEFAULT_PATCH_OUTPUT)
    parser.add_argument(
        "--units-output",
        default=str(DEFAULT_UNITS_OUTPUT),
        help="Optional selected annotation units JSONL; use empty string to disable",
    )
    parser.add_argument("--limit", type=int, default=20, help="Maximum DAG units to annotate")
    parser.add_argument("--max-requests", type=int, default=20, help="Maximum LLM requests")
    parser.add_argument("--units-per-request", type=int, default=10, help="DAG units per Responses API request")
    parser.add_argument(
        "--max-context-units",
        type=int,
        default=20,
        help="Maximum unit contexts included in one Responses API request",
    )
    parser.add_argument("--workers", type=int, default=8, help="Concurrent LLM workers")
    parser.add_argument(
        "--coverage-mode",
        choices=[COVERAGE_SHARED_ONLY, COVERAGE_ALL_TERMINALS],
        default=COVERAGE_SHARED_ONLY,
        help="Select only shared DAG units, or cover every terminal with terminal_cluster fallback units",
    )
    parser.add_argument("--min-incoming-count", type=int, default=2)
    parser.add_argument("--min-reachable-terminals", type=int, default=2)
    parser.add_argument(
        "--max-reachable-terminals",
        type=int,
        default=500,
        help="Skip very large shared units for sample runs; use 0 to disable",
    )
    parser.add_argument("--example-count", type=int, default=8)
    parser.add_argument("--max-records", type=int, default=200, help="Maximum dataset records to emit")
    parser.add_argument("--records-per-terminal", type=int, default=1)
    parser.add_argument(
        "--terminal-sink-threshold",
        type=int,
        default=1000,
        help="Do not collect terminals from no-outgoing sink nodes above this size",
    )
    parser.add_argument(
        "--terminal-cluster-size",
        type=int,
        default=1,
        help="Number of uncovered terminals per all-terminals fallback cluster",
    )
    parser.add_argument("--llm", action="store_true", help="Use the Responses API; otherwise heuristics only")
    parser.add_argument(
        "--base-url",
        default=os.environ.get("ANIFILEBERT_LLM_BASE_URL", DEFAULT_BASE_URL),
        help="OpenAI-compatible API base URL",
    )
    parser.add_argument(
        "--api-key",
        default=os.environ.get("ANIFILEBERT_LLM_API_KEY"),
        help="API key; defaults to ANIFILEBERT_LLM_API_KEY",
    )
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--http-timeout", type=int, default=120)
    parser.add_argument("--retries", type=int, default=3)
    parser.add_argument("--preserve-i-labels", action="store_true")
    parser.add_argument("--resume", action="store_true", help="Skip ok/fallback unit_ids already in patch output")
    ns = parser.parse_args()
    units_arg = str(ns.units_output).strip()
    return Args(
        dag=ns.dag,
        output=ns.output,
        patch_output=ns.patch_output,
        units_output=Path(units_arg) if units_arg else None,
        limit=max(0, ns.limit),
        max_requests=max(0, ns.max_requests),
        units_per_request=max(1, ns.units_per_request),
        max_context_units=max(1, ns.max_context_units),
        workers=max(1, ns.workers),
        coverage_mode=ns.coverage_mode,
        min_incoming_count=max(1, ns.min_incoming_count),
        min_reachable_terminals=max(1, ns.min_reachable_terminals),
        max_reachable_terminals=None if ns.max_reachable_terminals <= 0 else ns.max_reachable_terminals,
        example_count=max(1, ns.example_count),
        max_records=max(0, ns.max_records),
        records_per_terminal=max(1, ns.records_per_terminal),
        terminal_sink_threshold=max(1, ns.terminal_sink_threshold),
        terminal_cluster_size=max(1, ns.terminal_cluster_size),
        llm=ns.llm,
        base_url=ns.base_url,
        api_key=ns.api_key,
        model=ns.model,
        http_timeout=max(1, ns.http_timeout),
        retries=max(1, ns.retries),
        preserve_i_labels=ns.preserve_i_labels,
        resume=ns.resume,
    )


def load_dag(path: Path) -> dict[str, Any]:
    if not path.exists():
        raise SystemExit(f"DAG not found: {path}")
    try:
        dag = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise SystemExit(f"invalid DAG JSON in {path}: {exc}") from exc
    if not isinstance(dag, dict) or not isinstance(dag.get("nodes"), list):
        raise SystemExit(f"invalid DAG schema in {path}: missing nodes list")
    if not isinstance(dag.get("terminals"), list):
        raise SystemExit(f"invalid DAG schema in {path}: missing terminals list")
    return dag


def int_field(row: dict[str, Any], key: str, default: int = 0) -> int:
    try:
        return int(row.get(key, default) or default)
    except (TypeError, ValueError):
        return default


def node_id(node: dict[str, Any], fallback: int) -> int:
    try:
        return int(node.get("id", fallback))
    except (TypeError, ValueError):
        raise SystemExit(f"invalid node id: {node.get('id')!r}") from None


def terminal_id(terminal: dict[str, Any], fallback: int) -> str:
    value = terminal.get("terminal_id", terminal.get("id", fallback))
    return str(value)


def has_outgoing_edges(node: dict[str, Any]) -> bool:
    return any(isinstance(edge, dict) for edge in node.get("children") or [])


def edge_labels(node: dict[str, Any], limit: int) -> list[str]:
    labels: list[str] = []
    for edge in node.get("children") or []:
        if isinstance(edge, dict) and edge.get("label") is not None:
            labels.append(str(edge["label"]))
    return unique_keep_order(labels)[:limit]


def build_indexes(
    dag: dict[str, Any],
) -> tuple[dict[int, dict[str, Any]], dict[int, list[dict[str, Any]]]]:
    nodes: dict[int, dict[str, Any]] = {}
    for index, node in enumerate(dag["nodes"]):
        if isinstance(node, dict):
            nodes[node_id(node, index)] = node

    terminals_by_node: dict[int, list[dict[str, Any]]] = {}
    for index, terminal in enumerate(dag["terminals"]):
        if not isinstance(terminal, dict):
            continue
        try:
            terminal_node_id = int(terminal.get("node_id"))
        except (TypeError, ValueError):
            continue
        row = dict(terminal)
        row["_terminal_id"] = terminal_id(row, index)
        row["_terminal_index"] = index
        terminals_by_node.setdefault(terminal_node_id, []).append(row)
    return nodes, terminals_by_node


class ReachableTerminals:
    def __init__(
        self,
        nodes: dict[int, dict[str, Any]],
        terminals_by_node: dict[int, list[dict[str, Any]]],
        terminal_sink_threshold: int,
    ) -> None:
        self.nodes = nodes
        self.terminals_by_node = terminals_by_node
        self.terminal_sink_threshold = terminal_sink_threshold
        self.memo: dict[int, list[dict[str, Any]]] = {}
        self.visiting: set[int] = set()

    def get(self, start: int) -> list[dict[str, Any]]:
        if start in self.memo:
            return self.memo[start]
        if start in self.visiting:
            raise SystemExit(f"cycle detected while traversing DAG at node {start}")
        self.visiting.add(start)
        node = self.nodes.get(start, {})
        local_terminals = self.terminals_by_node.get(start, [])
        is_large_terminal_sink = not has_outgoing_edges(node) and len(local_terminals) >= self.terminal_sink_threshold
        found = [] if is_large_terminal_sink else list(local_terminals)
        for edge in node.get("children") or []:
            if not isinstance(edge, dict):
                continue
            try:
                found.extend(self.get(int(edge.get("target"))))
            except (TypeError, ValueError):
                continue
        self.visiting.remove(start)
        deduped = dedupe_terminals(found)
        self.memo[start] = deduped
        return deduped


def dedupe_terminals(terminals: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
    seen: set[str] = set()
    result: list[dict[str, Any]] = []
    for terminal in terminals:
        tid = str(terminal.get("_terminal_id") or terminal.get("terminal_id") or "")
        if not tid or tid in seen:
            continue
        seen.add(tid)
        result.append(terminal)
    return result


def aggregate_examples(terminals: Iterable[dict[str, Any]], key: str, limit: int) -> list[str]:
    values: list[str] = []
    for terminal in terminals:
        if key == "prefix":
            value = terminal.get("prefix")
            if value is not None:
                values.append(str(value))
        else:
            values.extend(string_list(terminal.get(key)))
    return unique_keep_order(values)[:limit]


def aggregate_terminal_field(terminals: Iterable[dict[str, Any]], key: str, limit: int) -> list[str]:
    values: list[str] = []
    for terminal in terminals:
        value = terminal.get(key)
        if value is not None:
            values.append(str(value))
    return unique_keep_order(values)[:limit]


def heuristic_unit_patch(unit: dict[str, Any]) -> dict[str, Any]:
    episode_titles: list[str] = []
    media_suffixes: list[str] = []
    title_candidates: list[str] = []
    needs_review = False
    for terminal in unit["_terminals"]:
        patch = heuristic_patch(terminal, int(terminal.get("_terminal_index", 0)))
        episode_titles.extend(patch["episode_title_suffixes"])
        media_suffixes.extend(patch["media_suffixes"])
        title_candidates.extend(patch["title_candidates"])
        needs_review = needs_review or bool(patch["needs_llm_review"])
    return {
        "unit_id": unit["unit_id"],
        "terminal_ids": unit["terminal_ids"],
        "episode_title_suffixes": unique_keep_order(episode_titles),
        "media_suffixes": unique_keep_order(media_suffixes),
        "title_candidates": unique_keep_order(title_candidates),
        "llm_label": None,
        "notes": f"heuristic fallback; terminals={len(unit['terminal_ids'])}; needs_review={needs_review}",
    }


def make_unit(
    node: dict[str, Any],
    node_id_value: int,
    terminals: list[dict[str, Any]],
    example_count: int,
) -> dict[str, Any]:
    terminal_ids = [str(terminal["_terminal_id"]) for terminal in terminals]
    reachable_weight = int_field(node, "reachable_weight")
    if reachable_weight <= 0:
        reachable_weight = sum(int_field(terminal, "weight", int_field(terminal, "count", 1)) for terminal in terminals)
    return {
        "unit_id": f"dag-node-{node_id_value}",
        "node_id": node_id_value,
        "kind": "shared_suffix",
        "incoming_count": int_field(node, "incoming_count"),
        "reachable_terminals": len(terminals),
        "reachable_weight": reachable_weight,
        "terminal_ids": terminal_ids,
        "prefix_examples": aggregate_examples(terminals, "prefix", example_count),
        "digit_skeleton_examples": aggregate_terminal_field(terminals, "digit_skeleton", example_count),
        "value_examples": aggregate_examples(terminals, "value_examples", example_count),
        "suffix_examples": aggregate_examples(terminals, "suffix_examples", example_count),
        "common_edge_labels": edge_labels(node, example_count),
        "_terminals": terminals,
    }


def make_terminal_cluster_unit(
    terminals: list[dict[str, Any]],
    cluster_index: int,
    example_count: int,
) -> dict[str, Any]:
    terminal_ids = [str(terminal["_terminal_id"]) for terminal in terminals]
    reachable_weight = sum(int_field(terminal, "weight", int_field(terminal, "count", 1)) for terminal in terminals)
    return {
        "unit_id": f"terminal-cluster-{cluster_index}",
        "node_id": None,
        "kind": "terminal_cluster",
        "incoming_count": 0,
        "reachable_terminals": len(terminals),
        "reachable_weight": reachable_weight,
        "terminal_ids": terminal_ids,
        "prefix_examples": aggregate_examples(terminals, "prefix", example_count),
        "digit_skeleton_examples": aggregate_terminal_field(terminals, "digit_skeleton", example_count),
        "value_examples": aggregate_examples(terminals, "value_examples", example_count),
        "suffix_examples": aggregate_examples(terminals, "suffix_examples", example_count),
        "common_edge_labels": [],
        "_terminals": terminals,
    }


def all_terminals(dag: dict[str, Any]) -> list[dict[str, Any]]:
    terminals: list[dict[str, Any]] = []
    for index, terminal in enumerate(dag["terminals"]):
        if not isinstance(terminal, dict):
            continue
        row = dict(terminal)
        row["_terminal_id"] = terminal_id(row, index)
        row["_terminal_index"] = index
        terminals.append(row)
    return dedupe_terminals(terminals)


def candidate_shared_units(
    dag: dict[str, Any],
    args: Args,
) -> list[dict[str, Any]]:
    nodes, terminals_by_node = build_indexes(dag)
    reachable = ReachableTerminals(nodes, terminals_by_node, args.terminal_sink_threshold)
    root = int(dag.get("root", 0) or 0)
    candidates: list[tuple[tuple[int, int, int, int], int, dict[str, Any]]] = []
    for current_id, node in nodes.items():
        if current_id == root or not has_outgoing_edges(node):
            continue
        incoming_count = int_field(node, "incoming_count")
        reachable_count = int_field(node, "reachable_terminals")
        if incoming_count < args.min_incoming_count:
            continue
        if reachable_count < args.min_reachable_terminals:
            continue
        if args.max_reachable_terminals is not None and reachable_count > args.max_reachable_terminals:
            continue
        sort_key = (-incoming_count, -int_field(node, "reachable_weight"), -reachable_count, current_id)
        candidates.append((sort_key, current_id, node))

    candidates.sort(key=lambda item: item[0])
    candidates.sort(key=lambda item: item[0])
    units: list[dict[str, Any]] = []
    for _sort_key, current_id, node in candidates:
        terminals = reachable.get(current_id)
        if not terminals:
            continue
        units.append(make_unit(node, current_id, terminals, args.example_count))
    return units


def unique_terminal_coverage(units: Iterable[dict[str, Any]]) -> set[str]:
    covered: set[str] = set()
    for unit in units:
        covered.update(str(terminal_id_value) for terminal_id_value in unit.get("terminal_ids") or [])
    return covered


def shared_only_units(dag: dict[str, Any], args: Args) -> list[dict[str, Any]]:
    limit = args.limit if args.limit > 0 else 0
    return candidate_shared_units(dag, args)[:limit]


def all_terminal_units(dag: dict[str, Any], args: Args) -> list[dict[str, Any]]:
    terminals = all_terminals(dag)
    terminals_by_id = {str(terminal["_terminal_id"]): terminal for terminal in terminals}
    selected: list[dict[str, Any]] = []
    covered: set[str] = set()

    for unit in candidate_shared_units(dag, args):
        unit_terminal_ids = set(unit["terminal_ids"])
        if not unit_terminal_ids:
            continue
        if unit_terminal_ids & covered:
            continue
        selected.append(unit)
        covered.update(unit_terminal_ids)

    uncovered = [terminal for terminal in terminals if str(terminal["_terminal_id"]) not in covered]
    for start in range(0, len(uncovered), args.terminal_cluster_size):
        cluster_terminals = uncovered[start : start + args.terminal_cluster_size]
        selected.append(
            make_terminal_cluster_unit(
                cluster_terminals,
                cluster_index=(start // args.terminal_cluster_size) + 1,
                example_count=args.example_count,
            )
        )
        covered.update(str(terminal["_terminal_id"]) for terminal in cluster_terminals)

    expected_ids = set(terminals_by_id)
    if covered != expected_ids:
        missing = sorted(expected_ids - covered)[:10]
        extra = sorted(covered - expected_ids)[:10]
        raise SystemExit(
            "all-terminals coverage mismatch: "
            f"expected={len(expected_ids)} covered={len(covered)} missing={missing} extra={extra}"
        )
    return selected


def select_units(dag: dict[str, Any], args: Args) -> list[dict[str, Any]]:
    if args.coverage_mode == COVERAGE_ALL_TERMINALS:
        return all_terminal_units(dag, args)
    return shared_only_units(dag, args)


def annotation_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "additionalProperties": False,
        "required": [
            "unit_id",
            "terminal_ids",
            "episode_title_suffixes",
            "media_suffixes",
            "title_candidates",
            "llm_label",
            "notes",
        ],
        "properties": {
            "unit_id": {"type": "string"},
            "terminal_ids": {"type": "array", "items": {"type": "string"}},
            "episode_title_suffixes": {"type": "array", "items": {"type": "string"}},
            "media_suffixes": {"type": "array", "items": {"type": "string"}},
            "title_candidates": {"type": "array", "items": {"type": "string"}},
            "llm_label": {"type": ["string", "null"]},
            "notes": {"type": "string"},
        },
    }


def batch_annotation_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "additionalProperties": False,
        "required": ["annotations"],
        "properties": {
            "annotations": {
                "type": "array",
                "items": annotation_schema(),
            }
        },
    }


def responses_url(base_url: str) -> str:
    return base_url.rstrip("/") + "/responses"


def unit_request_context(unit: dict[str, Any]) -> dict[str, Any]:
    return {
        "unit_id": unit["unit_id"],
        "node_id": unit["node_id"],
        "terminal_ids": unit["terminal_ids"],
        "incoming_count": unit["incoming_count"],
        "reachable_terminals": unit["reachable_terminals"],
        "reachable_weight": unit["reachable_weight"],
        "prefix_examples": unit["prefix_examples"],
        "value_examples": unit["value_examples"],
        "suffix_examples": unit["suffix_examples"],
        "common_edge_labels": unit["common_edge_labels"],
        "heuristic_patch": heuristic_unit_patch(unit),
    }


def request_variants(units: list[dict[str, Any]], args: Args, previous_error: str | None = None) -> list[dict[str, Any]]:
    if not units:
        raise ValueError("request_variants requires at least one unit")
    instructions = (
        "You annotate anime release filename DAG suffix units. Classify only suffix text that is "
        "a human episode title into episode_title_suffixes. Put resolution, codec, source, audio, "
        "subtitle, language, hash, release, and file-format fragments into media_suffixes. "
        "Use title_candidates for possible anime title strings found in prefix/value examples. Return "
        "strict JSON matching the supplied schema with one annotation per input unit. Keep unit_id and "
        "terminal_ids exactly unchanged for every unit."
    )
    if previous_error:
        instructions += f" Previous attempt failed validation: {previous_error}. Correct only that issue."
    input_obj = {
        "units": [unit_request_context(unit) for unit in units[: args.max_context_units]],
    }
    schema = batch_annotation_schema()
    base = {
        "model": args.model,
        "instructions": instructions,
        "input": json.dumps(input_obj, ensure_ascii=False),
    }
    return [
        {
            **base,
            "reasoning": {"effort": "medium"},
            "text": {
                "format": {
                    "type": "json_schema",
                    "name": "dmhy_dag_batch_annotation",
                    "strict": True,
                    "schema": schema,
                }
            },
        },
        {
            **base,
            "text": {
                "format": {
                    "type": "json_schema",
                    "name": "dmhy_dag_batch_annotation",
                    "strict": True,
                    "schema": schema,
                }
            },
        },
        {
            **base,
            "tools": [
                {
                    "type": "function",
                    "name": "submit_dmhy_dag_batch_annotation",
                    "strict": True,
                    "parameters": schema,
                    "description": "Submit strict DMHY DAG suffix annotations for all input units.",
                }
            ],
            "tool_choice": {"type": "function", "name": "submit_dmhy_dag_batch_annotation"},
        },
    ]


def http_post_json(payload: dict[str, Any], args: Args) -> dict[str, Any]:
    if not args.api_key:
        raise RuntimeError("--llm requires --api-key or ANIFILEBERT_LLM_API_KEY")
    request = urllib.request.Request(
        responses_url(args.base_url),
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {args.api_key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=args.http_timeout) as response:
            raw = response.read().decode("utf-8")
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"Responses API HTTP {exc.code}: {body[:500]}") from exc
    except (urllib.error.URLError, OSError) as exc:
        raise RuntimeError(f"Responses API request failed: {exc}") from exc
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Responses API returned invalid JSON envelope: {raw[:500]}") from exc
    if not isinstance(data, dict):
        raise RuntimeError("Responses API envelope must be a JSON object")
    return data


def extract_annotation_payload(data: dict[str, Any]) -> dict[str, Any]:
    output_text = data.get("output_text")
    if isinstance(output_text, str) and output_text.strip():
        return json.loads(strip_json_fence(output_text))

    text_chunks: list[str] = []
    for item in data.get("output") or []:
        if not isinstance(item, dict):
            continue
        if item.get("type") in {"function_call", "tool_call"}:
            arguments = item.get("arguments")
            if isinstance(arguments, str) and arguments.strip():
                return json.loads(arguments)
            if isinstance(arguments, dict):
                return arguments
        for content in item.get("content") or []:
            if not isinstance(content, dict):
                continue
            text = content.get("text")
            if isinstance(text, str):
                text_chunks.append(text)
            if content.get("type") in {"function_call", "tool_call"}:
                arguments = content.get("arguments")
                if isinstance(arguments, str) and arguments.strip():
                    return json.loads(arguments)
                if isinstance(arguments, dict):
                    return arguments
    if text_chunks:
        return json.loads(strip_json_fence("\n".join(text_chunks)))
    raise RuntimeError("Responses API did not include annotation text or tool arguments")


def strip_json_fence(text: str) -> str:
    text = text.strip()
    if text.startswith("```"):
        lines = text.splitlines()
        if lines:
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        return "\n".join(lines).strip()
    return text


def validate_annotation(annotation: Any, unit: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(annotation, dict):
        raise RuntimeError("annotation must be a JSON object")
    missing = [key for key in annotation_schema()["required"] if key not in annotation]
    if missing:
        raise RuntimeError(f"annotation missing required keys: {', '.join(missing)}")
    if annotation["unit_id"] != unit["unit_id"]:
        raise RuntimeError(f"unit_id mismatch: {annotation['unit_id']!r}")
    returned_ids = [str(item) for item in annotation.get("terminal_ids") or []]
    if returned_ids != unit["terminal_ids"]:
        raise RuntimeError("terminal_ids mismatch")
    cleaned = {
        "unit_id": unit["unit_id"],
        "terminal_ids": unit["terminal_ids"],
        "episode_title_suffixes": unique_keep_order(str(item) for item in annotation["episode_title_suffixes"]),
        "media_suffixes": unique_keep_order(str(item) for item in annotation["media_suffixes"]),
        "title_candidates": unique_keep_order(str(item) for item in annotation["title_candidates"]),
        "llm_label": annotation["llm_label"] if annotation["llm_label"] is None else str(annotation["llm_label"]),
        "notes": str(annotation["notes"]),
    }
    return cleaned


def validate_batch_annotation(payload: Any, units: list[dict[str, Any]]) -> list[dict[str, Any]]:
    if not isinstance(payload, dict):
        raise RuntimeError("batch annotation must be a JSON object")
    annotations = payload.get("annotations")
    if not isinstance(annotations, list):
        raise RuntimeError("batch annotation missing annotations list")
    expected_ids = [unit["unit_id"] for unit in units]
    returned_ids = [
        str(annotation.get("unit_id"))
        for annotation in annotations
        if isinstance(annotation, dict) and annotation.get("unit_id") is not None
    ]
    if sorted(returned_ids) != sorted(expected_ids) or len(returned_ids) != len(expected_ids):
        raise RuntimeError(f"unit_id set mismatch: expected={expected_ids!r} returned={returned_ids!r}")
    annotations_by_id: dict[str, Any] = {}
    for annotation in annotations:
        if not isinstance(annotation, dict):
            raise RuntimeError("annotations entries must be JSON objects")
        unit_id_value = str(annotation.get("unit_id"))
        if unit_id_value in annotations_by_id:
            raise RuntimeError(f"duplicate annotation unit_id: {unit_id_value}")
        annotations_by_id[unit_id_value] = annotation
    cleaned: list[dict[str, Any]] = []
    for unit in units:
        cleaned.append(validate_annotation(annotations_by_id[unit["unit_id"]], unit))
    return cleaned


def annotate_batch(units: list[dict[str, Any]], args: Args, use_llm: bool) -> list[dict[str, Any]]:
    if not use_llm:
        patches: list[dict[str, Any]] = []
        for unit in units:
            patch = heuristic_unit_patch(unit)
            patch.update({"source": HEURISTIC_SOURCE, "fallback": True, "attempts": 0, "errors": []})
            patches.append(patch)
        return patches

    errors: list[str] = []
    for attempt in range(args.retries):
        variants = request_variants(units, args, errors[-1] if errors else None)
        payload = variants[min(attempt, len(variants) - 1)]
        try:
            data = http_post_json(payload, args)
            annotations = validate_batch_annotation(extract_annotation_payload(data), units)
            for annotation in annotations:
                annotation.update({"source": LLM_SOURCE, "fallback": False, "attempts": attempt + 1, "errors": []})
            return annotations
        except (RuntimeError, json.JSONDecodeError, TypeError, ValueError) as exc:
            errors.append(str(exc))
            time.sleep(min(2.0, 0.25 * (attempt + 1)))

    patches = []
    for unit in units:
        patch = heuristic_unit_patch(unit)
        patch.update(
            {
                "source": HEURISTIC_SOURCE,
                "fallback": True,
                "attempts": args.retries,
                "errors": errors[-3:],
                "notes": f"{patch['notes']}; llm_failed_after={args.retries}",
            }
        )
        patches.append(patch)
    return patches


def annotate_unit(unit: dict[str, Any], args: Args, use_llm: bool) -> dict[str, Any]:
    return annotate_batch([unit], args, use_llm)[0]


def chunk_units(units: list[dict[str, Any]], size: int) -> list[list[dict[str, Any]]]:
    return [units[index : index + size] for index in range(0, len(units), size)]


def load_resume_patches(path: Path) -> dict[str, dict[str, Any]]:
    if not path.exists():
        return {}
    completed: dict[str, dict[str, Any]] = {}
    for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        if not line.strip():
            continue
        try:
            row = json.loads(line)
        except json.JSONDecodeError as exc:
            raise SystemExit(f"invalid resume patch JSON in {path}:{line_number}: {exc}") from exc
        if not isinstance(row, dict):
            continue
        unit_id_value = row.get("unit_id")
        if not unit_id_value:
            continue
        status = row.get("status")
        is_legacy_completed = status is None and ("source" in row or "fallback" in row)
        is_completed = status in {"ok", "fallback"} or is_legacy_completed
        if is_completed:
            completed[str(unit_id_value)] = row
    return completed


def patch_from_resume_row(unit: dict[str, Any], row: dict[str, Any]) -> dict[str, Any]:
    return {
        "unit_id": unit["unit_id"],
        "terminal_ids": unit["terminal_ids"],
        "episode_title_suffixes": unique_keep_order(str(item) for item in row.get("episode_title_suffixes") or []),
        "media_suffixes": unique_keep_order(str(item) for item in row.get("media_suffixes") or []),
        "title_candidates": unique_keep_order(str(item) for item in row.get("title_candidates") or []),
        "llm_label": row.get("llm_label"),
        "notes": str(row.get("notes", "")),
        "source": str(row.get("source") or HEURISTIC_SOURCE),
        "fallback": bool(row.get("fallback", row.get("status") == "fallback")),
        "attempts": int_field(row, "attempts"),
        "errors": list(row.get("errors") or []),
    }


def patch_row(unit: dict[str, Any], patch: dict[str, Any]) -> dict[str, Any]:
    fallback = bool(patch["fallback"])
    return {
        "unit_id": unit["unit_id"],
        "node_id": unit["node_id"],
        "kind": unit["kind"],
        "incoming_count": unit["incoming_count"],
        "reachable_terminals": unit["reachable_terminals"],
        "reachable_weight": unit["reachable_weight"],
        "terminal_ids": patch["terminal_ids"],
        "episode_title_suffixes": patch["episode_title_suffixes"],
        "media_suffixes": patch["media_suffixes"],
        "title_candidates": patch["title_candidates"],
        "llm_label": patch["llm_label"],
        "notes": patch["notes"],
        "source": patch["source"],
        "fallback": fallback,
        "status": "fallback" if fallback else "ok",
        "attempts": patch["attempts"],
        "errors": patch["errors"],
        "annotated_at": datetime.now(timezone.utc).isoformat(),
    }


def unit_source_id(unit: dict[str, Any]) -> str:
    node_id_value = unit.get("node_id")
    if node_id_value is None:
        return str(unit["unit_id"])
    return str(node_id_value)


def unit_public_row(unit: dict[str, Any]) -> dict[str, Any]:
    return {
        "unit_id": unit["unit_id"],
        "source_kind": "prefix_dag",
        "source_id": unit_source_id(unit),
        "terminal_ids": unit["terminal_ids"],
        "weight": int(unit["reachable_weight"]),
        "context": {
            "prefixes": unit["prefix_examples"],
            "digit_skeletons": unit["digit_skeleton_examples"],
            "edge_labels": unit["common_edge_labels"],
            "notes": (
                f"kind={unit['kind']}; incoming_count={unit['incoming_count']}; "
                f"reachable_terminals={unit['reachable_terminals']}"
            ),
        },
        "examples": {
            "values": unit["value_examples"],
            "suffixes": unit["suffix_examples"],
        },
        "expected_output": {
            "schema_version": "dmhy-annotation-v1",
        },
    }


def dataset_records(
    units: list[dict[str, Any]],
    patches: list[dict[str, Any]],
    args: Args,
) -> list[dict[str, Any]]:
    tokenizer = AnimeTokenizer()
    records: list[dict[str, Any]] = []
    seen: set[str] = set()
    for unit, patch in zip(units, patches):
        per_terminal_counts: dict[str, int] = {}
        for terminal in unit["_terminals"]:
            terminal_id_value = str(terminal["_terminal_id"])
            for source_index, filename in enumerate(string_list(terminal.get("value_examples"))):
                if len(records) >= args.max_records:
                    return records
                if per_terminal_counts.get(terminal_id_value, 0) >= args.records_per_terminal:
                    break
                if filename in seen:
                    continue
                seen.add(filename)
                sample = weak_label_filename(filename, tokenizer)
                if sample is None:
                    continue
                tokens, labels = normalize_generated_tokens(
                    sample["tokens"],
                    sample["labels"],
                    preserve_i_labels=args.preserve_i_labels,
                )
                records.append(
                    {
                        "file_id": f"prefix-dag:{unit['unit_id']}:{terminal_id_value}:{source_index}",
                        "filename": filename,
                        "tokens": tokens,
                        "labels": labels,
                        "terminal_id": terminal_id_value,
                        "terminal_index": terminal["_terminal_index"],
                        "unit_id": unit["unit_id"],
                        "source": patch["source"],
                        "needs_llm_review": patch["fallback"],
                        "episode_title_suffixes": patch["episode_title_suffixes"],
                        "media_suffixes": patch["media_suffixes"],
                        "title_candidates": patch["title_candidates"],
                        "annotations": {
                            "unit_id": unit["unit_id"],
                            "terminal_id": terminal_id_value,
                            "terminal_index": terminal["_terminal_index"],
                            "source": patch["source"],
                            "fallback": patch["fallback"],
                            "episode_title_suffixes": patch["episode_title_suffixes"],
                            "media_suffixes": patch["media_suffixes"],
                            "title_candidates": patch["title_candidates"],
                            "llm_label": patch["llm_label"],
                            "notes": patch["notes"],
                        },
                    }
                )
                per_terminal_counts[terminal_id_value] = per_terminal_counts.get(terminal_id_value, 0) + 1
    return records


def annotate_units(
    units: list[dict[str, Any]],
    args: Args,
) -> tuple[list[tuple[dict[str, Any], dict[str, Any]]], int, int, int]:
    resume_rows = load_resume_patches(args.patch_output) if args.resume else {}
    results: list[dict[str, Any] | None] = [None] * len(units)
    pending: list[tuple[int, dict[str, Any]]] = []
    skipped = 0
    for index, unit in enumerate(units):
        row = resume_rows.get(unit["unit_id"])
        if row is not None and [str(item) for item in row.get("terminal_ids") or []] == unit["terminal_ids"]:
            results[index] = patch_from_resume_row(unit, row)
            skipped += 1
        else:
            pending.append((index, unit))

    batch_size = max(1, min(args.units_per_request, args.max_context_units))
    batches = [
        ([index for index, _unit in batch], [unit for _index, unit in batch])
        for batch in chunk_units(pending, batch_size)
    ]
    if args.llm:
        runnable_batches = batches[: args.max_requests]
        llm_requested = len(runnable_batches)
    else:
        runnable_batches = batches
        llm_requested = 0
    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        futures = {
            executor.submit(annotate_batch, batch_units, args, args.llm): indices
            for indices, batch_units in runnable_batches
        }
        for future in as_completed(futures):
            indices = futures[future]
            for index, patch in zip(indices, future.result()):
                results[index] = patch
    completed = [
        (unit, patch)
        for unit, patch in zip(units, results)
        if patch is not None
    ]
    pending_unprocessed = len(units) - len(completed)
    return completed, skipped, llm_requested, pending_unprocessed


def main() -> None:
    args = parse_args()
    if args.llm and not args.api_key:
        raise SystemExit("--llm requires --api-key or ANIFILEBERT_LLM_API_KEY")

    dag = load_dag(args.dag)
    units = select_units(dag, args)
    if not units:
        raise SystemExit("no DAG units selected; adjust --limit/--min-incoming-count")

    completed_pairs, resume_skipped, llm_requested, pending_unprocessed = annotate_units(units, args)
    completed_units = [unit for unit, _patch in completed_pairs]
    patches = [patch for _unit, patch in completed_pairs]
    patch_rows = [patch_row(unit, patch) for unit, patch in completed_pairs]
    records = dataset_records(completed_units, patches, args)

    patch_count = write_jsonl(args.patch_output, patch_rows)
    record_count = write_jsonl(args.output, records)
    unit_count = 0
    if args.units_output is not None:
        unit_count = write_jsonl(args.units_output, (unit_public_row(unit) for unit in units))

    fallback_count = sum(1 for patch in patches if patch["fallback"])
    terminal_count = len(all_terminals(dag))
    unique_coverage_count = len(unique_terminal_coverage(units))
    summary = {
        "dag": str(args.dag),
        "coverage_mode": args.coverage_mode,
        "units_output": str(args.units_output) if args.units_output is not None else None,
        "patch_output": str(args.patch_output),
        "output": str(args.output),
        "selected_units": len(units),
        "dag_terminals": terminal_count,
        "unique_terminal_coverage": unique_coverage_count,
        "written_units": unit_count,
        "written_patches": patch_count,
        "written_records": record_count,
        "llm_requested": llm_requested,
        "pending_unprocessed": pending_unprocessed,
        "units_per_request": args.units_per_request,
        "max_context_units": args.max_context_units,
        "resume": args.resume,
        "resume_skipped": resume_skipped,
        "fallback_count": fallback_count,
        "workers": args.workers,
        "model": args.model if args.llm else None,
    }
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    if fallback_count:
        print(f"warning: {fallback_count} unit(s) used heuristic fallback", file=sys.stderr)


if __name__ == "__main__":
    main()