#!/usr/bin/env python3

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


messages = [
    'announce route 2.2.2.2/32 next-hop 22.22.22.22',
    'announce route 3.3.3.3/32 next-hop 33.33.33.33',
]

time.sleep(0.2)
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=2):
    # 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
