#!/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


# neighbor 127.0.0.1 local-ip 127.0.0.1 local-as 1 peer-as 1 router-id 1.2.3.4 family-allowed ipv4-unicast

messages = [
    'announce                                                                              route 1.0.0.0/24 next-hop 101.1.101.1 med 200 community [2:1]',
    'neighbor 127.0.0.1 announce                                                           route 1.1.0.0/24 next-hop 101.1.101.1',
    'neighbor 127.0.0.1 local-as 1 family-allowed ipv4-unicast announce                    route 1.2.0.0/24 next-hop 101.1.101.1',
    'neighbor 127.0.0.1 local-as 1 family-allowed in-open announce                         route 9.9.9.9/24 next-hop 101.1.101.1',
    'neighbor 127.0.0.1 local-as 1 peer-as 1 local-ip 127.0.0.1 router-id 1.2.3.4 announce route 1.3.0.0/24 next-hop 101.1.101.1',
]

while messages:
    message = messages.pop(0)
    sys.stdout.write(message + '\n')
    sys.stdout.flush()
    time.sleep(0.2)

if not wait_for_ack(expected_count=5):
    # 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 + 5:
        line = sys.stdin.readline().strip()
        if not line or 'shutdown' in line:
            break
        time.sleep(1)
except IOError:
    pass
