"""
highlight_seeds_dot.py

Highlight seed nodes (first cluster in a JSON seeds file) inside a GraphViz .dot file.

Usage
-----
    python highlight_seeds_dot.py <input.dot> <seeds.json> <output.dot>

Every node listed in seeds.json[clusters][0]["seed_nodes"] will be rendered
[color="red", style="filled", fillcolor="red", fontcolor="white"]
and all its incident edges colored red.
The script prints *nothing* to stdout – it only writes the modified .dot file.
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path


def load_seeds(seeds_path: Path) -> set[int]:
    """Return the seed-node indices from the first cluster."""
    with seeds_path.open("r", encoding="utf-8") as f:
        data = json.load(f)
    return set(data["clusters"][0]["seed_nodes"])


def highlight_dot(dot_path: Path, seeds: set[int], out_path: Path) -> None:
    """Read dot_path, highlight seeds and their edges, write to out_path."""
    content = dot_path.read_text(encoding="utf-8")

    # 1) highlight seed nodes
    def node_replace(m: re.Match) -> str:
        node_id = int(m.group(1))
        if node_id in seeds:
            return f'{m.group(0).rstrip(";")} [color="red", style="filled", fillcolor="red", fontcolor="white"];'
        return m.group(0)

    content = re.sub(rf'^\s*({"|".join(map(str, seeds))})\s*;', node_replace, content, flags=re.MULTILINE)

    # 2) highlight edges incident to seeds
    def edge_replace(m: re.Match) -> str:
        u, v = int(m.group(1)), int(m.group(2))
        if u in seeds or v in seeds:
            return f'{m.group(0).rstrip(";")} [color="red"];'
        return m.group(0)

    content = re.sub(rf'^\s*({"|".join(map(str, seeds))})\s*--\s*(\d+)\s*;', edge_replace, content, flags=re.MULTILINE)
    content = re.sub(rf'^\s*(\d+)\s*--\s*({"|".join(map(str, seeds))})\s*;', edge_replace, content, flags=re.MULTILINE)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(content, encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.strip(), formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("input_dot", type=str, help="Original GraphViz .dot file")
    parser.add_argument("seeds_json", type=str, help="Seeds JSON file")
    parser.add_argument("output_dot", type=str, help="Path to write highlighted .dot file")
    args = parser.parse_args()

    seeds = load_seeds(Path(args.seeds_json))
    highlight_dot(Path(args.input_dot), seeds, Path(args.output_dot))


if __name__ == "__main__":
    main()