#!/usr/bin/env python3
"""
inject.py — CLI wrapper for the ExaBGP Route Injection API

Usage:
  inject.py status
  inject.py routes
  inject.py scenarios
  inject.py announce <prefix> [<prefix>...] [--as-path ASN...] [--community STR...] [--med N] [--next-hop IP]
  inject.py withdraw <prefix> [<prefix>...]
  inject.py withdraw-all
  inject.py scenario <name>
  inject.py withdraw-scenario <name>
  inject.py churn [--count N] [--interval SEC]   # cycle announce/withdraw for ip_rib_log population

Environment:
  EXABGP_API=http://localhost:5050   API base URL
"""

import sys
import os
import json
import time
import argparse
import requests

API = os.environ.get('EXABGP_API', 'http://localhost:5050')


def _post(path, data=None):
    r = requests.post(f'{API}{path}', json=data or {}, timeout=10)
    r.raise_for_status()
    return r.json()


def _delete(path):
    r = requests.delete(f'{API}{path}', timeout=10)
    r.raise_for_status()
    return r.json()


def _get(path):
    r = requests.get(f'{API}{path}', timeout=10)
    r.raise_for_status()
    return r.json()


def _pp(data):
    print(json.dumps(data, indent=2))


def cmd_status(args):
    _pp(_get('/healthz'))


def cmd_routes(args):
    data = _get('/routes')
    print(f"Active routes: {data['count']}")
    for r in data['routes']:
        path = ' '.join(str(a) for a in r.get('as_path', []))
        comms = ' '.join(r.get('communities', []))
        parts = [r['prefix'], f'nh={r["next_hop"]}']
        if path:
            parts.append(f'as-path=[{path}]')
        if comms:
            parts.append(f'community=[{comms}]')
        if r.get('med') is not None:
            parts.append(f'med={r["med"]}')
        print('  ' + '  '.join(parts))


def cmd_scenarios(args):
    data = _get('/scenarios')
    print(f"{'Name':<20} {'Routes':>6}  Description")
    print('-' * 70)
    for name, info in data['scenarios'].items():
        print(f"{name:<20} {info['route_count']:>6}  {info['description']}")


def cmd_announce(args):
    payload = {
        'prefixes': args.prefixes,
        'next_hop': args.next_hop,
        'as_path': [int(a) for a in args.as_path] if args.as_path else [],
        'communities': args.community or [],
    }
    if args.med is not None:
        payload['med'] = args.med
    data = _post('/announce', payload)
    print(f"Announced {data['count']} route(s):")
    for p in data['announced']:
        print(f"  + {p}")


def cmd_withdraw(args):
    data = _post('/withdraw', {'prefixes': args.prefixes})
    print(f"Withdrew {data['count']} route(s):")
    for p in data['withdrawn']:
        print(f"  - {p}")


def cmd_withdraw_all(args):
    data = _post('/withdraw/all')
    print(f"Withdrew all {data['count']} route(s)")


def cmd_scenario(args):
    data = _post(f'/scenario/{args.name}')
    print(f"Loaded scenario '{args.name}': {data['count']} routes announced")


def cmd_withdraw_scenario(args):
    data = _delete(f'/scenario/{args.name}')
    print(f"Withdrew scenario '{args.name}': {data['count']} routes withdrawn")


def cmd_churn(args):
    """
    Cycle announce/withdraw on the 'churn' scenario to generate ip_rib_log
    entries and populate stats_chg_* tables in OpenBMP.
    """
    count = args.count
    interval = args.interval
    print(f"Starting churn: {count} cycles, {interval}s interval")
    print("This will populate ip_rib_log and stats_chg_* tables in OpenBMP.")
    print("Press Ctrl+C to stop.\n")

    cycle = 0
    try:
        while count == 0 or cycle < count:
            cycle += 1
            print(f"Cycle {cycle}: announcing churn scenario...")
            r = _post('/scenario/churn')
            print(f"  + {r['count']} routes announced")
            time.sleep(interval)

            print(f"Cycle {cycle}: withdrawing churn scenario...")
            r = _delete('/scenario/churn')
            print(f"  - {r['count']} routes withdrawn")
            time.sleep(interval)

            print()
    except KeyboardInterrupt:
        print("\nChurn stopped. Withdrawing any active routes...")
        _post('/withdraw/all')
        print("Done.")


def main():
    parser = argparse.ArgumentParser(
        description='ExaBGP Route Injection CLI',
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = parser.add_subparsers(dest='command')

    sub.add_parser('status', help='Show API health and peer states')
    sub.add_parser('routes', help='List active announced routes')
    sub.add_parser('scenarios', help='List available scenarios')
    sub.add_parser('withdraw-all', help='Withdraw all active routes')

    p = sub.add_parser('announce', help='Announce one or more prefixes')
    p.add_argument('prefixes', nargs='+')
    p.add_argument('--as-path', nargs='+', default=[], metavar='ASN')
    p.add_argument('--community', nargs='+', default=[], metavar='COMM')
    p.add_argument('--med', type=int, default=None)
    p.add_argument('--next-hop', default='self', metavar='IP')

    p = sub.add_parser('withdraw', help='Withdraw one or more prefixes')
    p.add_argument('prefixes', nargs='+')

    p = sub.add_parser('scenario', help='Load a named scenario')
    p.add_argument('name')

    p = sub.add_parser('withdraw-scenario', help='Withdraw a named scenario')
    p.add_argument('name')

    p = sub.add_parser('churn', help='Cycle announce/withdraw to populate ip_rib_log')
    p.add_argument('--count', type=int, default=0, metavar='N',
                   help='Number of cycles (0 = infinite)')
    p.add_argument('--interval', type=float, default=30.0, metavar='SEC',
                   help='Seconds between announce and withdraw (default: 30)')

    args = parser.parse_args()

    cmds = {
        'status': cmd_status,
        'routes': cmd_routes,
        'scenarios': cmd_scenarios,
        'announce': cmd_announce,
        'withdraw': cmd_withdraw,
        'withdraw-all': cmd_withdraw_all,
        'scenario': cmd_scenario,
        'withdraw-scenario': cmd_withdraw_scenario,
        'churn': cmd_churn,
    }

    if not args.command:
        parser.print_help()
        sys.exit(1)

    try:
        cmds[args.command](args)
    except requests.exceptions.ConnectionError:
        print(f"ERROR: Cannot connect to ExaBGP API at {API}")
        print("Is the exabgp container running?  docker compose logs exabgp")
        sys.exit(1)
    except requests.exceptions.HTTPError as e:
        print(f"ERROR: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()