#!/usr/bin/env python3
"""Fetch logs from Datadog using the Logs Search API.

Requires DD_API_KEY and DD_APP_KEY env vars. Optionally, set DD_SITE to override the Datadog site (defaults to datadoghq.com).

Usage:
    python dd_logs.py <query> [--range 15m] [--start 2024-01-01T00:00:00Z] [--end 2024-01-01T01:00:00Z]

Examples:
    python dd_logs.py "service:my-service status:error"
    python dd_logs.py "host:web-01" --range 1h
    python dd_logs.py "env:prod" --start 2024-01-01T00:00:00Z --end 2024-01-01T06:00:00Z
    python dd_logs.py "*" --range 30m --json
"""

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


PAGE_LIMIT = 1000

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


def parse_range(range_str):
    """Parse a relative range string like 15m, 2h, 3d into a timedelta."""
    units = {"s": 1, "m": 60, "h": 3600, "d": 86400, "w": 604800}
    suffix = range_str[-1]
    if suffix not in units:
        raise argparse.ArgumentTypeError(
            f"Invalid range '{range_str}'. Use a number followed by s/m/h/d/w (e.g. 15m, 2h, 7d)."
        )
    try:
        n = int(range_str[:-1])
    except ValueError:
        raise argparse.ArgumentTypeError(
            f"Invalid range '{range_str}'. Must be a number followed by s/m/h/d/w."
        )
    return timedelta(seconds=n * units[suffix])


def parse_datetime(dt_str):
    """Parse an ISO8601-ish datetime string into a UTC-aware datetime."""
    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",
    ]
    for fmt in formats:
        try:
            dt = datetime.strptime(dt_str, fmt)
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt
        except ValueError:
            continue
    raise argparse.ArgumentTypeError(
        f"Cannot parse datetime '{dt_str}'. Use ISO8601 format e.g. 2024-01-01T00:00:00Z"
    )


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


def format_log(log):
    a = log.get("attributes", {})
    ts = a.get("timestamp", "")
    status = (a.get("status") or "").upper()
    service = a.get("service") or ""
    host = a.get("host") or ""
    message = a.get("message") or ""

    parts = [ts]
    if status:
        parts.append(f"[{status}]")
    if service:
        parts.append(f"service:{service}")
    if host:
        parts.append(f"host:{host}")
    parts.append(message)

    return "  ".join(parts)


def stream_logs(query, start, end, on_page):
    """Fetch all log pages and call on_page(logs) for each. Never holds all logs in memory."""
    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/logs/events/search"
    headers = {
        "DD-API-KEY": api_key,
        "DD-APPLICATION-KEY": app_key,
        "Content-Type": "application/json",
    }

    start_str = to_iso(start)
    end_str = to_iso(end)

    logging.info(f"Searching logs from {start_str} to {end_str}")
    logging.info(f"Query: {query}")

    cursor = None
    page_num = 0
    total = 0

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

        body = {
            "filter": {
                "query": query,
                "from": start_str,
                "to": end_str,
            },
            "sort": "timestamp",
            "page": page,
        }

        logging.debug(f"POST {url} (page {page_num})")

        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)

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

        on_page(page_logs)

        cursor = ((data.get("meta") or {}).get("page") or {}).get("after")
        if not cursor or len(page_logs) < PAGE_LIMIT:
            break

    return total


def main():
    parser = argparse.ArgumentParser(description="Search Datadog logs")
    parser.add_argument("query", help="Log search query (e.g. 'service:my-app status:error')")

    time_group = parser.add_argument_group("Time range (use --range OR --start/--end)")
    time_group.add_argument(
        "--range", "-r",
        dest="range",
        metavar="RANGE",
        help="Relative time range ending now (e.g. 15m, 2h, 7d). Mutually exclusive with --start/--end.",
    )
    time_group.add_argument(
        "--start",
        dest="start",
        metavar="DATETIME",
        help="Start datetime in ISO8601 format (e.g. 2024-01-01T00:00:00Z)",
    )
    time_group.add_argument(
        "--end",
        dest="end",
        metavar="DATETIME",
        help="End datetime in ISO8601 format. Defaults to now when --start is given.",
    )

    parser.add_argument("--json", action="store_true", help="Output raw JSON instead of formatted text")
    parser.add_argument("--output", "-o", default="./logs.txt", metavar="FILE",
                        help="Output file path (default: ./logs.txt)")
    args = parser.parse_args()

    # Validate time args
    if args.range and (args.start or args.end):
        parser.error("--range cannot be used together with --start or --end")

    now = datetime.now(timezone.utc)

    if args.range:
        delta = parse_range(args.range)
        start = now - delta
        end = now
    elif args.start:
        start = parse_datetime(args.start)
        end = parse_datetime(args.end) if args.end else now
    elif args.end:
        parser.error("--end requires --start")
    else:
        # Default: last 15 minutes
        start = now - timedelta(minutes=15)
        end = now

    with open(args.output, "w") as f:
        if args.json:
            # Write a valid JSON array incrementally: one object per line inside [ ]
            f.write("[\n")
            first_log = [True]  # mutable for closure

            def write_json_page(logs):
                for log in logs:
                    if not first_log[0]:
                        f.write(",\n")
                    f.write(json.dumps(log))
                    first_log[0] = False
                f.flush()

            total = stream_logs(args.query, start, end, write_json_page)
            f.write("\n]\n")
        else:
            empty = [True]

            def write_text_page(logs):
                for log in logs:
                    f.write(format_log(log) + "\n")
                    empty[0] = False
                f.flush()

            total = stream_logs(args.query, start, end, write_text_page)
            if empty[0]:
                f.write("No logs found.\n")

    logging.info(f"Wrote {total} logs to {args.output}")


if __name__ == "__main__":
    main()