#!/usr/bin/env python3
"""Analyze a saved Datadog trace for performance insights.

Reads raw.json or trace.json and surfaces:
- Trace header (root span info)
- Time budget by service
- OWS endpoint breakdown
- Top N slowest spans (global and OWS-only)

Usage:
    python dd_analyze.py trace.json
    python dd_analyze.py trace.json --top 50
    python dd_analyze.py trace.json --ows-only
"""

import argparse
import datetime
import json
import os
import re
import sys
from collections import defaultdict
from typing import Any, IO


class _Tee:
    """Write to multiple file-like objects simultaneously."""

    def __init__(self, *files: IO[str]) -> None:
        self.files = files

    def write(self, data: str) -> None:
        for f in self.files:
            f.write(data)

    def flush(self) -> None:
        for f in self.files:
            f.flush()


def format_dur(ns: float) -> str:
    """Format a nanosecond duration as a human-readable string (ms or s)."""
    ms = ns / 1e6
    if ms >= 1000:
        return f"{round(ms / 1000, 2)}s"
    return f"{round(ms, 2)}ms"


def get_attr(s: dict[str, Any], key: str, default: Any = None) -> Any:
    """Extract a normalized attribute from a span, handling both API formats.

    Datadog spans come in two shapes:
    - Raw API format: top-level keys with an "attributes" wrapper
    - Normalized format: flat keys like "operation", "resource", "start", "end"
    """
    if "attributes" in s:
        a = s["attributes"]
        if key == "duration":
            return (a.get("custom") or {}).get("duration") or 0
        return a.get(key, default)
    else:
        if key == "duration":
            return s.get("duration_ms", 0) * 1e6
        if key == "operation_name": return s.get("operation", default)
        if key == "resource_name": return s.get("resource", default)
        if key == "start_timestamp": return s.get("start", default)
        if key == "end_timestamp": return s.get("end", default)
        return s.get(key, default)


def parse_time(ts_str: str) -> float:
    """Parse a timestamp string into a Unix timestamp (seconds since epoch).

    Accepts ISO 8601 strings (with or without "Z" suffix, with or without
    timezone offset) and falls back to bare float strings (epoch seconds).
    Returns 0.0 if parsing fails.
    """
    if not ts_str:
        return 0.0
    try:
        # replace("Z", "+00:00") is a no-op when "Z" is absent,
        # so this handles both "…Z" and already-offset ISO strings.
        ts = ts_str.replace("Z", "+00:00")
        dt = datetime.datetime.fromisoformat(ts)
        return dt.timestamp()
    except Exception:
        try:
            return float(ts_str)
        except Exception:
            return 0.0


def get_root(spans: list[dict[str, Any]]) -> dict[str, Any] | None:
    """Return the root span of the trace.

    A span is a root if it has no parent_id, its parent_id is "0", or its
    parent_id does not appear in the trace. When multiple roots exist (e.g.
    a Step Function with parallel Lambda invocations), the span with the
    longest duration is returned as the most significant entry point.
    """
    span_ids = {get_attr(s, "span_id") for s in spans}
    roots = [
        s for s in spans
        if not get_attr(s, "parent_id")
        or get_attr(s, "parent_id") == "0"
        or get_attr(s, "parent_id") not in span_ids
    ]
    if not roots:
        return spans[0] if spans else None
    # When multiple roots exist (e.g. Step Function with parallel lambdas),
    # pick the one with the longest duration as the most significant entry point.
    return max(roots, key=span_dur)


def span_dur(s: dict[str, Any]) -> float:
    """Return the duration of a span in nanoseconds."""
    return get_attr(s, "duration") or 0


def span_svc(s: dict[str, Any]) -> str:
    """Return the service name for a span."""
    return get_attr(s, "service") or ""


def span_op(s: dict[str, Any]) -> str:
    """Return the operation name for a span."""
    return get_attr(s, "operation_name") or ""


def span_res(s: dict[str, Any]) -> str:
    """Return the resource name for a span."""
    return get_attr(s, "resource_name") or ""


