#!/usr/bin/env python3
"""Attribute AWS account cost to VPCs, EKS clusters and teams without cost allocation tags.

Why this exists: 0414 is a member account, so activating cost allocation tags needs the
payer account, and resource-level billing is off (get-cost-and-usage-with-resources returns
only NoResourceId). So we join the two sides ourselves:

  inventory:  resource -> VPC -> EKS cluster -> team    (read-only APIs, no tags needed)
  billing:    Cost Explorer grouped by USAGE_TYPE        (no resource ids needed)

SAVINGS PLANS: the default metric is amortized, not unblended, because this account does not
buy its own Savings Plans -- the payer account does. Unblended therefore reports SP-covered
usage as $0 here (SavingsPlanCoveredUsage is cancelled by an equal SavingsPlanNegation), so
an unblended audit silently omits it, and omits it unevenly: the whole gap lands on EC2
compute, which is also what drives per-cluster attribution. Amortized attributes the plan's
actual cost back to the usage that consumed it, which is the company-wide view. The two
answer different questions -- turning off an SP-covered instance saves the company money but
does not lower this account's own unblended bill. tools/README.md carries the measured
impact and the guidance on which metric to use when.

Uniform-rate charges are split exactly, not estimated: NAT-gateway hours divided by NAT
count is the per-gateway cost, EBS scales with each volume's size, EKS control-plane hours
are a flat per-cluster charge. Bytes-based charges (NAT/data transfer) depend on traffic
rather than resource count, and S3/ECR/KMS are not VPC-scoped, so those stay in an explicit
SHARED bucket. Spreading them would produce a precise-looking, misleading table.

HISTORICAL PERIODS: inventory is a snapshot of now, so every resource is filtered by its
creation time -- anything created after the period is excluded. What cannot be recovered is
resources that existed then and have since been deleted. The script quantifies that gap per
period by comparing billed NAT-gateway hours against the hours today's inventory could have
produced, and reports it as inventory drift. Treat a period with large drift as indicative
rather than exact.

Usage:
  python3 cost_by_cluster.py --profile PingCAP --start 2026-07-01 --end 2026-07-25
  python3 cost_by_cluster.py --profile PingCAP --months 3
"""

import argparse
import collections
import datetime as dt
import json
import subprocess
import sys

REGIONS = ["us-west-2", "us-east-1"]
REGION_PREFIX = {"us-west-2": "USW2-", "us-east-1": ""}

# Cluster -> team. Confirmed by SabaPing 2026-07-25; unmapped clusters show as UNCONFIRMED.
TEAM_MAP = {
    "staging-us-east-1-f01": "Starter",
    "staging-seed-us-east-1": "Starter",
    "staging-seed-us-east-1-ess": "Starter",
    "staging-seed-us-west-2": "Dedicated",
    "staging-seed-us-west-2-ng-0711": "Premium",
    "staging-seed-usw2-byoc-pm": "Premium",
    "dev-seed-us-west-2-byoc-0303": "Premium",
    "clawmem-staging-us-west-2-eks": "AI Innovation",
    "staging-seed-us-west-2-lyc": "QA",
    "staging-seed-us-west-2-mjp": "QA",
    "staging-base-us-west-2": "SHARED-global-cp",
    "dev-base-us-west-2": "SHARED-dev-cp",
}


def aws(profile, region, *args):
    cmd = ["aws", "--profile", profile, "--output", "json"]
    if region:
        cmd += ["--region", region]
    cmd += list(args)
    out = subprocess.run(cmd, capture_output=True, text=True, timeout=180)
    if out.returncode != 0:
        print(f"  ! {' '.join(args[:2])} ({region}): {out.stderr.strip()[:110]}", file=sys.stderr)
        return None
    return json.loads(out.stdout or "null")


def ts(value):
    """Parse an AWS timestamp to a date, tolerating the several shapes AWS returns."""
    if not value:
        return dt.date(2000, 1, 1)
    if isinstance(value, (int, float)):
        return dt.datetime.utcfromtimestamp(value).date()
    return dt.date.fromisoformat(str(value)[:10])


