#!/usr/bin/env python3
"""Reconstruct the *root entry services* per endpoint for a Datadog APM service.

This reproduces the leftmost (upstream) column of Datadog's Resource-page
Dependency Map: for each of a service's server endpoints (resources), the
distinct set of services that sit at the **root** of traces flowing through
that endpoint.

There is no API that returns this directly. Datadog's own map is UI-only and
is itself built from a *sample* of ingested spans. We replicate that by
sampling traces per endpoint and reading each trace's root span. Output is
therefore an ESTIMATE — raise --traces-per-endpoint until the root set stops
growing (that is your convergence / confidence signal).

Algorithm:
  1. Enumerate the service's server endpoints by aggregating spans grouped by
     resource_name (filtered to server-entry spans).
  2. For each endpoint, page the Spans Search API to collect distinct trace_ids.
  3. For each trace, find the root span (parent_id empty / "0" / not present
     among the trace's span IDs) and record its service.
  4. Aggregate into { endpoint: { root_service: count } } and emit JSON.

Requires DD_API_KEY, DD_APP_KEY, DD_SITE (default datadoghq.com) env vars.

Usage:
    python dd_root_entry_services.py <service> [options]

Examples:
    python dd_root_entry_services.py ows-product --env prod
    python dd_root_entry_services.py ows-product --from now-7d --to now
    python dd_root_entry_services.py ows-product --traces-per-endpoint 200
    python dd_root_entry_services.py ows-product --dry-run        # first endpoint only
    DD_DEBUG=1 python dd_root_entry_services.py ows-product --dry-run  # dump raw event
"""

import argparse
import json
import logging
import os
import re
import sys
import time
import urllib.error
import urllib.request
from collections import defaultdict
from datetime import datetime, timedelta, timezone


PAGE_LIMIT = 1000

# Server-entry span filters. We probe span.kind:server first (the generic,
# preferred tag); if the org doesn't populate it we fall back to the
# framework entry operations that this org is known to use.
SERVER_KIND_FILTER = "span.kind:server"
OPERATION_FILTER = (
    "(operation_name:flask.request OR operation_name:fastapi.request "
    "OR operation_name:tornado.request)"
)

# Beyond this many spans for a single batched trace_id query we stop paging,
# to guard against a runaway query. 20k spans comfortably covers a batch of
# normal traces.
BATCH_SPAN_CAP = 20000

logging.basicConfig(
    level=logging.DEBUG if os.environ.get("DD_DEBUG") else logging.INFO,
    format="%(levelname)s %(message)s",
    stream=sys.stderr,
)


# ---------------------------------------------------------------------------
# Time helpers
# ---------------------------------------------------------------------------

_UNITS = {"s": 1, "m": 60, "h": 3600, "d": 86400, "w": 604800}

_DT_FORMATS = [
    "%Y-%m-%dT%H:%M:%SZ",
    "%Y-%m-%dT%H:%M:%S%z",
    "%Y-%m-%dT%H:%M:%S",
    "%Y-%m-%d %H:%M:%S",
    "%Y-%m-%d",
]


def to_iso(dt):
    return dt.strftime("%Y-%m-%dT%H:%M:%SZ")


def resolve_when(token, now):
    """Resolve a time token to a tz-aware datetime.

    Accepts:
      - "now"                → now
      - "now-30d" / "now-2h" → relative to now (s/m/h/d/w units)
      - epoch milliseconds   → all-digit string (13ish digits)
      - ISO 8601             → 2024-01-01T00:00:00Z and friends
    """
    token = token.strip()
    if token == "now":
        return now

    m = re.fullmatch(r"now-(\d+)([smhdw])", token)
    if m:
        n, unit = int(m.group(1)), m.group(2)
        return now - timedelta(seconds=n * _UNITS[unit])

    if token.isdigit():
        # Treat as epoch milliseconds (Datadog convention).
        return datetime.fromtimestamp(int(token) / 1000.0, tz=timezone.utc)

    for fmt in _DT_FORMATS:
        try:
            dt = datetime.strptime(token, fmt)
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt
        except ValueError:
            continue

    raise argparse.ArgumentTypeError(
        f"Cannot parse time '{token}'. Use now, now-30d, epoch-ms, or ISO 8601."
    )


# ---------------------------------------------------------------------------
# Span field helpers
# ---------------------------------------------------------------------------