def span_err(s: dict[str, Any]) -> bool:
    """Return True if the span has an error flag set."""
    return bool(get_attr(s, "error"))


def get_trace_duration_ns(spans: list[dict[str, Any]]) -> float:
    """Compute the wall-clock duration of the full trace in nanoseconds.

    Uses the earliest start and latest end timestamps across all spans.
    Falls back to the maximum single-span duration if timestamps are missing.
    """
    if not spans:
        return 0.0

    earliest = min((parse_time(get_attr(s, "start_timestamp")) for s in spans if get_attr(s, "start_timestamp")), default=0.0)
    latest = max((parse_time(get_attr(s, "end_timestamp")) for s in spans if get_attr(s, "end_timestamp")), default=0.0)

    if earliest and latest and latest >= earliest:
        # Return in nanoseconds to match span_dur format
        return (latest - earliest) * 1e9

    return max(span_dur(s) for s in spans)


def print_header(root: dict[str, Any] | None, trace_dur_ns: float) -> None:
    """Print a summary header for the trace using the root span."""
    if not root:
        return
    dur = span_dur(root)
    print("=" * 70)
    print("TRACE HEADER")
    print("=" * 70)
    print(f"  Service:   {span_svc(root)}")
    print(f"  Resource:  {span_res(root)}")
    print(f"  Operation: {span_op(root)}")
    print(f"  Root Dur:  {format_dur(dur)}")
    print(f"  Trace Dur: {format_dur(trace_dur_ns)} (wall-clock)")
    print(f"  Start:     {get_attr(root, 'start_timestamp')}")
    print(f"  End:       {get_attr(root, 'end_timestamp')}")
    print(f"  Env:       {get_attr(root, 'env')}")
    if span_err(root):
        print("  Status:    ERROR")


def print_service_budget(spans: list[dict[str, Any]], trace_dur_ns: float) -> None:
    """Print total, average, and max duration per service, sorted by max duration."""
    totals: dict[str, dict[str, float]] = defaultdict(lambda: {"total": 0, "count": 0, "max": 0})
    for s in spans:
        svc = span_svc(s)
        dur = span_dur(s)
        totals[svc]["total"] += dur
        totals[svc]["count"] += 1
        totals[svc]["max"] = max(totals[svc]["max"], dur)

    sorted_svcs = sorted(totals.items(), key=lambda x: x[1]["max"], reverse=True)
    trace_ms = trace_dur_ns / 1e6 or 1

    print("\n" + "=" * 70)
    print("TIME BUDGET BY SERVICE")
    print("=" * 70)
    print(f"  {'SERVICE':<30} {'TOTAL':>10}  {'%':>5}  {'COUNT':>6}  {'AVG':>10}  {'MAX':>10}")
    print(f"  {'-'*30} {'-'*10}  {'-'*5}  {'-'*6}  {'-'*10}  {'-'*10}")
    for svc, d in sorted_svcs:
        total_ms = d["total"] / 1e6
        avg_ms = total_ms / d["count"] if d["count"] else 0
        max_ms = d["max"] / 1e6
        pct = (total_ms / trace_ms) * 100
        print(f"  {svc:<30} {format_dur(d['total']):>10}  {pct:>4.1f}%  {d['count']:>6}  {format_dur(avg_ms * 1e6):>10}  {format_dur(d['max']):>10}")


def is_ows_entry(s: dict[str, Any]) -> bool:
    """Return True if the span is an OWS HTTP entry point (flask/fastapi/tornado request)."""
    svc = span_svc(s)
    if not svc.startswith("ows-"):
        return False
    op = span_op(s)
    return "flask.request" in op or "fastapi.request" in op or "tornado.request" in op


