#!/usr/bin/env python3
"""
Run 'Update RADIUS Account Groups' for every account in Gaiia.
Status must be UPPERCASE (ACTIVE, SUSPENDED, etc.)

Usage:
  source .envrc
  uv run python3 sync_all_radius_groups.py --dry-run    (preview)
  uv run python3 sync_all_radius_groups.py               (execute)
"""

from gaiia import GaiiaClient, GaiiaAPIError, GaiiaRateLimitedError
import json
import sys
import time

WF_ID = "workflow_fLXTNCvridPqwEEXxf3CK1"


def list_accounts(client):
    """Paginate through all accounts, return [{id, name, status, productIds}]."""
    accounts = []
    after = None
    page = 0

    while True:
        page += 1
        data = client.query("""
            query Q($after: String) {
                accounts(first: 250, after: $after) {
                    nodes {
                        id
                        name
                        status { name }
                        billingSubscriptions {
                            nodes {
                                productVersion {
                                    product { id name slug }
                                }
                            }
                        }
                    }
                    pageInfo { hasNextPage endCursor }
                    totalCount
                }
            }
        """, {"after": after} if after else {})

        result = data['accounts']

        for acct in result['nodes']:
            status_name = acct['status']['name'] if acct['status'] else ''
            products = []
            for sub in (acct.get('billingSubscriptions') or {}).get('nodes', []):
                pv = sub.get('productVersion') or {}
                p = pv.get('product') if pv else None
                if p and p['id'] not in products:
                    products.append(p['id'])

            accounts.append({
                'id': acct['id'],
                'name': acct.get('name', '') or '',
                'status': status_name,
                'productIds': products,
            })

        sys.stderr.write(f"\r  Page {page}: {len(accounts)} accounts so far...")
        sys.stderr.flush()

        page_info = result['pageInfo']
        if not page_info.get('hasNextPage'):
            break
        after = page_info['endCursor']

    sys.stderr.write("\n")
    return accounts


def sync_account(client, acct, dry_run=False):
    """Run the Update RADIUS Account Groups workflow for one account."""
    status_upper = acct['status'].upper()
    wf_input = {
        "accountId": acct['id'],
        "status": status_upper,
        "productIds": acct['productIds'],
    }

    if dry_run:
        name = acct['name'] or '(no name)'
        print(f"  [DRY-RUN] {name}  status={status_upper}  products={len(acct['productIds'])}")
        return True

    try:
        result = client.query("""
            mutation Run($input: StartWorkflowExecutionInput!) {
                startWorkflowExecution(input: $input) {
                    workflowExecution { id status }
                }
            }
        """, {"input": {"workflowId": WF_ID, "input": wf_input}})

        we = result['startWorkflowExecution']['workflowExecution']
        return we is not None
    except GaiiaAPIError as e:
        print(f"  FAILED {acct.get('name','?')}: API error - {e}")
        return False
    except GaiiaRateLimitedError as e:
        print(f"  RATE LIMITED - sleeping until {e.retry_at}")
        if e.retry_at:
            wait = max(0, (e.retry_at.timestamp() - time.time()))
            time.sleep(wait + 1)
        return False


def main():
    dry_run = "--dry-run" in sys.argv

    with GaiiaClient(timezone='America/Chicago') as g:
        print("Listing all accounts...")
        accounts = list_accounts(g)
        print(f"\nFound {len(accounts)} accounts")

        # Stats
        statuses = {}
        for a in accounts:
            s = a['status'].upper()
            statuses[s] = statuses.get(s, 0) + 1
        print(f"Statuses: {statuses}")

        if dry_run:
            print("\n--dry-run mode: previewing --")
        else:
            print(f"\nSyncing via workflow {WF_ID}...")

        succeeded = 0
        failed = 0
        for i, a in enumerate(accounts):
            name = a['name'] or f"(id={a['id']})"
            if not dry_run:
                sys.stderr.write(f"\r  {i+1}/{len(accounts)}: {name[:40]}...")
                sys.stderr.flush()

            if sync_account(g, a, dry_run):
                succeeded += 1
            else:
                failed += 1

            if not dry_run:
                time.sleep(0.15)  # avoid rate limiting (~7/sec)

        if not dry_run:
            sys.stderr.write("\n")

        print(f"\nDone: {succeeded} started, {failed} failed")
        if dry_run:
            print("(nothing executed — remove --dry-run to run)")


if __name__ == '__main__':
    main()
