#!/usr/bin/env python3
"""Fetch and summarize a Datadog APM trace by ID.

Uses the Datadog spans API directly with cursor pagination to fetch all spans
(not capped at pup's 1000-span limit). Requires DD_API_KEY, DD_APP_KEY, and
DD_SITE env vars (same ones pup uses).

Usage:
    python dd_trace.py <trace_id> [--from <time>]

Examples:
    python dd_trace.py 7892843842668647916
    python dd_trace.py 7892843842668647916 --from 30d
"""

import argparse
import json
import logging
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone


PAGE_LIMIT = 1000

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


def ms(ns):
    return f"{round(ns / 1e6, 2)}ms"


def parse_tags(span_tags, keys):
    result = {}
    for t in span_tags:
        if ":" in t:
            k, v = t.split(":", 1)
            if k in keys:
                result[k] = v
    return result


def resolve_from_time(from_time):
    """Convert relative time string to ISO8601 timestamp."""
    now = datetime.now(timezone.utc)
    units = {"s": 1, "m": 60, "h": 3600, "d": 86400, "w": 604800}
    for suffix, secs in units.items():
        if from_time.endswith(suffix):
            n = int(from_time[:-1])
            return (now - timedelta(seconds=n * secs)).strftime("%Y-%m-%dT%H:%M:%SZ")
    # Assume already absolute (RFC3339 or similar)
    return from_time


def fetch_trace(trace_id, from_time, output_dir=None):
    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)

    url = f"https://api.{site}/api/v2/spans/events/search"
    headers = {
        "DD-API-KEY": api_key,
        "DD-APPLICATION-KEY": app_key,
        "Content-Type": "application/json",
    }

    start = resolve_from_time(from_time)
    end = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    all_spans = []
    cursor = None
    page_num = 0

    logging.info(f"Fetching trace {trace_id} from {start} to {end}")

    pages_dir = os.path.join(output_dir, "span_data") if output_dir else str(trace_id)

    while True:
        page_num += 1
        page = {"limit": PAGE_LIMIT}
        if cursor:
            page["cursor"] = cursor

        body = {
            "data": {
                "attributes": {
                    "filter": {
                        "query": f"trace_id:{trace_id}",
                        "from": start,
                        "to": end,
                    },
                    "page": page,
                    "sort": "timestamp",
                },
                "type": "search_request",
            }
        }

        logging.debug(f"POST {url} (page {page_num})")
        logging.debug(f"Request body: {json.dumps(body)}")

        req = urllib.request.Request(
            url,
            data=json.dumps(body).encode(),
            headers=headers,
            method="POST",
        )
        try:
            with urllib.request.urlopen(req) as resp:
                raw = resp.read()
                logging.debug(f"Response ({resp.status}): {raw[:500]}")
                data = json.loads(raw)
        except urllib.error.HTTPError as e:
            body_text = e.read().decode()
            logging.error(f"API error {e.code}: {body_text}")
            sys.exit(1)

        if not data:
            logging.warning("Empty response from API")
            break

        os.makedirs(pages_dir, exist_ok=True)
        page_file = os.path.join(pages_dir, f"page_{page_num}.json")
        with open(page_file, "w") as f:
            json.dump(data, f)
        logging.debug(f"Saved raw page {page_num} to {page_file}")

        page_spans = data.get("data", [])
        logging.info(f"Page {page_num}: got {len(page_spans)} spans (total so far: {len(all_spans) + len(page_spans)})")
        all_spans.extend(page_spans)

        cursor = ((data.get("meta") or {}).get("page") or {}).get("after")
        logging.debug(f"Next cursor: {cursor}")
        if not cursor or len(page_spans) < PAGE_LIMIT:
            break

    return {"data": all_spans}


def span_to_dict(s):
    a = s.get("attributes", s)
    c = a.get("custom", {})
    is_raw = "attributes" in s
    return {
        "span_id": a.get("span_id"),
        "parent_id": a.get("parent_id"),
        "service": a.get("service"),
        "operation": a.get("operation_name") if is_raw else a.get("operation"),
        "resource": a.get("resource_name") if is_raw else a.get("resource"),
        "status": a.get("status"),
        "error": bool(a.get("error")),
        "duration_ms": round((c.get("duration") or 0) / 1e6, 3) if is_raw else round((a.get("duration_ms") or 0), 3),
        "start": a.get("start_timestamp") if is_raw else a.get("start"),
        "end": a.get("end_timestamp") if is_raw else a.get("end"),
        "env": a.get("env"),
        "children": [],
    }