def print_ows_endpoints(spans: list[dict[str, Any]]) -> None:
    """Print call counts and timings for each OWS HTTP endpoint, sorted by total duration."""
    groups: dict[tuple[str, str], dict[str, float]] = defaultdict(lambda: {"total": 0, "count": 0, "max": 0, "errors": 0})
    http_verbs = ("GET ", "POST ", "PUT ", "DELETE ", "PATCH ", "OPTIONS ", "HEAD ")

    for s in spans:
        if not is_ows_entry(s):
            continue

        res = span_res(s)

        # Filter out sentry and grass handlers, only keep HTTP verb + path
        if "sentry_sdk.integrations" in res or "grass.api.Handler" in res:
            continue

        if not any(res.startswith(verb) for verb in http_verbs):
            continue

        key = (span_svc(s), res)
        dur = span_dur(s)
        groups[key]["total"] += dur
        groups[key]["count"] += 1
        groups[key]["max"] = max(groups[key]["max"], dur)
        if span_err(s):
            groups[key]["errors"] += 1

    if not groups:
        return

    sorted_eps = sorted(groups.items(), key=lambda x: x[1]["total"], reverse=True)

    print("\n" + "=" * 70)
    print("OWS ENDPOINTS")
    print("=" * 70)
    print(f"  {'SERVICE':<22} {'RESOURCE':<40} {'CALLS':>5}  {'TOTAL':>10}  {'AVG':>10}  {'MAX':>10}  {'ERR':>4}")
    print(f"  {'-'*22} {'-'*40} {'-'*5}  {'-'*10}  {'-'*10}  {'-'*10}  {'-'*4}")
    for (svc, res), d in sorted_eps:
        avg_ms = (d["total"] / 1e6) / d["count"] if d["count"] else 0
        print(
            f"  {svc:<22} {res[:40]:<40} {d['count']:>5}  "
            f"{format_dur(d['total']):>10}  {format_dur(avg_ms * 1e6):>10}  {format_dur(d['max']):>10}  {d['errors']:>4}"
        )


def normalize_resource(resource: str) -> str:
    """Strip path parameters to group calls to the same endpoint."""
    r = resource
    # Replace UUIDs (8-4-4-4-12 hex, or plain long hex strings)
    r = re.sub(r'[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}', '{uuid}', r, flags=re.IGNORECASE)
    r = re.sub(r'[0-9a-f-]{32,}', '{uuid}', r, flags=re.IGNORECASE)
    # Replace Flask/Werkzeug path params like <identity_id>, <int:id>
    r = re.sub(r'<[^>]+>', '{param}', r)
    # Replace bare numeric IDs in path segments
    r = re.sub(r'(?<=/)\d+(?=/|$)', '{id}', r)
    return r


def detect_n1(spans: list[dict[str, Any]], threshold: int = 5) -> list[dict[str, Any]]:
    """Detect N+1 call patterns: many calls to the same parameterized OWS endpoint.

    Groups OWS entry spans by (service, normalized_resource). Any group with
    at least `threshold` calls is flagged. For each candidate, walks up the
    parent chain to find the nearest graphql.resolve ancestor and checks
    whether a batch/dataloader endpoint exists for the same service.
    """
    parent_map = {get_attr(s, "span_id"): s for s in spans if get_attr(s, "span_id")}

    # Group OWS entry spans by (service, normalized_resource)
    groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for s in spans:
        if not is_ows_entry(s):
            continue
        res = span_res(s)
        http_verbs = ("GET ", "POST ", "PUT ", "DELETE ", "PATCH ")
        if not any(res.startswith(v) for v in http_verbs):
            continue
        if "sentry_sdk" in res or "grass.api.Handler" in res:
            continue
        key = (span_svc(s), normalize_resource(res))
        groups[key].append(s)

    # Build set of all resources for batch detection
    all_resources = {span_res(s).lower() for s in spans}

    candidates: list[dict[str, Any]] = []
    for (svc, norm_res), group_spans in groups.items():
        if len(group_spans) < threshold:
            continue

        total_dur = sum(span_dur(s) for s in group_spans)
        avg_dur = total_dur / len(group_spans)

        # Walk up parent_map to find closest graphql.resolve ancestor
        resolver_caller = None
        for s in group_spans[:10]:  # sample up to 10 spans for performance
            cur_id = get_attr(s, "parent_id")
            depth = 0
            while cur_id and cur_id in parent_map and depth < 30:
                ancestor = parent_map[cur_id]
                if span_op(ancestor) == "graphql.resolve":
                    resolver_caller = span_res(ancestor)
                    break
                cur_id = get_attr(ancestor, "parent_id")
                depth += 1
            if resolver_caller:
                break

        # Check if any batch/dataloader endpoint exists for same service
        has_batch = any(
            ("dataloader" in r or "batch" in r or "lookup" in r)
            and svc in r
            for r in all_resources
        )

        candidates.append({
            "service": svc,
            "normalized_resource": norm_res,
            "count": len(group_spans),
            "total_dur": total_dur,
            "avg_dur": avg_dur,
            "resolver_caller": resolver_caller or "unknown",
            "has_batch": has_batch,
        })

    candidates.sort(key=lambda x: x["total_dur"], reverse=True)
    return candidates


