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


# Test sequence:
# 1. Send a regular announce command and expect ACK (baseline test)
# 2. Send disable-ack command and expect ACK (last ACK before disabling)
# 3. Send another announce command and expect NO ACK (should timeout)
# 4. Send enable-ack command and expect ACK (re-enabling)
# 5. Send final announce command and expect ACK (verify re-enabled)

time.sleep(0.2)  # let the EOR pass

# Step 1: Baseline - send command and expect ACK
sys.stdout.write('announce route 1.1.0.0/24 next-hop 101.1.101.1\n')
sys.stdout.flush()
time.sleep(0.1)

if not wait_for_ack(expected_count=1):
    sys.stderr.write('ERROR: Baseline test failed - no ACK received for announce\n')
    sys.exit(1)

# Step 2: Send disable-ack command and expect final ACK
sys.stdout.write('disable-ack\n')
sys.stdout.flush()
time.sleep(0.1)

if not wait_for_ack(expected_count=1):
    sys.stderr.write('ERROR: No ACK received for disable-ack command\n')
    sys.exit(1)

# Step 3: Send command and expect NO ACK (should timeout)
sys.stdout.write('announce route 1.2.0.0/24 next-hop 101.1.101.1\n')
sys.stdout.flush()
time.sleep(0.1)

if wait_for_ack(expected_count=1, timeout=1):
    sys.stderr.write('ERROR: Received unexpected ACK after disable-ack\n')
    sys.exit(1)

# Step 4: Send enable-ack command and expect ACK
sys.stdout.write('enable-ack\n')
sys.stdout.flush()
time.sleep(0.1)

if not wait_for_ack(expected_count=1):
    sys.stderr.write('ERROR: No ACK received for enable-ack command\n')
    sys.exit(1)

# Step 5: Send final command and expect ACK (verify re-enabled)
sys.stdout.write('announce route 1.3.0.0/24 next-hop 101.1.101.1\n')
sys.stdout.flush()
time.sleep(0.1)

if not wait_for_ack(expected_count=1):
    sys.stderr.write('ERROR: No ACK received after enable-ack\n')
    sys.exit(1)

# All tests passed
sys.stderr.write('SUCCESS: All ACK control tests passed\n')

# 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
