#!/usr/bin/env python3
"""
service_subgraphs.py

For every service in a graph file (nodes + dependencies) create two
independent Graphviz diagrams:

* downstream – everything that the service can reach.
* upstream   – everything that can reach the service.

The input file must be in the *new* format (produced by
convert_service_yaml.py).  The output is written to a directory you
specify, one file per graph.

Usage:
    python service_subgraphs.py -i services-graph.yaml -o subgraphs -f svg
"""

import argparse
import json
import sys
from pathlib import Path
from typing import Dict, List, Set

import yaml
from graphviz import Digraph


# --------------------------------------------------------------------
# 1. Load the graph (nodes + edges)
# --------------------------------------------------------------------
def load_graph(path: Path) -> Dict[str, Dict]:
    """
    Returns a dict:
        {
            "services": { <id>: {"type":..., "meta":...} },
            "adj":      { <id>: set([child1, child2, ...]) }
        }
    """
    # Detect YAML vs JSON
    if path.suffix.lower() == ".json":
        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)
    else:
        with path.open("r", encoding="utf-8") as f:
            data = yaml.safe_load(f)

    if not data or "services" not in data or "dependencies" not in data:
        raise ValueError(
            "Input file must contain top‑level 'services' and 'dependencies'"
        )

    services: Dict[str, Dict] = {}
    for svc in data["services"]:
        svc_id = svc.get("id") or svc.get("name")
        if not svc_id:
            raise ValueError("Each service must have an 'id'")
        services[svc_id] = {"type": svc.get("type", "unknown"), "meta": svc.get("meta", {})}

    # Build adjacency lists
    adj: Dict[str, Set[str]] = {k: set() for k in services}
    for edge in data["dependencies"]:
        src = edge.get("source") or edge.get("name")
        tgt = edge.get("target") or edge.get("name")
        if src not in services or tgt not in services:
            continue  # silently ignore unknown references
        adj[src].add(tgt)

    return {"services": services, "adj": adj}


# --------------------------------------------------------------------
# 2. Graph traversal helpers
# --------------------------------------------------------------------
def reachable_from(adj: Dict[str, Set[str]], start: str) -> Set[str]:
    """Return all nodes that can be reached from `start` (DFS)."""
    visited = set()
    stack = [start]
    while stack:
        node = stack.pop()
        for nxt in adj.get(node, []):
            if nxt not in visited:
                visited.add(nxt)
                stack.append(nxt)
    return visited


def reverse_adj(adj: Dict[str, Set[str]]) -> Dict[str, Set[str]]:
    """Return the transpose of `adj`."""
    rev: Dict[str, Set[str]] = {k: set() for k in adj}
    for src, targets in adj.items():
        for tgt in targets:
            rev[tgt].add(src)
    return rev


# --------------------------------------------------------------------
# 3. Build a Graphviz subgraph
# --------------------------------------------------------------------
TYPE_STYLE = {
    "airflow":         {"shape": "diamond",       "color": "#8a2be2"},
    "dynamodb":        {"shape": "cylinder",      "color": "#00ced1"},
    "ecs":             {"shape": "box",           "color": "#6cb8e5"},
    "elasticsearch":   {"shape": "rectangle",     "color": "#20b2aa"},
    "kafka":           {"shape": "parallelogram", "color": "#ff6347"},
    "lambda":          {"shape": "ellipse",       "color": "#ff7f50"},
    "neo4j":           {"shape": "pentagon",      "color": "#ff1493"},
    "rds":             {"shape": "cylinder",      "color": "#ffa500"},
    "redis":           {"shape": "octagon",       "color": "#ff4500"},
    "s3":              {"shape": "oval",          "color": "#32cd32"},
    "schema-registry": {"shape": "box",           "color": "#ff8c00"},
    "snowflake":       {"shape": "hexagon",       "color": "#4b0082"},
    "swf":             {"shape": "parallelogram", "color": "#1e90ff"},
    "unknown":         {"shape": "box",           "color": "#777777"},
}

def node_attrs(name: str, svc_meta: Dict) -> Dict[str, str]:
    style = TYPE_STYLE.get(svc_meta["type"], TYPE_STYLE["unknown"])
    tooltip_parts = [f"{k}: {v}" for k, v in svc_meta["meta"].items()]
    tooltip = "\n".join(tooltip_parts)
    return {
        "label": name,
        "shape": style["shape"],
        "style": "filled",
        "fillcolor": style["color"],
        "tooltip": tooltip,
        "title": tooltip,
    }

def build_subgraph(
    root: str,
    nodes_of_interest: Set[str],
    edges_of_interest: Set[tuple],
    services_meta: Dict[str, Dict],
) -> Digraph:
    dot = Digraph(comment=f"Sub‑graph for {root}")
    dot.attr(rankdir="LR")

    # Add only the nodes that belong to this sub‑graph
    for node in nodes_of_interest | {root}:
        meta = services_meta[node]
        dot.node(node, **node_attrs(node, meta))

    # Add edges that stay inside the sub‑graph
    for src, tgt in edges_of_interest:
        if src in nodes_of_interest and tgt in nodes_of_interest:
            rel_type = "depends_on"  # we don’t keep the original label here; you can adjust
            dot.edge(src, tgt, rel_type)

    return dot


# --------------------------------------------------------------------
# 4. CLI
# --------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Create upstream/downstream sub‑graphs for each service."
    )
    parser.add_argument(
        "-i", "--input",
        required=True,
        type=Path,
        help="Input graph file (YAML or JSON) in the new format.",
    )
    parser.add_argument(
        "-o", "--output-dir",
        required=True,
        type=Path,
        help="Directory to store all generated sub‑graphs.",
    )
    parser.add_argument(
        "-f", "--format",
        choices=["svg", "png", "pdf", "dot", "mermaid"],
        default="svg",
        help="Output format for the generated graphs.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    # Load graph
    try:
        data = load_graph(args.input)
    except Exception as exc:
        print(f"[ERROR] Failed to read input: {exc}", file=sys.stderr)
        sys.exit(1)

    services_meta = data["services"]
    adj = data["adj"]
    rev_adj = reverse_adj(adj)

    # Create output directory
    args.output_dir.mkdir(parents=True, exist_ok=True)

    # Pre‑compute all edges as a set of tuples for fast lookup
    all_edges = {(src, tgt) for src, targets in adj.items() for tgt in targets}

    for svc_id in services_meta:
        # ------------------- downstream -------------------
        downstream_nodes = reachable_from(adj, svc_id)
        downstream_dot = build_subgraph(
            root=svc_id,
            nodes_of_interest=downstream_nodes,
            edges_of_interest=all_edges,
            services_meta=services_meta,
        )
        downstream_dot.format = args.format
        downstream_file = args.output_dir / f"{svc_id}-downstream.{args.format}"
        downstream_dot.render(filename=str(downstream_file), cleanup=True)

        # ------------------- upstream -------------------
        upstream_nodes = reachable_from(rev_adj, svc_id)
        upstream_dot = build_subgraph(
            root=svc_id,
            nodes_of_interest=upstream_nodes,
            edges_of_interest=all_edges,
            services_meta=services_meta,
        )
        upstream_dot.format = args.format
        upstream_file = args.output_dir / f"{svc_id}-upstream.{args.format}"
        upstream_dot.render(filename=str(upstream_file), cleanup=True)

    print(f"Created {len(services_meta)*2} sub‑graphs in {args.output_dir}")


if __name__ == "__main__":
    main()