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

Fetches Datadog service‑dependency data for a single service and visualises it with GraphViz.

Usage:
    DD_API_KEY=... DD_APP_KEY=... \
        python service_deps.py <service> -e <environment> [options]

Options:
    -e, --env          Required.  Datadog environment name (e.g. prod, staging)
    --start            Epoch timestamp to start the query from.
                       Default: now – 7 days (one week of data).
    -o, --output       Output filename without extension (default: service_deps).
    -f, --format       GraphViz output format: pdf | png | dot   (default: pdf)
    -d, --depth        How many levels deep to recurse into dependencies.
                       1 = only the immediate calls/called_by. 2+ will walk the graph.
                       Default: 1
"""

from __future__ import annotations

import argparse
import datetime
import json
import os
import sys
import time
from collections import defaultdict, deque
from typing import Dict, List, Set

try:
    from graphviz import Digraph
except ImportError as exc:
    print("graphviz package missing – install with `pip install graphviz`", file=sys.stderr)
    raise SystemExit(1)

# --------------------------------------------------------------------------- #
#  Configuration & helpers
# --------------------------------------------------------------------------- #

def get_api_keys() -> tuple[str, str]:
    """Return DD_API_KEY and DD_APP_KEY from the environment."""
    api_key = os.getenv("DD_API_KEY")
    app_key = os.getenv("DD_APP_KEY")
    if not api_key or not app_key:
        raise RuntimeError(
            "Datadog API keys missing. Set DD_API_KEY and DD_APP_KEY env vars."
        )
    return api_key, app_key


def build_datadog_url(service: str) -> str:
    """Return the full Datadog service‑dependency URL."""
    base = os.getenv("DD_BASE_URL", "https://api.datadoghq.com")
    return f"{base}/api/v1/service_dependencies/{service}"


# --------------------------------------------------------------------------- #
#  API client
# --------------------------------------------------------------------------- #

def fetch_service_dependency(
    service: str,
    env: str,
    start_epoch: int,
    api_key: str,
    app_key: str,
) -> dict:
    """
    Call Datadog’s endpoint and return the parsed JSON.
    Raises RuntimeError on HTTP errors or bad JSON.
    """
    import requests

    url = build_datadog_url(service)
    headers = {
        "Accept": "application/json",
        "DD-API-KEY": api_key,
        "DD-APPLICATION-KEY": app_key,
    }
    params = {"env": env, "start": start_epoch}

    resp = requests.get(url, headers=headers, params=params, timeout=30)
    if resp.status_code != 200:
        raise RuntimeError(
            f"Datadog API error {resp.status_code}: {resp.text}"
        )
    try:
        return resp.json()
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Invalid JSON from Datadog: {exc}") from exc


# --------------------------------------------------------------------------- #
#  Graph building
# --------------------------------------------------------------------------- #

def add_node(
    dot: Digraph, name: str, internal_services: Set[str]
) -> None:
    """
    Add a node to the graph. Internal services (ours) are coloured green,
    external libs/other services are grey.
    """
    if name in internal_services:
        shape = "ellipse"
        fillcolor = "lightgreen"
    else:
        shape = "box"
        fillcolor = "lightgrey"

    dot.node(name, shape=shape, style="filled", fillcolor=fillcolor)


def build_graph(
    root_service: str,
    data_map: Dict[str, dict],
    internal_services: Set[str],
) -> Digraph:
    """
    Build a directed graph from the collected dependency data.
    Each edge is labelled either 'calls' (outgoing) or 'called by' (incoming).
    """
    dot = Digraph(comment=f"Dependencies for {root_service}", format="pdf")
    dot.attr(rankdir="LR", splines="ortho")

    # Root node
    add_node(dot, root_service, internal_services)

    # Walk the collected data map once – we already have all nodes.
    for svc, payload in data_map.items():
        calls = payload.get("calls", [])
        called_by = payload.get("called_by", [])

        for dep in calls:
            add_node(dot, dep, internal_services)
            dot.edge(svc, dep, label="calls")

        for caller in called_by:
            add_node(dot, caller, internal_services)
            dot.edge(caller, svc, label="called by")

    return dot


# --------------------------------------------------------------------------- #
#  Recursive fetching
# --------------------------------------------------------------------------- #

def collect_dependencies(
    service: str,
    env: str,
    start_epoch: int,
    api_key: str,
    app_key: str,
    depth: int,
) -> Dict[str, dict]:
    """
    Walk the dependency graph up to *depth* levels.
    Returns a mapping from service name → payload dict.
    """
    visited: Set[str] = set()
    queue: deque[tuple[str, int]] = deque([(service, 1)])
    data_map: Dict[str, dict] = {}

    while queue:
        current_service, cur_depth = queue.popleft()
        if current_service in visited or cur_depth > depth:
            continue
        visited.add(current_service)

        payload = fetch_service_dependency(
            current_service,
            env,
            start_epoch,
            api_key,
            app_key,
        )
        data_map[current_service] = payload

        # Enqueue the next level (outgoing calls + incoming callers)
        if cur_depth < depth:
            deps = set(payload.get("calls", [])) | set(payload.get("called_by", []))
            for d in deps:
                queue.append((d, cur_depth + 1))

    return data_map


# --------------------------------------------------------------------------- #
#  Main
# --------------------------------------------------------------------------- #

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Show Datadog service dependencies for a week."
    )
    parser.add_argument("service", help="The root service to query")
    parser.add_argument("-e", "--env", required=True, help="Datadog environment name")
    parser.add_argument("--start",
                        type=int,
                        default=None,
                        help="Epoch start timestamp (default: now – 7 days)")
    parser.add_argument("-o", "--output", default="service_deps",
                        help="Base filename for the graph output")
    parser.add_argument("-f", "--format", choices=["pdf", "png", "dot"], default="pdf",
                        help="GraphViz output format")
    parser.add_argument("-d", "--depth", type=int, default=1,
                        help="Recursion depth (1 = only direct dependencies)")

    args = parser.parse_args()

    # Compute start epoch
    if args.start is None:
        now = int(time.time())
        start_epoch = now - 7 * 24 * 60 * 60  # one week ago
    else:
        start_epoch = args.start

    api_key, app_key = get_api_keys()
    env = args.env
    root_service = args.service

    print(
        f"Fetching dependencies for service '{root_service}' "
        f"(env={env}) from {datetime.datetime.utcfromtimestamp(start_epoch)} UTC"
    )

    # ------------------------------------------------------------------- #
    #  Collect data (optionally recursive)
    # ------------------------------------------------------------------- #
    try:
        data_map = collect_dependencies(
            root_service,
            env,
            start_epoch,
            api_key,
            app_key,
            depth=args.depth,
        )
    except Exception as exc:
        print(f"Error while fetching dependencies: {exc}", file=sys.stderr)
        sys.exit(1)

    # --------------------------------------------------------------- #
    #  Build graph
    # --------------------------------------------------------------- #
    internal_services = set(data_map.keys())
    dot_graph = build_graph(root_service, data_map, internal_services)

    # Render
    try:
        output_path = dot_graph.render(
            filename=args.output,
            format=args.format,
            cleanup=True,
        )
        print(f"Graph written to {output_path}")
    except Exception as exc:
        print(f"Failed to render graph: {exc}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()