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

Visualise dependencies between ECS services, Lambda functions,
RDS, S3, Airflow DAGs and Snowflake warehouses.

The input file must be in the *new* graph format (nodes + edges)
produced by `convert_service_yaml.py`.

Usage:
    python dependency_graph.py --input services-graph.yaml --output graph.svg
"""

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

import yaml
from graphviz import Digraph

# --------------------------------------------------------------------
# 1.  Data structures
# --------------------------------------------------------------------
class Service:
    """Represents one node in the dependency graph."""
    def __init__(
        self,
        name: str,
        svc_type: str,
        meta: Dict[str, str],
        dependencies: List[Dict[str, str]],
    ):
        self.name = name
        self.type = svc_type
        self.meta = meta or {}
        # each dependency is a dict: {"name": <target>, "type": <relation>}
        self.dependencies = dependencies or []

    def __repr__(self) -> str:
        return f"Service(name={self.name!r}, type={self.type!r})"


# --------------------------------------------------------------------
# 2.  Parsing the YAML/JSON file
# --------------------------------------------------------------------
def load_graph(path: Path) -> Dict[str, Service]:
    """Load a graph from a *new* format (nodes + edges)."""
    # Decide whether to parse as YAML or 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' keys"
        )

    # Build node map first
    services: Dict[str, Service] = {}
    for svc_dict in data["services"]:
        name = svc_dict.get("id") or svc_dict.get("name")
        if not name:
            raise ValueError("Each service must have an 'id' (or legacy 'name')")
        svc_type = svc_dict.get("type", "unknown")
        meta = svc_dict.get("meta", {})
        services[name] = Service(name, svc_type, meta, dependencies=[])

    # Build edges: add them to the appropriate source node
    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:
            raise ValueError(f"Edge source {src!r} is unknown")
        if tgt not in services:
            raise ValueError(f"Edge target {tgt!r} is unknown")

        services[src].dependencies.append({"name": tgt, "type": rel})

    return services


# --------------------------------------------------------------------
# 3.  Graph theory helpers
# --------------------------------------------------------------------
def detect_cycles(services: Dict[str, Service]) -> List[List[str]]:
    """Return a list of cycles (each a list of node names)."""
    visited: Set[str] = set()
    stack: List[str] = []
    cycles: List[List[str]] = []

    def dfs(node: str):
        visited.add(node)
        stack.append(node)
        for dep in services[node].dependencies:
            neigh = dep["name"]
            if neigh not in visited:
                dfs(neigh)
            elif neigh in stack:
                idx = stack.index(neigh)
                cycles.append(stack[idx:] + [neigh])
        stack.pop()

    for svc in services:
        if svc not in visited:
            dfs(svc)
    return cycles


def topological_order(services: Dict[str, Service]) -> List[str]:
    """Return a topological ordering of services (empty list if cycle)."""
    indegree: Dict[str, int] = {n: 0 for n in services}
    for svc in services.values():
        for d in svc.dependencies:
            indegree[d["name"]] += 1

    queue = [n for n, deg in indegree.items() if deg == 0]
    order: List[str] = []

    while queue:
        n = queue.pop()
        order.append(n)
        for d in services[n].dependencies:
            indegree[d["name"]] -= 1
            if indegree[d["name"]] == 0:
                queue.append(d["name"])

    if len(order) != len(services):
        return []  # cycle detected
    return order


# --------------------------------------------------------------------
# 4.  Graphviz rendering
# --------------------------------------------------------------------
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(service: Service) -> Dict[str, str]:
    """Return the Graphviz attributes for a node."""
    style = TYPE_STYLE.get(service.type, TYPE_STYLE["unknown"])
    tooltip_parts = [f"{k}: {v}" for k, v in service.meta.items()]
    tooltip = "\n".join(tooltip_parts) if tooltip_parts else ""
    return {
        "label": service.name,
        "shape": style["shape"],
        "style": "filled",
        "fillcolor": style["color"],
        "tooltip": tooltip,
        "title": tooltip,
    }

def build_graph(services: Dict[str, Service]) -> Digraph:
    """Return a graphviz.Digraph object representing the dependencies."""
    dot = Digraph(comment="Service dependency graph")
    dot.attr(rankdir="LR")  # left → right

    # Add all nodes first
    for svc in services.values():
        dot.node(svc.name, **node_attrs(svc))

    # Add edges – label is the relation type
    for svc in services.values():
        for dep in svc.dependencies:
            tgt = dep["name"]
            rel_type = dep.get("type", "depends_on")
            dot.edge(svc.name, tgt, rel_type)

    return dot


# --------------------------------------------------------------------
# 5.  CLI
# --------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Visualise service dependencies (ecs, lambda, rds, s3, airflow, snowflake)."
    )
    parser.add_argument(
        "--input",
        "-i",
        required=True,
        type=Path,
        help="Path to YAML/JSON file describing services in the new format.",
    )
    parser.add_argument(
        "--output",
        "-o",
        required=True,
        type=Path,
        help="Path of the output file (SVG, PNG, PDF, DOT).",
    )
    parser.add_argument(
        "--format",
        "-f",
        choices=["svg", "png", "pdf", "dot", "mermaid"],
        default="svg",
        help="Graphviz output format.",
    )
    parser.add_argument(
        "--detect-cycles",
        "-c",
        action="store_true",
        help="Print any detected dependency cycles.",
    )
    parser.add_argument(
        "--topo",
        action="store_true",
        help="Print a topological ordering (empty if cycles exist).",
    )
    return parser.parse_args()


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

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

    if args.detect_cycles:
        cycles = detect_cycles(services)
        if cycles:
            print("Detected cycles:")
            for cycle in cycles:
                print("  → ".join(cycle))
        else:
            print("No cycles detected.")

    if args.topo:
        order = topological_order(services)
        if order:
            print("Topological order:")
            print(" → ".join(order))
        else:
            print("Cannot compute topological order (cycle detected).")

    # Rendering
    if args.format == "mermaid":
        mermaid_lines = ["graph TD"]
        for svc in services.values():
            for dep in svc.dependencies:
                rel = dep.get("type", "depends_on")
                mermaid_lines.append(f"    {svc.name} -->|{rel}| {dep['name']}")
        mermaid_text = "\n".join(mermaid_lines)
        try:
            with args.output.open("w", encoding="utf-8") as f:
                f.write(mermaid_text)
            print(f"Mermaid diagram written to {args.output}")
        except Exception as exc:
            print(f"[ERROR] Failed to write Mermaid diagram: {exc}", file=sys.stderr)
            sys.exit(1)
    else:
        dot = build_graph(services)
        dot.format = args.format
        try:
            dot.render(filename=str(args.output), cleanup=True)
            print(f"Graph written to {args.output}")
        except Exception as exc:
            print(f"[ERROR] Failed to render graph: {exc}", file=sys.stderr)
            sys.exit(1)


if __name__ == "__main__":
    main()