def build_tree(spans):
    """Build a nested span tree rooted at a parent span."""
    if not spans:
        return {}

    nodes = {}
    for s in spans:
        node = span_to_dict(s)
        sid = node.get("span_id")
        if sid:
            nodes[sid] = node
    if not nodes:
        return {}
    span_ids = set(nodes)

    # Prefer a natural root (no/missing parent), otherwise earliest span.
    root_candidates = []
    for node in nodes.values():
        pid = node.get("parent_id")
        if not pid or pid == "0" or pid not in span_ids:
            root_candidates.append(node)
    sortable_key = lambda node: (
        node.get("start") or "",
        node.get("span_id") or "",
    )
    root_node = min(root_candidates, key=sortable_key) if root_candidates else min(nodes.values(), key=sortable_key)
    root_id = root_node["span_id"]

    # Build a single-parent mapping, falling back to root when parent is invalid.
    parent_by_child = {}
    for sid, node in nodes.items():
        if sid == root_id:
            continue
        pid = node.get("parent_id")
        if not pid or pid == "0" or pid not in nodes or pid == sid:
            parent_by_child[sid] = root_id
        else:
            parent_by_child[sid] = pid

    # Break parent cycles in O(n) by walking parent chains once.
    state = {sid: 0 for sid in parent_by_child}  # 0=unseen, 1=visiting, 2=done
    for sid in parent_by_child:
        if state[sid] != 0:
            continue
        path = []
        index_by_sid = {}
        cur = sid
        while cur in parent_by_child:
            cur_state = state[cur]
            if cur_state == 2:
                break
            if cur_state == 1:
                cycle_start = index_by_sid.get(cur)
                if cycle_start is not None:
                    # Reattach one node from the cycle to root to break it.
                    parent_by_child[cur] = root_id
                break
            state[cur] = 1
            index_by_sid[cur] = len(path)
            path.append(cur)
            nxt = parent_by_child.get(cur)
            if not nxt or nxt == root_id or nxt not in parent_by_child:
                break
            cur = nxt
        for path_sid in path:
            state[path_sid] = 2

    for sid, pid in parent_by_child.items():
        nodes[pid]["children"].append(nodes[sid])

    # Sort children by start time (iterative to avoid recursion limit on deep traces)
    root_node = nodes[root_id]
    visited = set()
    stack = [root_node]
    while stack:
        node = stack.pop()
        if node["span_id"] in visited:
            continue
        visited.add(node["span_id"])
        node["children"].sort(key=lambda x: (x.get("start") or "", x.get("span_id") or ""))
        stack.extend(node["children"])
    return root_node


def print_json(data):
    spans = data if isinstance(data, list) else data.get("data", [])
    if not spans:
        print("[]")
        return

    tree = build_tree(spans)
    print(json.dumps(tree, indent=2))


def get_root(spans):
    span_ids = {s["attributes"]["span_id"] for s in spans}
    return next(
        (s for s in spans
         if not s["attributes"].get("parent_id")
         or s["attributes"]["parent_id"] not in span_ids),
        spans[0]
    )


