#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = ["boto3>=1.34"]
# ///
"""Cost Explorer breakdown with baseline-vs-period deltas, as compact JSON.

Fetches daily UnblendedCost for [start - baseline_days, end] in a single
Cost Explorer request, then computes per-group deltas between the baseline
window and the analysis period. Output is capped to the top movers so it
stays readable.

Usage:
    uv run scripts/cost_breakdown.py --start 2026-06-01 --end 2026-06-10
    uv run scripts/cost_breakdown.py --start ... --end ... --group-by USAGE_TYPE --filter-service "Amazon Elastic Compute Cloud - Compute"
    uv run scripts/cost_breakdown.py --start ... --end ... --group-by TAG:team
"""
import argparse
import json
from collections import defaultdict
from datetime import datetime, timedelta

import boto3

GROUP_DIMENSIONS = {"SERVICE", "USAGE_TYPE", "REGION"}


def parse_date(s: str):
    return datetime.strptime(s, "%Y-%m-%d").date()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--start", required=True, type=parse_date, help="analysis period start (YYYY-MM-DD, inclusive)")
    parser.add_argument("--end", required=True, type=parse_date, help="analysis period end (YYYY-MM-DD, inclusive)")
    parser.add_argument("--group-by", default="SERVICE",
                        help="SERVICE | USAGE_TYPE | REGION | TAG:<tag-key> (default: SERVICE)")
    parser.add_argument("--filter-service", help="restrict to one service (exact Cost Explorer service name)")
    parser.add_argument("--baseline-days", type=int, default=14,
                        help="days immediately before --start used as baseline (default: 14)")
    parser.add_argument("--top", type=int, default=20, help="max groups returned, by |delta| (default: 20)")
    parser.add_argument("--profile", help="AWS profile name")
    args = parser.parse_args()

    if args.end < args.start:
        raise SystemExit("--end must be >= --start")

    session = boto3.Session(profile_name=args.profile)
    ce = session.client("ce", region_name="us-east-1")

    baseline_start = args.start - timedelta(days=args.baseline_days)
    fetch_end_exclusive = args.end + timedelta(days=1)  # CE TimePeriod End is exclusive

    group_by = args.group_by.upper()
    if group_by in GROUP_DIMENSIONS:
        group_spec = {"Type": "DIMENSION", "Key": group_by}
    elif group_by.startswith("TAG:"):
        group_spec = {"Type": "TAG", "Key": args.group_by.split(":", 1)[1]}
    else:
        raise SystemExit(f"unsupported --group-by {args.group_by!r}; use SERVICE, USAGE_TYPE, REGION, or TAG:<key>")

    kwargs = {
        "TimePeriod": {"Start": baseline_start.isoformat(), "End": fetch_end_exclusive.isoformat()},
        "Granularity": "DAILY",
        "Metrics": ["UnblendedCost"],
        "GroupBy": [group_spec],
    }
    if args.filter_service:
        kwargs["Filter"] = {"Dimensions": {"Key": "SERVICE", "Values": [args.filter_service]}}

    # group key -> {iso date -> cost}
    daily = defaultdict(dict)
    while True:
        resp = ce.get_cost_and_usage(**kwargs)
        for day in resp["ResultsByTime"]:
            d = day["TimePeriod"]["Start"]
            for g in day.get("Groups", []):
                key = " | ".join(g["Keys"])
                daily[key][d] = daily[key].get(d, 0.0) + float(g["Metrics"]["UnblendedCost"]["Amount"])
        token = resp.get("NextPageToken")
        if not token:
            break
        kwargs["NextPageToken"] = token

    period_days = (args.end - args.start).days + 1
    period_dates = {(args.start + timedelta(days=i)).isoformat() for i in range(period_days)}
    baseline_dates = {(baseline_start + timedelta(days=i)).isoformat() for i in range(args.baseline_days)}

    groups = []
    for key, series in daily.items():
        baseline_total = sum(v for d, v in series.items() if d in baseline_dates)
        period_total = sum(v for d, v in series.items() if d in period_dates)
        baseline_daily_avg = baseline_total / args.baseline_days if args.baseline_days else 0.0
        expected_period_total = baseline_daily_avg * period_days
        delta = period_total - expected_period_total
        groups.append(
            {
                "group": key,
                "baseline_daily_avg_usd": round(baseline_daily_avg, 2),
                "period_daily_avg_usd": round(period_total / period_days, 2),
                "period_total_usd": round(period_total, 2),
                "delta_vs_baseline_usd": round(delta, 2),
                "pct_change": round(100 * delta / expected_period_total, 1) if expected_period_total > 0.005 else None,
            }
        )

    groups.sort(key=lambda g: abs(g["delta_vs_baseline_usd"]), reverse=True)
    truncated = len(groups) > args.top
    groups = groups[: args.top]

    # daily grand total across all groups, so the spike day is visible
    totals = defaultdict(float)
    for series in daily.values():
        for d, v in series.items():
            totals[d] += v
    daily_totals = {d: round(v, 2) for d, v in sorted(totals.items())}

    print(json.dumps(
        {
            "group_by": args.group_by,
            "filter_service": args.filter_service,
            "baseline_window": {"start": baseline_start.isoformat(), "end": (args.start - timedelta(days=1)).isoformat()},
            "analysis_period": {"start": args.start.isoformat(), "end": args.end.isoformat()},
            "truncated_to_top": args.top if truncated else None,
            "groups": groups,
            "daily_totals_usd": daily_totals,
        },
        indent=1,
    ))


if __name__ == "__main__":
    main()