def get_attr(span, key):
    """Read a span attribute, tolerating both raw API and flat shapes."""
    return span.get("attributes", span).get(key)


def span_duration_ns(span):
    """Root-span duration in nanoseconds (custom.duration in the raw API)."""
    a = span.get("attributes", span)
    return (a.get("custom") or {}).get("duration") or 0


def _escape_query_value(value):
    """Escape a value for a double-quoted Datadog query term.

    Backslash and double-quote must be escaped so a resource/path containing
    them (e.g. a user-supplied --path) can't break or alter the query.
    """
    return value.replace("\\", "\\\\").replace('"', '\\"')


# ---------------------------------------------------------------------------
# Endpoint sources (source code / explicit) — alternative to aggregate enumeration
# ---------------------------------------------------------------------------

def make_resource(method, path):
    """Build the Datadog resource_name string for a Flask route.

    DD records server-entry resources as 'METHOD /path' with Flask's param
    syntax preserved, e.g. 'GET /product/<int:product_id>'.
    """
    return f"{method.upper().strip()} {path.strip()}"


HTTP_VERBS = ("get", "post", "put", "patch", "delete", "head", "options")


def parse_routes(filepath):
    """Parse route decorators from a Flask or FastAPI/Starlette handlers file.

    Handles:
      - Flask:   @app.route("/path", methods=["GET", "POST"])
      - FastAPI: @router.get("/path", ...) / @app.post(...) / etc.

    Returns 'METHOD /path' resource strings matching the format Datadog uses in
    span resource_name. Path-param syntax is preserved as written — Flask
    '<int:id>' and FastAPI '{id}' both already match Datadog's resource_name.
    """
    with open(filepath) as f:
        content = f.read()

    resources = []
    seen = set()

    def add(method, path):
        resource = make_resource(method, path)
        if resource not in seen:
            seen.add(resource)
            resources.append(resource)

    # FastAPI / Starlette verb decorators: @router.get("/path", ...). The method
    # is the decorator attribute; the path is the first string literal.
    verb_re = re.compile(
        r'@\w+(?:\.\w+)*\.(' + "|".join(HTTP_VERBS) + r')\(\s*["\']([^"\']+)["\']'
    )
    for m in verb_re.finditer(content):
        add(m.group(1), m.group(2))

    # Flask route decorators: @app.route("/path", methods=["GET", ...]).
    flask_re = re.compile(
        r'@\w+(?:\.\w+)*\.route\(\s*["\']([^"\']+)["\']'
        r'(?:[^)]*?methods\s*=\s*\[([^\]]+)\])?',
        re.DOTALL,
    )
    for m in flask_re.finditer(content):
        path = m.group(1)
        methods_str = m.group(2) or '"GET"'
        methods = re.findall(r'["\']([A-Z]+)["\']', methods_str) or ["GET"]
        for method in methods:
            add(method, path)

    return resources


# ---------------------------------------------------------------------------
# Datadog API
# ---------------------------------------------------------------------------

def _client():
    api_key = os.environ.get("DD_API_KEY")
    app_key = os.environ.get("DD_APP_KEY")
    site = os.environ.get("DD_SITE", "datadoghq.com")
    if not api_key or not app_key:
        print("Error: DD_API_KEY and DD_APP_KEY must be set.", file=sys.stderr)
        sys.exit(1)
    base = f"https://api.{site}"
    headers = {
        "DD-API-KEY": api_key,
        "DD-APPLICATION-KEY": app_key,
        "Content-Type": "application/json",
    }
    return base, headers


def dd_post(path, body, retries=6):
    """POST to a Datadog v2 endpoint with exponential backoff on 429.

    Honors the x-ratelimit-reset header when present, else 15 * 2^attempt.
    """
    base, headers = _client()
    url = f"{base}{path}"
    for attempt in range(retries):
        req = urllib.request.Request(
            url, data=json.dumps(body).encode(), headers=headers, method="POST"
        )
        try:
            with urllib.request.urlopen(req) as resp:
                return json.loads(resp.read())
        except urllib.error.HTTPError as e:
            if e.code == 429:
                reset = (e.headers or {}).get("x-ratelimit-reset")
                try:
                    wait = float(reset) if reset else 15.0 * (2 ** attempt)
                except (TypeError, ValueError):
                    wait = 15.0 * (2 ** attempt)
                logging.warning(
                    f"Rate limited (429). Waiting {wait:.0f}s before retry "
                    f"{attempt + 1}/{retries}."
                )
                time.sleep(wait)
            else:
                logging.error(f"API error {e.code}: {e.read().decode()}")
                sys.exit(1)
    logging.error("Exhausted retries after repeated 429s.")
    sys.exit(1)


