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

Given a graph file (the “new” format) and a service name, render a
one‑level dependency diagram with that service in the centre.

The resulting diagram will contain only:
    * the target service
    * its immediate upstream services (incoming edges)
    * its immediate downstream services (outgoing edges)
"""

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

import yaml
from graphviz import Digraph


# --------------------------------------------------------------------
# 1. Load the graph and keep edge labels
# --------------------------------------------------------------------
def load_graph(path: Path) -> Tuple[Dict[str, Dict], Dict[str, Set[str]], Dict[Tuple[str, str], str]]:
    """
    Returns:
        services_meta : {id: {"type":..., "meta":...}}
        adj          : {src: set([tgt1, tgt2])}
        edge_labels  : {(src, tgt): relation_type}
    """
    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_meta: 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_meta[svc_id] = {"type": svc.get("type", "unknown"), "meta": svc.get("meta", {})}

    adj: Dict[str, Set[str]] = {k: set() for k in services_meta}
    edge_labels: Dict[Tuple[str, str], str] = {}

    for edge in data["dependencies"]:
        src = edge.get("source") or edge.get("name")
        tgt = edge.get("target") or edge.get("name")
        rel = edge.get("relation", "depends_on")
        if src not in services_meta or tgt not in services_meta:
            continue  # ignore unknown references
        adj[src].add(tgt)
        edge_labels[(src, tgt)] = rel

    return services_meta, adj, edge_labels


# --------------------------------------------------------------------
# 2. Helper to build the one‑level graph
# --------------------------------------------------------------------
TYPE_STYLE = {
    "airflow":         {"shape": "Msquare",       "color": "#8a2be2"},
    "dynamodb":        {"shape": "cylinder",      "color": "#00ced1"},
    "ecs":             {"shape": "box",           "color": "#6cb8e5"},
    "elasticsearch":   {"shape": "septagon",      "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": "folder",        "color": "#ff8c00"},
    "snowflake":       {"shape": "hexagon",       "color": "#4b0082"},
    "swf":             {"shape": "invhouse",      "color": "#1e90ff"},
    "unknown":         {"shape": "egg",           "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_one_level(
    target: str,
    services_meta: Dict[str, Dict],
    adj: Dict[str, Set[str]],
    edge_labels: Dict[Tuple[str, str], str],
) -> Digraph:
    """
    Returns a Graphviz Digraph containing only the target service
    and its immediate parents / children.
    """
    dot = Digraph(comment=f"One‑level view of {target}")
    dot.attr(rankdir="LR", overlap="false")

    # -------------------------------------------------
    # 1️⃣ Collect node sets
    # -------------------------------------------------
    upstream = adj.get(target, set())          # outgoing edges (downstream)
    downstream = {src for src, tgts in adj.items() if target in tgts}  # incoming

    all_nodes = {target} | upstream | downstream

    # -------------------------------------------------
    # 2️⃣ Add nodes
    # -------------------------------------------------
    for n in all_nodes:
        dot.node(n, **node_attrs(n, services_meta[n]))

    # -------------------------------------------------
    # 3️⃣ Add edges (only those touching the target)
    # -------------------------------------------------
    for src, tgt in edge_labels:
        if src == target or tgt == target:
            rel = edge_labels[(src, tgt)]
            dot.edge(src, tgt, label=rel)

    # -------------------------------------------------
    # 4️⃣ Force layout: upstream on left, downstream on right
    # -------------------------------------------------
    with dot.subgraph(name="cluster_up") as s:
        s.attr(rank="source")
        for n in downstream:
            s.node(n)   # nodes already added – just assign rank

    with dot.subgraph(name="cluster_down") as s:
        s.attr(rank="sink")
        for n in upstream:
            s.node(n)

    return dot


# --------------------------------------------------------------------
# 3. CLI
# --------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Render a one‑level dependency diagram around a single service."
    )
    parser.add_argument(
        "-i", "--input",
        required=True,
        type=Path,
        help="Input graph file (YAML or JSON) in the new format.",
    )
    parser.add_argument(
        "-s", "--service",
        required=True,
        type=str,
        help="The service name/id to centre in the diagram.",
    )
    parser.add_argument(
        "-o", "--output",
        required=True,
        type=Path,
        help="Output file path (any .svg/.png/.pdf/.dot/.mermaid).",
    )
    parser.add_argument(
        "-f", "--format",
        choices=["svg", "png", "pdf", "dot", "mermaid"],
        default="svg",
        help="Graphviz output format.",
    )
    return parser.parse_args()


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

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

    if args.service not in services_meta:
        print(f"[ERROR] Service {args.service!r} not found in the graph.", file=sys.stderr)
        sys.exit(1)

    dot = build_one_level(args.service, services_meta, adj, edge_labels)
    dot.format = args.format
    try:
        dot.render(filename=str(args.output), cleanup=True)
    except Exception as exc:
        print(f"[ERROR] Failed to render graph: {exc}", file=sys.stderr)
        sys.exit(1)

    print(f"Created one‑level diagram for {args.service} → {args.output}")


if __name__ == "__main__":
    main()