def print_n1_report(candidates: list[dict[str, Any]], threshold: int) -> None:
    """Print the N+1 detection report, listing flagged endpoints sorted by total duration."""
    print("\n" + "=" * 70)
    print(f"N+1 CANDIDATES  (threshold: {threshold}+ calls to parameterized endpoint)")
    print("=" * 70)

    if not candidates:
        print("  No N+1 patterns detected above threshold.")
        return

    print(f"  {'SERVICE':<18} {'RESOURCE':<38} {'CALLS':>5}  {'TOTAL':>10}  {'AVG':>10}  {'BATCH?':>6}  CALLER RESOLVER")
    print(f"  {'-'*18} {'-'*38} {'-'*5}  {'-'*10}  {'-'*10}  {'-'*6}  ---------------")
    for c in candidates:
        batch = "YES" if c["has_batch"] else "NO"
        svc = c["service"][:18]
        res = c["normalized_resource"][:38]
        caller = c["resolver_caller"][:40]
        print(
            f"  {svc:<18} {res:<38} {c['count']:>5}  "
            f"{format_dur(c['total_dur']):>10}  {format_dur(c['avg_dur']):>10}  {batch:>6}  {caller}"
        )


def leaf_span_ids(spans: list[dict[str, Any]]) -> set[str]:
    """Return the set of span IDs that have no children (no other span lists them as parent)."""
    parent_ids = {get_attr(s, "parent_id") for s in spans if get_attr(s, "parent_id")}
    return {get_attr(s, "span_id") for s in spans if get_attr(s, "span_id") not in parent_ids}


def top_spans_as_json(
    spans: list[dict[str, Any]],
    top: int,
    ows_only: bool = False,
    leaf_only: bool = False,
) -> list[dict[str, Any]]:
    """Return the top N slowest spans as a list of dicts, suitable for JSON output.

    Optionally filter to OWS spans only and/or leaf spans only (no children).
    """
    filtered = [s for s in spans if span_svc(s).startswith("ows-")] if ows_only else spans
    if leaf_only:
        leaves = leaf_span_ids(spans)
        filtered = [s for s in filtered if get_attr(s, "span_id") in leaves]
    sorted_spans = sorted(filtered, key=span_dur, reverse=True)[:top]
    result = []
    for i, s in enumerate(sorted_spans, 1):
        result.append({
            "rank": i,
            "service": span_svc(s),
            "operation": span_op(s),
            "resource": span_res(s),
            "duration_ms": round(span_dur(s) / 1e6, 2),
            "error": span_err(s),
        })
    return result


def print_top_spans(spans: list[dict[str, Any]], top: int, ows_only: bool = False) -> None:
    """Print the top N slowest spans in a formatted table."""
    label = "OWS SPANS" if ows_only else "SPANS"
    filtered = [s for s in spans if span_svc(s).startswith("ows-")] if ows_only else spans
    sorted_spans = sorted(filtered, key=span_dur, reverse=True)[:top]

    print(f"\n{'='*70}")
    print(f"TOP {top} SLOWEST {label}")
    print("=" * 70)
    print(f"  {'SERVICE':<22} {'OPERATION':<35} {'RESOURCE':<35} {'DUR':>10}  ERR")
    print(f"  {'-'*22} {'-'*35} {'-'*35} {'-'*10}  ---")
    for s in sorted_spans:
        svc = span_svc(s)[:22]
        op = span_op(s)[:35]
        res = span_res(s)[:35]
        err = " X" if span_err(s) else ""
        print(f"  {svc:<22} {op:<35} {res:<35} {format_dur(span_dur(s)):>10}{err}")