def search_spans_page(query, frm, to, cursor=None, sort="-timestamp", limit=PAGE_LIMIT):
    """One page of /api/v2/spans/events/search. Returns (spans, next_cursor)."""
    page = {"limit": limit}
    if cursor:
        page["cursor"] = cursor
    body = {
        "data": {
            "type": "search_request",
            "attributes": {
                "filter": {"query": query, "from": frm, "to": to},
                "sort": sort,
                "page": page,
            },
        }
    }
    data = dd_post("/api/v2/spans/events/search", body)
    spans = data.get("data", [])
    next_cursor = ((data.get("meta") or {}).get("page") or {}).get("after")
    return spans, next_cursor


_RAW_EVENT_DUMPED = False


def _maybe_dump_raw_event(spans):
    """Under DD_DEBUG, print one raw event so field paths can be verified."""
    global _RAW_EVENT_DUMPED
    if _RAW_EVENT_DUMPED or not os.environ.get("DD_DEBUG") or not spans:
        return
    _RAW_EVENT_DUMPED = True
    print("\n=== RAW data[0] (verify field paths) ===", file=sys.stderr)
    print(json.dumps(spans[0], indent=2)[:4000], file=sys.stderr)
    print("--- resolved via get_attr() ---", file=sys.stderr)
    for k in ("trace_id", "span_id", "parent_id", "service", "resource_name"):
        print(f"  {k:<13}: {get_attr(spans[0], k)!r}", file=sys.stderr)
    print("=========================================\n", file=sys.stderr)


# ---------------------------------------------------------------------------
# Step 1 — enumerate server endpoints
# ---------------------------------------------------------------------------

def _aggregate_resources(query, frm, to, limit):
    body = {
        "data": {
            "type": "aggregate_request",
            "attributes": {
                "filter": {"query": query, "from": frm, "to": to},
                "group_by": [{"facet": "resource_name", "limit": limit}],
                "compute": [{"aggregation": "count", "type": "total"}],
            },
        }
    }
    resp = dd_post("/api/v2/spans/analytics/aggregate", body)
    # Response shape: data is a LIST of buckets, each
    #   {"attributes": {"by": {"resource_name": ...}, "compute": {"c0": <count>}}}
    buckets = resp.get("data") or []
    rows = []
    for b in buckets:
        attrs = b.get("attributes") or {}
        name = (attrs.get("by") or {}).get("resource_name")
        if name:
            count = (attrs.get("compute") or {}).get("c0") or 0
            rows.append((name, count))
    # Order by volume desc (the API-side sort schema is finicky; sort locally).
    rows.sort(key=lambda r: -r[1])
    return [name for name, _ in rows]


def probe_entry_filter(service, env, frm, to, override=None):
    """Pick the server-entry span filter the org actually populates.

    Returns the override unchanged; otherwise probes span.kind:server then the
    framework operation filter (one cheap aggregate each) and returns the first
    that matches any spans, falling back to OPERATION_FILTER. Used for the
    explicit-endpoints path so it doesn't blindly assume operation_name and
    return 0 traces in orgs that only populate span.kind:server.
    """
    if override:
        return override
    for entry_filter in (SERVER_KIND_FILTER, OPERATION_FILTER):
        q = f"service:{service} env:{env} {entry_filter}".strip()
        if _aggregate_resources(q, frm, to, limit=1):
            logging.info(f"Probed entry filter: '{entry_filter}'.")
            return entry_filter
    logging.info(f"No server-entry spans matched a probe; using {OPERATION_FILTER}.")
    return OPERATION_FILTER