def collect(profile):
    """One record per billable resource: kind, vpc, cluster, created, and a weight."""
    rows = []
    vpc_cluster = {}
    for region in REGIONS:
        clusters = (aws(profile, region, "eks", "list-clusters") or {}).get("clusters", [])
        for name in clusters:
            d = aws(profile, region, "eks", "describe-cluster", "--name", name)
            if d:
                c = d["cluster"]
                vpc_cluster[c["resourcesVpcConfig"]["vpcId"]] = name
                rows.append(dict(kind="eks", region=region, vpc=c["resourcesVpcConfig"]["vpcId"],
                                 created=ts(c.get("createdAt")), weight=1, sub=""))

        res = aws(profile, region, "ec2", "describe-instances",
                  "--filters", "Name=instance-state-name,Values=running") or {}
        inst_vpc = {}
        for r in res.get("Reservations", []):
            for i in r["Instances"]:
                inst_vpc[i["InstanceId"]] = i.get("VpcId")
                rows.append(dict(kind="compute", region=region, vpc=i.get("VpcId"),
                                 created=ts(i.get("LaunchTime")), weight=1, sub=i["InstanceType"]))

        for v in (aws(profile, region, "ec2", "describe-volumes") or {}).get("Volumes", []):
            att = v.get("Attachments") or []
            rows.append(dict(kind="ebs", region=region,
                             vpc=inst_vpc.get(att[0]["InstanceId"]) if att else None,
                             created=ts(v.get("CreateTime")), weight=v["Size"], sub=v["VolumeType"]))

        for n in (aws(profile, region, "ec2", "describe-nat-gateways",
                      "--filter", "Name=state,Values=available") or {}).get("NatGateways", []):
            rows.append(dict(kind="nat", region=region, vpc=n.get("VpcId"),
                             created=ts(n.get("CreateTime")), weight=1, sub=""))

        for lb in (aws(profile, region, "elbv2", "describe-load-balancers") or {}).get("LoadBalancers", []):
            rows.append(dict(kind="loadbalancer", region=region, vpc=lb.get("VpcId"),
                             created=ts(lb.get("CreatedTime")), weight=1, sub=""))
        for lb in (aws(profile, region, "elb", "describe-load-balancers") or {}).get("LoadBalancerDescriptions", []):
            rows.append(dict(kind="loadbalancer", region=region, vpc=lb.get("VPCId"),
                             created=ts(lb.get("CreatedTime")), weight=1, sub=""))

        subnet_vpc = {s["SubnetId"]: s["VpcId"]
                      for s in (aws(profile, region, "ec2", "describe-subnets") or {}).get("Subnets", [])}

        for db in (aws(profile, region, "rds", "describe-db-instances") or {}).get("DBInstances", []):
            sn = (db.get("DBSubnetGroup") or {}).get("Subnets") or []
            rows.append(dict(kind="rds", region=region,
                             vpc=subnet_vpc.get(sn[0]["SubnetIdentifier"]) if sn else None,
                             created=ts(db.get("InstanceCreateTime")), weight=1,
                             sub=db["DBInstanceClass"]))

        for e in (aws(profile, region, "ec2", "describe-vpc-endpoints") or {}).get("VpcEndpoints", []):
            rows.append(dict(kind="vpc-endpoint", region=region, vpc=e.get("VpcId"),
                             created=ts(e.get("CreationTimestamp")), weight=1, sub=""))

        for e in (aws(profile, region, "ec2", "describe-network-interfaces") or {}).get("NetworkInterfaces", []):
            if e.get("Association", {}).get("PublicIp"):
                rows.append(dict(kind="public-ipv4", region=region, vpc=e.get("VpcId"),
                                 created=dt.date(2000, 1, 1), weight=1, sub=""))

        for m in (aws(profile, region, "kafka", "list-clusters-v2") or {}).get("ClusterInfoList", []):
            prov = m.get("Provisioned") or {}
            subs = ((prov.get("BrokerNodeGroupInfo") or {}).get("ClientSubnets")) or []
            rows.append(dict(kind="msk", region=region,
                             vpc=subnet_vpc.get(subs[0]) if subs else None,
                             created=ts(m.get("CreationTime")), weight=prov.get("NumberOfBrokerNodes", 1),
                             sub=""))

    for r in rows:
        r["cluster"] = vpc_cluster.get(r["vpc"], "")
        r["vpc"] = r["vpc"] or f"(no-vpc-{r['region']})"
    return rows, vpc_cluster


# The unblended metric each amortized metric is compared against. Pairing net with net and
# gross with gross keeps the difference between them purely the effect of amortizing a
# Savings Plan or reserved instance; comparing across the pair would fold private-rate
# discounts and credits into a figure labelled as plan cost.
UNBLENDED_SIBLING = {"AmortizedCost": "UnblendedCost", "NetAmortizedCost": "NetUnblendedCost"}


