#!/usr/bin/env python3
"""Detect and activate every cost allocation tag on the current AWS account.

Self-contained: needs only Python 3 (boto3 is installed automatically if
missing). Uses whatever credentials the environment already provides, so
upload it to CloudShell and run:

    python3 enable_cost_tags.py            # list tags, then activate all inactive ones
    python3 enable_cost_tags.py --dry-run  # list only, change nothing

Notes:
  * Cost Explorer lists a tag key only after it has appeared on a resource that
    incurred cost, so brand-new tags can take up to ~24h to show up.
  * Activated tags take up to 24h to appear in billing/Cost Explorer data.
  * In an AWS Organization this must run in the management (payer) account.
  * Needs ce:ListCostAllocationTags and ce:UpdateCostAllocationTagsStatus.
"""
import argparse
import importlib
import subprocess
import sys

print("=== Cost allocation tag activator ===", flush=True)


def ensure_boto3():
    """Import boto3, installing it for the current user first if it is missing."""
    try:
        importlib.import_module("boto3")
        return
    except ImportError:
        pass
    print("boto3 is not installed - installing it now (this takes a few seconds)...", flush=True)
    cmd = [sys.executable, "-m", "pip", "install", "--quiet", "--user", "boto3"]
    if subprocess.call(cmd) != 0:
        # Some Pythons lack pip; try to bootstrap it once, then retry.
        subprocess.call([sys.executable, "-m", "ensurepip", "--user"])
        if subprocess.call(cmd) != 0:
            sys.exit("Could not install boto3 automatically. Please contact your administrator.")
    importlib.invalidate_caches()
    import site
    site.main()  # pick up the freshly created user site-packages directory
    importlib.import_module("boto3")
    print("boto3 installed.\n", flush=True)


ensure_boto3()
import boto3  # noqa: E402
from botocore.exceptions import ClientError  # noqa: E402

BATCH = 20  # UpdateCostAllocationTagsStatus accepts at most 20 entries per call


def list_tags(ce):
    tags, token = [], None
    while True:
        kw = {"MaxResults": 1000}
        if token:
            kw["NextToken"] = token
        r = ce.list_cost_allocation_tags(**kw)
        tags += r["CostAllocationTags"]
        token = r.get("NextToken")
        if not token:
            return tags


def main():
    ap = argparse.ArgumentParser(description="Activate all cost allocation tags.")
    ap.add_argument("--dry-run", action="store_true", help="show tags, change nothing")
    args = ap.parse_args()

    session = boto3.Session()
    print("Detecting account...", flush=True)
    try:
        ident = session.client("sts").get_caller_identity()
    except ClientError as e:
        sys.exit(f"Cannot determine account (no valid credentials?): {e}")
    print(f"Account: {ident['Account']}\nCaller:  {ident['Arn']}\n", flush=True)

    ce = session.client("ce", region_name="us-east-1")  # Cost Explorer is global, served from us-east-1
    print("Looking up cost allocation tags...", flush=True)
    try:
        tags = list_tags(ce)
    except ClientError as e:
        sys.exit(f"Could not list cost allocation tags: {e}\n"
                 "In an Organization, run this in the management account.")

    if not tags:
        print("No cost allocation tags detected yet. Tags appear here once they are on "
              "resources that have incurred cost (can take up to 24h).")
        return

    tags.sort(key=lambda t: (t["Type"], t["TagKey"].lower()))
    w = max(len(t["TagKey"]) for t in tags)
    print(f"\n{'TAG KEY':<{w}}  {'TYPE':<12}  STATUS")
    for t in tags:
        print(f"{t['TagKey']:<{w}}  {t['Type']:<12}  {t['Status']}")

    todo = [t["TagKey"] for t in tags if t["Status"] != "Active"]
    print(f"\n{len(tags)} tag(s) detected, {len(tags) - len(todo)} already active, "
          f"{len(todo)} to activate.", flush=True)
    if not todo:
        return
    if args.dry_run:
        print("Dry run: nothing changed.")
        return

    print("Activating...", flush=True)
    failed = 0
    for i in range(0, len(todo), BATCH):
        chunk = todo[i:i + BATCH]
        try:
            r = ce.update_cost_allocation_tags_status(
                CostAllocationTagsStatus=[{"TagKey": k, "Status": "Active"} for k in chunk])
        except ClientError as e:
            sys.exit(f"Activation failed: {e}")
        for err in r.get("Errors", []):
            failed += 1
            print(f"  FAILED {err['TagKey']}: {err.get('Code')} {err.get('Message')}")
    ok = len(todo) - failed
    print(f"Activated {ok} tag(s)." + (f" {failed} failed." if failed else ""))
    if ok:
        print("They may take up to 24h to appear in billing data.")
    print("=== Done ===")
    sys.exit(1 if failed else 0)


main()