def print_trace(data):
    spans = data.get("data", [])
    if not spans:
        print("No spans found for this trace ID and time range.")
        return

    root = get_root(spans)

    a = root["attributes"]
    c = a.get("custom", {})
    dur_ns = c.get("duration", 0)

    print("=" * 60)
    print(f"Trace: {a.get('trace_id')}")
    print("=" * 60)
    print(f"  Service:   {a.get('service')}")
    print(f"  Resource:  {a.get('resource_name')}")
    print(f"  Operation: {a.get('operation_name')}")
    print(f"  Status:    {a.get('status')}{'  [ERROR]' if a.get('error') else ''}")
    print(f"  Duration:  {ms(dur_ns)}")
    print(f"  Start:     {a.get('start_timestamp')}")
    print(f"  End:       {a.get('end_timestamp')}")
    print(f"  Env:       {a.get('env')}")

    http_tags = parse_tags(c.get("span_tags", []), {"http.method", "http.url", "http.status_code"})
    for k, v in http_tags.items():
        print(f"  {k}: {v}")

    errors = [s for s in spans if s["attributes"].get("error")]
    if errors:
        print(f"\n{'='*60}")
        print(f"ERRORS ({len(errors)})")
        print("=" * 60)
        for s in errors:
            ea = s["attributes"]
            ec = ea.get("custom", {})
            print(f"  [{ea.get('service')}] {ea.get('operation_name')} | {ea.get('resource_name')}")
            for t in ec.get("span_tags", []):
                if "error" in t.lower():
                    print(f"    {t}")

    print(f"\n{'='*60}")
    print(f"ALL SPANS ({len(spans)})")
    print("=" * 60)
    print(f"  {'SERVICE':<20} {'OPERATION':<50} {'RESOURCE':<45} {'DUR':>8}  STATUS")
    print(f"  {'-'*20} {'-'*50} {'-'*45} {'-'*8}  ------")
    for s in sorted(spans, key=lambda x: x["attributes"].get("start_timestamp") or ""):
        sa = s["attributes"]
        sc = sa.get("custom", {})
        dur = sc.get("duration", 0)
        err = " [ERROR]" if sa.get("error") else ""
        svc = (sa.get("service") or "")[:20]
        op  = (sa.get("operation_name") or "")[:50]
        res = (sa.get("resource_name") or "")[:45]
        print(f"  {svc:<20} {op:<50} {res:<45} {ms(dur):>8}  {sa.get('status','')}{err}")


def main():
    parser = argparse.ArgumentParser(description="Fetch a Datadog trace by ID")
    parser.add_argument("trace_id", help="Datadog trace ID")
    parser.add_argument("--from", dest="from_time", default="7d",
                        help="Time range to search (default: 7d). Examples: 1h, 4h, 7d, 30d")
    parser.add_argument("--json", action="store_true",
                        help="Output as nested JSON tree instead of table")
    parser.add_argument("-o", "--out", dest="output_file",
                        help="Save output to file (e.g. trace.json)")
    parser.add_argument("--raw", dest="raw_file",
                        help="Load spans from a raw JSON file instead of fetching (saved with --save-raw)")
    parser.add_argument("--save-raw", dest="save_raw",
                        help="Save raw fetched spans to file before processing")
    parser.add_argument("--output-dir", dest="output_dir",
                        help="Directory to save all output: span_data/page_N.json, full_<trace_id>.json")
    args = parser.parse_args()

    if args.raw_file:
        logging.info(f"Loading spans from {args.raw_file}")
        with open(args.raw_file) as f:
            data = json.load(f)
    else:
        data = fetch_trace(args.trace_id, args.from_time, output_dir=args.output_dir)
        if args.save_raw:
            with open(args.save_raw, "w") as f:
                json.dump(data, f)
            logging.info(f"Raw spans saved to {args.save_raw}")

    if args.json:
        spans = data if isinstance(data, list) else data.get("data", [])
        content = json.dumps(build_tree(spans), indent=2) if spans else "[]"
        saved = []
        if args.output_dir:
            os.makedirs(args.output_dir, exist_ok=True)
            full_path = os.path.join(args.output_dir, f"full_{args.trace_id}.json")
            with open(full_path, "w") as f:
                f.write(content)
            saved.append(full_path)
        if args.output_file:
            with open(args.output_file, "w") as f:
                f.write(content)
            saved.append(args.output_file)
        if saved:
            for p in saved:
                logging.info(f"Saved to {p}")
        else:
            print(content)
    else:
        spans = data if isinstance(data, list) else data.get("data", [])
        if spans and "attributes" not in spans[0]:
            logging.error("Table output requires raw spans with attributes. Use --json for normalized span input.")
            sys.exit(1)
        if isinstance(data, list):
            data = {"data": data}
        print_trace(data)


if __name__ == "__main__":
    main()