def costs_for(profile, start, end, metric):
    """Cost and usage quantity per usage type, plus the matching unblended cost.

    The unblended sibling is fetched only for an amortized metric, and is what the
    savings-plan line is measured against. Returns (cost, qty, unblended-or-None).
    """
    sibling = UNBLENDED_SIBLING.get(metric)
    metrics = [metric, "UsageQuantity"] + ([sibling] if sibling else [])
    d = aws(profile, None, "ce", "get-cost-and-usage",
            "--time-period", f"Start={start},End={end}", "--granularity", "MONTHLY",
            "--metrics", *metrics, "--group-by", "Type=DIMENSION,Key=USAGE_TYPE")
    if not d:
        sys.exit("Cost Explorer query failed -- the profile needs ce:GetCostAndUsage.")
    cost, qty = collections.Counter(), collections.Counter()
    unblended = collections.Counter() if sibling else None
    for p in d["ResultsByTime"]:
        for g in p["Groups"]:
            usage = g["Keys"][0]
            cost[usage] += float(g["Metrics"][metric]["Amount"])
            qty[usage] += float(g["Metrics"]["UsageQuantity"]["Amount"])
            if sibling:
                unblended[usage] += float(g["Metrics"][sibling]["Amount"])
    return cost, qty, unblended


def bucket_for(usage, region):
    """Map a usage type to the inventory kind that carries it, or None if not VPC-scoped."""
    p = REGION_PREFIX[region]
    if p and not usage.startswith(p):
        return None
    b = usage[len(p):] if p else usage
    if b.startswith(("BoxUsage:", "SpotUsage:")):
        return ("compute", b.split(":", 1)[1])
    if b.startswith("NatGateway-Hours"):
        return ("nat", None)
    if "LoadBalancerUsage" in b or "LCUUsage" in b:
        return ("loadbalancer", None)
    if b.startswith("EBS:Volume"):
        return ("ebs", b.split(".")[-1] if "." in b else None)
    if b.startswith(("InstanceUsage:", "Multi-AZUsage:", "Aurora:")):
        return ("rds", None)
    if b.startswith("AmazonEKS-Hours"):
        return ("eks", None)
    if b.startswith("VpcEndpoint-Hours"):
        return ("vpc-endpoint", None)
    if b.startswith("PublicIPv4:"):
        return ("public-ipv4", None)
    if b.startswith("Kafka."):
        return ("msk", None)
    return None


def shared_label(usage):
    u = usage.lower()
    if "natgateway-bytes" in u or "datatransfer" in u or "-bytes" in u:
        return "data-transfer (needs VPC Flow Logs)"
    if "timedstorage" in u or "requests-tier" in u:
        return "S3 (not VPC-scoped)"
    if "ecr" in u or "repositorystorage" in u:
        return "ECR (shared registry)"
    if "kms" in u:
        return "KMS keys"
    if "certificate" in u or "acm" in u or "privateca" in u:
        return "ACM / Private CA"
    if "spotusage" in u or "boxusage" in u:
        return "deleted instances (not in inventory)"
    return "other-unattributed"


def attribute(rows, cost, end_date):
    """Split each usage type over the resources that already existed in the period."""
    live = [r for r in rows if r["created"] <= end_date]
    per = collections.defaultdict(collections.Counter)   # vpc -> category -> $
    shared = collections.Counter()
    for usage, amount in cost.items():
        if amount <= 0:
            continue
        hit = None
        for region in REGIONS:
            b = bucket_for(usage, region)
            if not b:
                continue
            kind, sub = b
            sel = [r for r in live if r["kind"] == kind and r["region"] == region
                   and (sub is None or r["sub"] == sub)]
            if not sel and sub is not None:      # unknown sub-type: fall back to the kind
                sel = [r for r in live if r["kind"] == kind and r["region"] == region]
            if sel:
                denom = sum(r["weight"] for r in sel)
                for r in sel:
                    per[r["vpc"]][kind] += amount * r["weight"] / denom
                hit = True
                break
        if not hit:
            shared[shared_label(usage)] += amount
    return per, shared, live


def nat_drift(rows, qty, end_date, days):
    """Billed NAT hours vs what today's inventory could have produced -> deleted resources."""
    billed = sum(v for k, v in qty.items() if "NatGateway-Hours" in k)
    live = len([r for r in rows if r["kind"] == "nat" and r["created"] <= end_date])
    possible = live * days * 24
    if possible <= 0:
        return None
    return billed / possible