def list_endpoints(service, env, frm, to, limit, entry_filter_override):
    """Return (endpoints, entry_filter_used).

    Probes span.kind:server first; falls back to framework entry operations
    if that yields nothing (unless the user forced --entry-filter).
    """
    if entry_filter_override:
        q = f"service:{service} env:{env} {entry_filter_override}".strip()
        eps = _aggregate_resources(q, frm, to, limit)
        return eps, entry_filter_override

    for entry_filter in (SERVER_KIND_FILTER, OPERATION_FILTER):
        q = f"service:{service} env:{env} {entry_filter}".strip()
        eps = _aggregate_resources(q, frm, to, limit)
        if eps:
            logging.info(f"Entry filter '{entry_filter}' → {len(eps)} endpoints.")
            return eps, entry_filter
        logging.info(f"Entry filter '{entry_filter}' → 0 endpoints, trying next.")

    return [], OPERATION_FILTER


# ---------------------------------------------------------------------------
# Step 2 — collect distinct trace_ids per endpoint
# ---------------------------------------------------------------------------

def _time_buckets(frm, to, n):
    """Split an ISO window [frm, to] into n contiguous sub-windows.

    Returns a list of (sub_frm_iso, sub_to_iso). Falls back to a single
    window if the ISO strings can't be parsed.
    """
    try:
        start = datetime.strptime(frm, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
        end = datetime.strptime(to, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
    except ValueError:
        return [(frm, to)]
    total = (end - start).total_seconds()
    if n <= 1 or total <= 0:
        return [(frm, to)]
    step = total / n
    windows = []
    for i in range(n):
        a = start + timedelta(seconds=step * i)
        b = start + timedelta(seconds=step * (i + 1)) if i < n - 1 else end
        windows.append((to_iso(a), to_iso(b)))
    return windows


def _collect_trace_ids(query, frm, to, want, seen):
    """Page the search API into `seen` (a dict used as an ordered set)."""
    cursor = None
    start_count = len(seen)
    while len(seen) - start_count < want:
        limit = min(PAGE_LIMIT, want - (len(seen) - start_count))
        spans, cursor = search_spans_page(query, frm, to, cursor=cursor, limit=limit)
        _maybe_dump_raw_event(spans)
        for ev in spans:
            tid = get_attr(ev, "trace_id")
            if tid and tid not in seen:
                seen[tid] = True
        if not cursor or not spans:
            break


def trace_ids_for_endpoint(service, env, entry_filter, resource, frm, to, want,
                           spread_buckets=1):
    """Collect up to `want` distinct trace_ids for one endpoint.

    Default: the most-recent `want` traces (sort:-timestamp). With
    spread_buckets > 1, split the window into that many sub-windows and pull
    a share from each — de-biases recency so older/bursty upstream roots
    (which the most-recent slice misses) can surface.
    """
    query = (
        f'service:{service} env:{env} {entry_filter} '
        f'resource_name:"{_escape_query_value(resource)}"'
    ).strip()
    seen = {}
    if spread_buckets <= 1:
        _collect_trace_ids(query, frm, to, want, seen)
        return list(seen.keys())

    buckets = min(spread_buckets, want)
    per_bucket = -(-want // buckets)  # ceil
    for a, b in _time_buckets(frm, to, buckets):
        _collect_trace_ids(query, a, b, per_bucket, seen)
    return list(seen.keys())[:want]


# ---------------------------------------------------------------------------
# Step 3 — resolve each trace's root service
# ---------------------------------------------------------------------------

def _root_service_for_group(spans):
    """Given all sampled spans of one trace, return the root span's service."""
    if not spans:
        return "<no spans>"
    ids = {get_attr(s, "span_id") for s in spans}
    roots = [
        s for s in spans
        if (lambda p: not p or p == "0" or p not in ids)(get_attr(s, "parent_id"))
    ]
    if not roots:
        # Every span's parent is present (a cycle, or the true root wasn't
        # sampled and no topmost span surfaced). Can't determine origin.
        return "<root not in sample>"
    root = max(roots, key=span_duration_ns)
    return get_attr(root, "service") or "<unknown>"


def _fetch_batch(trace_ids, frm, to):
    """Fetch all spans for a set of trace_ids, grouped by trace_id.

    Returns (by_trace, cap_hit) where cap_hit is True only if paging stopped
    at BATCH_SPAN_CAP with more spans still available (i.e. truncated).
    """
    query = (f"trace_id:{trace_ids[0]}" if len(trace_ids) == 1
             else "trace_id:(" + " OR ".join(trace_ids) + ")")
    by_trace = defaultdict(list)
    cursor = None
    fetched = 0
    cap_hit = False
    while True:
        spans, cursor = search_spans_page(query, frm, to, cursor=cursor, sort="timestamp")
        _maybe_dump_raw_event(spans)
        for s in spans:
            tid = get_attr(s, "trace_id")
            if tid:
                by_trace[tid].append(s)
        fetched += len(spans)
        if not cursor or not spans:
            break                    # got everything for this batch
        if fetched >= BATCH_SPAN_CAP:
            cap_hit = True           # stopped early — more spans remain
            break
    return by_trace, cap_hit


def _resolve_batch(trace_ids, frm, to, counts, auto_shrink):
    """Resolve roots for a batch; on a span-cap trip, split and retry.

    Splitting keeps each sub-query under the cap so truncated traces don't get
    mislabeled `<no spans>` / a wrong (too-deep) root. A single trace that still
    exceeds the cap is accepted: spans are paged timestamp-ascending, so the
    earliest (root) span is almost always within the kept window.
    """
    by_trace, cap_hit = _fetch_batch(trace_ids, frm, to)
    if cap_hit and auto_shrink and len(trace_ids) > 1:
        mid = len(trace_ids) // 2
        logging.info(
            f"Span cap {BATCH_SPAN_CAP} hit for batch of {len(trace_ids)}; "
            f"splitting into {mid} + {len(trace_ids) - mid} and retrying."
        )
        _resolve_batch(trace_ids[:mid], frm, to, counts, auto_shrink)
        _resolve_batch(trace_ids[mid:], frm, to, counts, auto_shrink)
        return
    if cap_hit:
        logging.warning(
            f"Span cap {BATCH_SPAN_CAP} hit for {len(trace_ids)} trace(s); "
            f"some roots may be approximate."
            + ("" if auto_shrink else " (--no-auto-shrink)")
        )
    for tid in trace_ids:
        counts[_root_service_for_group(by_trace.get(tid, []))] += 1


def root_services_for_traces(trace_ids, frm, to, batch_size, auto_shrink=True):
    """Map a list of trace_ids to a {root_service: count} dict.

    Batches trace_ids into OR-queries to cut API calls. Each batch is fully
    paged, spans grouped by trace_id locally to find each root. When a batch
    overruns the span cap and auto_shrink is on, it's split and retried so no
    trace is silently dropped.
    """
    counts = defaultdict(int)
    for i in range(0, len(trace_ids), batch_size):
        _resolve_batch(trace_ids[i:i + batch_size], frm, to, counts, auto_shrink)
    return dict(counts)


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------

def build_map(service, env, frm, to, traces_per_endpoint, max_endpoints,
              batch_size, entry_filter_override, dry_run, explicit_endpoints=None,
              spread_buckets=1, auto_shrink=True):
    if explicit_endpoints is not None:
        # Endpoints supplied from source / CLI — skip aggregate enumeration, but
        # still probe for the entry-span filter this org populates.
        endpoints = explicit_endpoints
        entry_filter_used = probe_entry_filter(
            service, env, frm, to, entry_filter_override
        )
        logging.info(
            f"Using {len(endpoints)} supplied endpoint(s); skipping aggregate "
            f"enumeration (entry filter: {entry_filter_used})."
        )
    else:
        endpoints, entry_filter_used = list_endpoints(
            service, env, frm, to, limit=1000,
            entry_filter_override=entry_filter_override,
        )

    if not endpoints:
        logging.warning(
            "No server endpoints found. The time window may exceed raw-span "
            "retention for the Spans Search API (often ~15d) — try a shorter "
            "--from (e.g. now-7d), or check --service/--env/--entry-filter."
        )
        return {}, entry_filter_used

    if dry_run:
        endpoints = endpoints[:1]
        logging.info(f"--dry-run: processing only first endpoint '{endpoints[0]}'.")
    elif max_endpoints:
        endpoints = endpoints[:max_endpoints]

    result = {}
    empty_endpoints = 0
    for idx, ep in enumerate(endpoints, 1):
        logging.info(f"[{idx}/{len(endpoints)}] {ep}")
        tids = trace_ids_for_endpoint(
            service, env, entry_filter_used, ep, frm, to, traces_per_endpoint,
            spread_buckets=spread_buckets,
        )
        logging.info(f"    {len(tids)} trace(s) sampled.")
        if not tids:
            empty_endpoints += 1
            result[ep] = {}
            continue
        counts = root_services_for_traces(tids, frm, to, batch_size, auto_shrink)
        result[ep] = dict(sorted(counts.items(), key=lambda kv: -kv[1]))

    if empty_endpoints == len(endpoints):
        if explicit_endpoints is not None:
            logging.warning(
                "Every supplied endpoint returned 0 traces. Check the resource "
                "format matches Datadog ('METHOD /path' with Flask params, e.g. "
                "'GET /product/<int:product_id>'), the --method, and that the "
                "--from window is within raw-span retention (try now-7d)."
            )
        else:
            logging.warning(
                "Every endpoint returned 0 traces. The window likely exceeds "
                "raw-span retention for Spans Search — try a shorter --from (e.g. now-7d)."
            )

    return result, entry_filter_used


# ---------------------------------------------------------------------------
# Output
# ---------------------------------------------------------------------------

def _is_meta(svc):
    """True for diagnostic buckets like <no spans> / <root not in sample>."""
    return svc.startswith("<") and svc.endswith(">")


def warn_low_count(result, threshold):
    """Log a warning listing low-count roots that should be verified.

    Low-count roots are real but sampled few times — confidence is low, so the
    user should confirm them against Datadog's dependency map before trusting.
    Diagnostic <...> buckets are excluded.
    """
    if threshold <= 0:
        return
    for ep, roots in result.items():
        low = [(svc, c) for svc, c in roots.items()
               if c <= threshold and not _is_meta(svc)]
        if low:
            listed = ", ".join(f"{svc} ({c})" for svc, c in low)
            logging.warning(
                f"{ep}: {len(low)} low-confidence root(s) with count <= {threshold} "
                f"— verify against the Datadog dependency map before trusting: {listed}"
            )


def to_csv(result):
    rows = ["endpoint,root_service,count"]
    for ep in sorted(result):
        for svc, count in sorted(result[ep].items(), key=lambda kv: -kv[1]):
            rows.append(f"{_csv(ep)},{_csv(svc)},{count}")
    return "\n".join(rows)


def _csv(val):
    val = str(val)
    if "," in val or '"' in val:
        val = '"' + val.replace('"', '""') + '"'
    return val


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

def main():
    parser = argparse.ArgumentParser(
        description="Estimate root entry services per endpoint for a Datadog APM service."
    )
    parser.add_argument("service", help="Target service name (e.g. ows-product)")
    parser.add_argument("--env", default="prod", help="Environment tag (default: prod)")
    parser.add_argument("--from", dest="frm", default="now-30d",
                        help="Window start: now-30d | epoch-ms | ISO 8601 (default: now-30d)")
    parser.add_argument("--to", dest="to", default="now",
                        help="Window end: now | epoch-ms | ISO 8601 (default: now)")
    parser.add_argument("--traces-per-endpoint", dest="traces_per_endpoint",
                        type=int, default=240,
                        help="Traces to sample per endpoint (default: 240)")
    parser.add_argument("--max-endpoints", dest="max_endpoints", type=int, default=None,
                        help="Only process the top N endpoints by volume")
    parser.add_argument("--spread", dest="spread", action=argparse.BooleanOptionalAction,
                        default=True,
                        help="Sample traces spread across the time window instead of "
                             "just the most recent — surfaces older/bursty upstream roots "
                             "(default: on; use --no-spread for recent-only)")
    parser.add_argument("--spread-buckets", dest="spread_buckets", type=int, default=24,
                        help="Number of sub-windows to spread across (default: 24)")
    parser.add_argument("--batch-size", dest="batch_size", type=int, default=10,
                        help="trace_ids per OR-query when resolving roots (default: 10; "
                             "use 1 for simple per-trace fetches)")
    parser.add_argument("--auto-shrink", dest="auto_shrink",
                        action=argparse.BooleanOptionalAction, default=True,
                        help="On a per-batch span-cap trip, split the batch and retry "
                             "so large traces aren't dropped as <no spans> "
                             "(default: on; --no-auto-shrink to disable)")
    parser.add_argument("--low-count-threshold", dest="low_count_threshold",
                        type=int, default=2,
                        help="Warn to verify root services sampled <= N times "
                             "(default: 2; 0 disables the warning)")
    parser.add_argument("--entry-filter", dest="entry_filter", default=None,
                        help="Override the server-entry span filter "
                             "(default: probe span.kind:server, fall back to "
                             "operation_name:flask/fastapi/tornado.request)")

    src_group = parser.add_argument_group(
        "Endpoint source (optional — skip aggregate enumeration)"
    )
    src_group.add_argument("--handlers", dest="handlers", metavar="FILE",
                           help="Read endpoints from a Flask handlers.py "
                                "(@app.route decorators) instead of querying DD")
    src_group.add_argument("--path", dest="path", metavar="PATH",
                           help="Look up a single endpoint by path, e.g. "
                                "'/product/<int:product_id>'")
    src_group.add_argument("--method", dest="method", default="GET", metavar="METHOD",
                           help="HTTP method for --path (default: GET)")
    parser.add_argument("--dry-run", action="store_true",
                        help="Process only the first endpoint (cheap smoke test)")
    parser.add_argument("--csv", action="store_true",
                        help="Emit CSV (endpoint,root_service,count) instead of JSON")
    parser.add_argument("--out", "-o", dest="out", default=None, metavar="FILE",
                        help="Write output to FILE in addition to stdout")
    args = parser.parse_args()

    if args.batch_size < 1:
        parser.error("--batch-size must be >= 1")
    if args.handlers and args.path:
        parser.error("Use either --handlers or --path, not both.")

    # Spread is on by default; --no-spread falls back to recent-only (1 bucket).
    if args.spread:
        if args.spread_buckets < 1:
            parser.error("--spread-buckets must be >= 1")
        spread_buckets = args.spread_buckets
    else:
        spread_buckets = 1

    # Resolve endpoints from source/CLI when provided (skips aggregate enumeration).
    explicit_endpoints = None
    endpoint_source = "aggregate"
    if args.handlers:
        explicit_endpoints = parse_routes(args.handlers)
        endpoint_source = f"handlers:{args.handlers}"
        logging.info(f"Parsed {len(explicit_endpoints)} routes from {args.handlers}")
        if not explicit_endpoints:
            parser.error(f"No Flask/FastAPI route decorators found in {args.handlers}")
    elif args.path:
        explicit_endpoints = [make_resource(args.method, args.path)]
        endpoint_source = "explicit"
        logging.info(f"Single endpoint: {explicit_endpoints[0]}")

    now = datetime.now(timezone.utc)
    start = resolve_when(args.frm, now)
    end = resolve_when(args.to, now)
    frm_iso, to_iso_str = to_iso(start), to_iso(end)

    result, entry_filter_used = build_map(
        service=args.service,
        env=args.env,
        frm=frm_iso,
        to=to_iso_str,
        traces_per_endpoint=args.traces_per_endpoint,
        max_endpoints=args.max_endpoints,
        batch_size=args.batch_size,
        entry_filter_override=args.entry_filter,
        dry_run=args.dry_run,
        explicit_endpoints=explicit_endpoints,
        spread_buckets=spread_buckets,
        auto_shrink=args.auto_shrink,
    )

    warn_low_count(result, args.low_count_threshold)

    if args.csv:
        output = to_csv(result)
    else:
        caveat = (
            "Sampled reconstruction of the APM Resource Dependency Map's "
            "upstream column. Raise --traces-per-endpoint until the root "
            "set stops growing (convergence)."
        )
        if args.low_count_threshold > 0:
            caveat += (
                f" Roots with count <= {args.low_count_threshold} are "
                "low-confidence — verify them against the Datadog dependency map."
            )
        envelope = {
            "estimate": True,
            "caveat": caveat,
            "service": args.service,
            "env": args.env,
            "from": frm_iso,
            "to": to_iso_str,
            "traces_per_endpoint": args.traces_per_endpoint,
            "sampling": ("recent" if spread_buckets <= 1
                         else f"spread/{spread_buckets} buckets"),
            "low_count_threshold": args.low_count_threshold,
            "entry_filter_used": entry_filter_used,
            "endpoint_source": endpoint_source,
            "endpoints": result,
        }
        output = json.dumps(envelope, indent=2)

    print(output)
    if args.out:
        with open(args.out, "w") as f:
            f.write(output + "\n")
        logging.info(f"Saved to {args.out}")


if __name__ == "__main__":
    main()