#!/usr/pkg/bin/python3.13

import os
import sys
import time
import select
import json


def wait_for_ack(expected_count=1, timeout=2):
    """
    Wait for ACK responses from ExaBGP.

    Polls STDIN until all expected ACK messages are received.
    ExaBGP may not respond immediately, so we loop with sleep.

    Handles both text and JSON encoder formats:
    - Text: "done", "error", "shutdown"
    - JSON: {"answer": "done|error|shutdown", "message": "..."}

    Args:
        expected_count: Number of ACK messages expected (default: 1)
        timeout: Total timeout in seconds (default: 2)

    Returns:
        True if all ACKs received successfully
        False if any command failed or timeout occurred

    Raises:
        SystemExit: If ExaBGP sends shutdown message
    """
    import json
    received = 0
    start_time = time.time()

    while received < expected_count:
        # Check if we've exceeded timeout
        elapsed = time.time() - start_time
        if elapsed >= timeout:
            return False

        # Poll for data (non-blocking with short timeout)
        ready, _, _ = select.select([sys.stdin], [], [], 0.1)

        if ready:
            line = sys.stdin.readline().strip()

            # Parse response (could be text or JSON)
            answer = None
            if line.startswith('{'):
                # JSON format: {"answer": "done|error|shutdown", ...}
                try:
                    data = json.loads(line)
                    answer = data.get('answer')
                except:
                    pass
            else:
                # Text format: done|error|shutdown
                answer = line

            if answer == "done":
                received += 1
            elif answer == "error":
                return False
            elif answer == "shutdown":
                raise SystemExit(0)
            # Ignore other messages (could be BGP updates if receiving)
        else:
            # No data yet, sleep briefly before next poll
            time.sleep(0.1)

    return True


time.sleep(2)  # let the EOR pass


messages = [
    'announce flow route { match { source 4.4.4.4/32; } then { rate-limit 0; } }',
    'announce flow route { match { source 5.5.5.5/32; } then { discard; } }',
    'announce flow route { match { source 6.6.6.6/32; } scope { interface-set [ non-transitive:input:3405770241:1 ]; } then { discard; } }',
    'announce flow route { match { source 7.7.7.7/32; } scope { interface-set [ non-transitive:input:3405770241:1 transitive:output:254:254 ]; } then { discard; } }',
    'announce flow route { match { source 8.8.8.8/32; } then { discard; } scope { interface-set [ non-transitive:input:3405770241:1 transitive:output:254:254 ]; } }',
    'announce flow route { match { source 9.9.9.9/32; } then { discard; } scope { interface-set [ non-transitive:input:3405770241:1 transitive:output:254:254 ]; } }',
    'announce flow route { match { source 10.10.10.10/32; } scope { interface-set [ transitive:input-output:1234:10 transitive:input:1234:10 transitive:output:0:0]; } then { discard; } }',
    'announce flow route destination 133.130.1.219/32 discard interface-set [ transitive:input:1234:10]',
    'announce flow route destination 133.130.1.19/32 interface-set [ transitive:input:1234:10 ] discard',
]

# Send all messages first
message_count = 0
for message in messages:
    try:
        sys.stdout.write(message + '\n')
        sys.stdout.flush()
        message_count += 1
    except IOError:
        sys.stderr.write('IOError - ExaBGP exited ? ')
        sys.stderr.write('Could not write message: %s\n' % message)
        sys.stderr.flush()

# Wait for all ACKs at once
if message_count > 0:
    if not wait_for_ack(expected_count=message_count):
        # ACK not received, continuing anyway (ACKs may be disabled in config)
        pass

# Wait for shutdown signal
try:
    now = time.time()
    while os.getppid() != 1 and time.time() < now + 1:
        line = sys.stdin.readline().strip()
        if not line or 'shutdown' in line:
            break
        time.sleep(0.5)
except IOError:
    pass