def report(profile, start, end, metric, rows, vpc_cluster):
    cost, qty, unblended = costs_for(profile, start, end, metric)
    end_date = dt.date.fromisoformat(end)
    days = (end_date - dt.date.fromisoformat(start)).days or 30
    per, shared, live = attribute(rows, cost, end_date)

    total = sum(cost.values())
    attributed = sum(sum(c.values()) for c in per.values())
    drift = nat_drift(rows, qty, end_date, days)

    print(f"\n{'='*104}\n{metric}  {start} -> {end}  ({days}d)   TOTAL ${total:,.0f}")
    print(f"attributed ${attributed:,.0f} ({attributed/total*100:.0f}%)   "
          f"shared/unattributable ${sum(shared.values()):,.0f} ({sum(shared.values())/total*100:.0f}%)")
    covered = total - sum(unblended.values()) if unblended is not None else None
    if covered is not None:
        # Printed as "of which", not as a third bucket: this spend already sits inside the
        # attributed and shared figures above, mostly in compute. Adding all three would
        # exceed the total.
        print(f"  of which savings plan ${covered:,.0f} ({covered/total*100:.0f}%) -- already "
              f"counted above; real company spend that this account's own bill shows as $0, "
              f"because the plan is paid for by the payer account")
    if drift:
        note = "inventory covers the period" if drift <= 1.15 else \
               f"~{(drift-1)*100:.0f}% MORE NAT hours were billed than today's inventory explains"
        print(f"inventory drift check (NAT hours billed / possible): {drift:.2f}x -- {note}")

    print(f"\n{'VPC':<26}{'EKS cluster':<44}{'team':<18}{'$':>10}")
    by_cluster, by_team = collections.Counter(), collections.Counter()
    for vpc, cats in sorted(per.items(), key=lambda kv: -sum(kv[1].values())):
        amt = sum(cats.values())
        cluster = vpc_cluster.get(vpc, "")
        team = TEAM_MAP.get(cluster, "UNCONFIRMED" if cluster else "non-cluster VPC")
        by_cluster[cluster or "(no EKS cluster)"] += amt
        by_team[team] += amt
        print(f"{vpc:<26}{(cluster or '—')[:42]:<44}{team:<18}{amt:>10,.0f}")

    print(f"\n{'BY EKS CLUSTER':<44}{'$':>10}")
    for c, amt in by_cluster.most_common():
        print(f"  {c[:56]:<56}{amt:>10,.0f}")
    print(f"\n{'BY TEAM':<44}{'$':>10}")
    for t, amt in by_team.most_common():
        print(f"  {t:<56}{amt:>10,.0f}   ({amt/total*100:>4.1f}% of account)")
    print(f"\n{'SHARED / NOT ATTRIBUTABLE':<44}{'$':>10}")
    for k, amt in shared.most_common():
        print(f"  {k:<56}{amt:>10,.0f}")

    return dict(start=start, end=end, metric=metric, total=total, savings_plan_share=covered,
                attributed=attributed, drift=drift,
                by_vpc={k: dict(v) for k, v in per.items()},
                by_cluster=dict(by_cluster), by_team=dict(by_team), shared=dict(shared))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--profile", required=True)
    ap.add_argument("--start")
    ap.add_argument("--end")
    ap.add_argument("--months", type=int, help="audit the last N whole months plus month-to-date")
    ap.add_argument("--today", default=dt.date.today().isoformat())
    ap.add_argument("--metric", default="NetAmortizedCost",
                    choices=["UnblendedCost", "NetUnblendedCost",
                             "AmortizedCost", "NetAmortizedCost"])
    ap.add_argument("--json")
    args = ap.parse_args()

    print("collecting inventory ...", file=sys.stderr)
    rows, vpc_cluster = collect(args.profile)
    print(f"  {len(rows)} resources, {len(vpc_cluster)} EKS clusters", file=sys.stderr)

    periods = []
    if args.months:
        today = dt.date.fromisoformat(args.today)
        first = today.replace(day=1)
        for _ in range(args.months - 1):
            prev_end = first
            first = (first - dt.timedelta(days=1)).replace(day=1)
            periods.append((first.isoformat(), prev_end.isoformat()))
        periods.reverse()
        periods.append((today.replace(day=1).isoformat(), today.isoformat()))
    else:
        periods.append((args.start, args.end))

    out = [report(args.profile, s, e, args.metric, rows, vpc_cluster) for s, e in periods]
    if args.json:
        with open(args.json, "w") as fh:
            json.dump(out, fh, indent=2, default=str)
        print(f"\nwrote {args.json}", file=sys.stderr)


if __name__ == "__main__":
    main()