#!/usr/bin/env python3
"""
Clear all MAC->IP address assignments from customer inventory items.

Two-phase approach:
  1. List all active IP assignments via paginated V1 public API
  2. Expire each one via startWorkflowExecution mutation

Before running:
  1. Create the "Expire Single IP Assignment" workflow in Gaiia UI
     (use expire_one_workflow.json)
  2. Copy its workflow ID and set EXPIRE_WF_ID below
  3. source .envrc
  4. uv run python3 clear_all_ips.py --dry-run   (preview)
  5. uv run python3 clear_all_ips.py              (execute)
"""

from gaiia import GaiiaClient, GaiiaMutationError
import json
import os
import sys
import time

EXPIRE_WF_ID = os.environ.get("EXPIRE_WF_ID", "")

def list_all_assignments(client):
    """Paginate through all inventory items and return active MAC->IP pairs."""
    all_assignments = []
    after = None
    page = 0

    while True:
        page += 1
        query = """
        query Q($after: String) {
            inventoryItems(
                filters: {assigneeType: {equals: ACCOUNT}}
                first: 250
                after: $after
            ) {
                nodes {
                    ipAddressAssignments {
                        nodes { ipAddress macAddress isActive }
                    }
                }
                pageInfo { hasNextPage endCursor }
                totalCount
            }
        }
        """
        data = client.query(query, {"after": after} if after else {})
        result = data['inventoryItems']

        if page == 1:
            print(f"Total inventory items: {result['totalCount']}")

        for item in result['nodes']:
            for a in item.get('ipAddressAssignments', {}).get('nodes', []):
                if a.get('isActive') != False:
                    all_assignments.append({
                        'macAddress': a['macAddress'],
                        'ipAddress': a['ipAddress']
                    })

        page_info = result['pageInfo']
        sys.stderr.write(f"\r  Page {page}: {len(all_assignments)} assignments so far...")
        sys.stderr.flush()

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

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


def expire_one(client, mac_address, ip_address, dry_run=False):
    """Call the Expire Single IP Assignment workflow via startWorkflowExecution."""
    if dry_run:
        print(f"  [DRY-RUN] {mac_address} -> {ip_address}")
        return True

    mutation = """
    mutation Expire($input: StartWorkflowExecutionInput!) {
        startWorkflowExecution(input: $input) {
            workflowExecution { id status }
        }
    }
    """
    try:
        client.query(mutation, {
            "input": {
                "workflowId": EXPIRE_WF_ID,
                "input": {"macAddress": mac_address, "ipAddress": ip_address}
            }
        })
        return True
    except GaiiaMutationError as e:
        print(f"  FAILED {mac_address} -> {ip_address}: {e}")
        return False


def main():
    if not EXPIRE_WF_ID:
        print("ERROR: Set EXPIRE_WF_ID env var to the workflow ID of 'Expire Single IP Assignment'")
        print("Example: EXPIRE_WF_ID=workflow_xxx uv run python3 clear_all_ips.py --dry-run")
        sys.exit(1)

    dry_run = "--dry-run" in sys.argv

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

        if dry_run:
            print("\n--dry-run mode: showing what would be expired --")
        else:
            print(f"\nExpiring via workflow {EXPIRE_WF_ID}...")

        succeeded = 0
        failed = 0
        for i, a in enumerate(assignments):
            if not dry_run:
                sys.stderr.write(f"\r  {i+1}/{len(assignments)}: {a['macAddress']} -> {a['ipAddress']}")
                sys.stderr.flush()
            if expire_one(g, a['macAddress'], a['ipAddress'], dry_run):
                succeeded += 1
            else:
                failed += 1
            if not dry_run:
                time.sleep(0.1)  # avoid rate limiting

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

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


if __name__ == '__main__':
    main()