def flatten_spans(data: dict[str, Any] | list) -> list[dict[str, Any]]:
    """Flatten raw data into a flat list of span dicts.

    Handles three input shapes:
    - Raw API format: {"data": [...]}
    - Array of spans (normalized or raw)
    - Nested span tree with "children" keys (recursively flattened)
    """
    flat: list[dict[str, Any]] = []

    # Raw API format: {"data": [...]}
    if isinstance(data, dict) and "data" in data and isinstance(data["data"], list):
        items = data["data"]
    # Array of spans (normalized or raw)
    elif isinstance(data, list):
        items = data
    # Root of a nested span tree
    elif isinstance(data, dict):
        items = [data]
    else:
        return []

    # Flatten nested trees recursively
    def _flatten(nodes: list) -> None:
        for node in nodes:
            # We copy to avoid mutating the original data and remove 'children' so we don't hold big structures
            flat_node = dict(node)
            children = flat_node.pop("children", [])
            flat.append(flat_node)
            if children:
                _flatten(children)

    _flatten(items)
    return flat


def main() -> None:
    """Entry point: parse args, load trace file, and run all analysis sections."""
    parser = argparse.ArgumentParser(description="Analyze a saved Datadog trace for slowness")
    parser.add_argument("file", help="Path to raw.json or trace.json (saved via dd_trace.py)")
    parser.add_argument("--top", type=int, default=20, help="Number of slowest spans to show (default: 20)")
    parser.add_argument("--ows-only", action="store_true", help="Skip sections 1 and 4; show only OWS analysis")
    parser.add_argument("--n1-threshold", type=int, default=5, help="Min calls to flag as N+1 candidate (default: 5)")
    parser.add_argument("--spans-json", action="store_true",
        help="Output top spans as JSON only (for interactive drill-down)")
    parser.add_argument("--leaf-only", action="store_true",
        help="When used with --spans-json, restrict to leaf spans (no children) to surface actual bottlenecks")
    parser.add_argument("--output-dir", dest="output_dir",
        help="Directory to save trace_analysis.txt alongside the trace data")
    args = parser.parse_args()

    try:
        with open(args.file) as f:
            data = json.load(f)
    except FileNotFoundError:
        print(f"Error: file not found: {args.file}", file=sys.stderr)
        sys.exit(1)

    spans = flatten_spans(data)
    if not spans:
        print("No spans found in file.", file=sys.stderr)
        sys.exit(1)

    if args.spans_json:
        print(json.dumps(top_spans_as_json(spans, args.top, leaf_only=args.leaf_only), indent=2))
        return

    # Set up tee to file if output_dir is given
    analysis_file = None
    analysis_path = None
    if args.output_dir:
        os.makedirs(args.output_dir, exist_ok=True)
        analysis_path = os.path.join(args.output_dir, "trace_analysis.txt")
        analysis_file = open(analysis_path, "w")
        ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        header_line = f"Analysis generated: {ts}\n\n"
        analysis_file.write(header_line)
        sys.stdout = _Tee(sys.__stdout__, analysis_file)

    root = get_root(spans)
    trace_dur_ns = get_trace_duration_ns(spans)

    if not args.ows_only:
        print_header(root, trace_dur_ns)

    print_service_budget(spans, trace_dur_ns)
    print_ows_endpoints(spans)

    n1_candidates = detect_n1(spans, threshold=args.n1_threshold)
    print_n1_report(n1_candidates, threshold=args.n1_threshold)

    if not args.ows_only:
        print_top_spans(spans, args.top, ows_only=False)

    print_top_spans(spans, args.top, ows_only=True)

    if analysis_file:
        sys.stdout = sys.__stdout__
        analysis_file.close()
        print(f"\nAnalysis saved to: {analysis_path}")


if __name__ == "__main__":
    main()