#!/bin/bash
# End-to-End Test Script for AAA Framework
# This script uses nohup to start agents in the background and sends a task

set -e

TASK_ID="${1:-task-cve-2024-32964-ssrf}"
BASE_DIR="$(cd "$(dirname "$0")" && pwd)"
VENV_PATH="$BASE_DIR/.venv/bin/activate"

# Check if API key is set
if [ -z "$OPENAI_API_KEY" ]; then
    echo "Error: OPENAI_API_KEY environment variable is not set"
    echo "Please set it with: export OPENAI_API_KEY='your-key-here'"
    exit 1
fi

echo "============================================================"
echo "AAA Security Assessment Framework - End-to-End Test"
echo "============================================================"
echo "Task ID: $TASK_ID"
echo "Base Directory: $BASE_DIR"
echo

# Clean up any existing agents
echo "[0/4] Cleaning up existing agent processes..."
pkill -f "python.*main.py green" || true
pkill -f "python.*main.py white" || true
sleep 2

# Activate virtual environment
source "$VENV_PATH"

# Start Green Agent in background
echo "[1/4] Starting Green Agent on port 9001..."
nohup python "$BASE_DIR/main.py" green --port 9001 > /tmp/green_agent.log 2>&1 &
GREEN_PID=$!
echo "      Green Agent PID: $GREEN_PID"

# Wait for Green Agent to start
sleep 5

# Check if Green Agent is running
if ! ps -p $GREEN_PID > /dev/null; then
    echo "      Error: Green Agent failed to start"
    cat /tmp/green_agent.log
    exit 1
fi

# Test Green Agent
if curl -s -m 3 http://localhost:9001/.well-known/agent-card.json > /dev/null 2>&1; then
    echo "      Green Agent ready at http://localhost:9001"
else
    echo "      Warning: Green Agent not responding to HTTP requests"
    echo "      Continuing anyway..."
fi

# Start White Agent in background
echo "[2/4] Starting White Agent on port 9002..."
nohup python "$BASE_DIR/main.py" white --port 9002 > /tmp/white_agent.log 2>&1 &
WHITE_PID=$!
echo "      White Agent PID: $WHITE_PID"

# Wait for White Agent to start
sleep 5

# Check if White Agent is running
if ! ps -p $WHITE_PID > /dev/null; then
    echo "      Error: White Agent failed to start"
    cat /tmp/white_agent.log
    kill $GREEN_PID
    exit 1
fi

if curl -s -m 3 http://localhost:9002/.well-known/agent-card.json > /dev/null 2>&1; then
    echo "      White Agent ready at http://localhost:9002"
else
    echo "      Warning: White Agent not responding to HTTP requests"
    echo "      Continuing anyway..."
fi

# Send task to Green Agent
echo "[3/4] Sending task to Green Agent..."
echo "      This may take several minutes (Docker startup + LLM calls + verification)..."
echo

python -c "
import asyncio
import sys
from src.my_util import my_a2a

async def send_task():
    task_id = '$TASK_ID'
    green_url = 'http://localhost:9001'
    white_url = 'http://localhost:9002/'

    task_message = f'''
<task_id>{task_id}</task_id>
<white_agent_url>{white_url}</white_agent_url>
'''

    try:
        response = await my_a2a.send_message(green_url, task_message)

        from a2a.types import SendMessageSuccessResponse, Message
        from a2a.utils import get_text_parts

        res_root = response.root
        if isinstance(res_root, SendMessageSuccessResponse):
            res_result = res_root.result
            if isinstance(res_result, Message):
                text_parts = get_text_parts(res_result.parts)
                if text_parts:
                    print()
                    print('=' * 60)
                    print('[4/4] EVALUATION RESULTS')
                    print('=' * 60)
                    print(text_parts[0])
                    return 0
        print('Error: Unexpected response format')
        return 1
    except Exception as e:
        print(f'Error sending task: {e}')
        import traceback
        traceback.print_exc()
        return 1

exit_code = asyncio.run(send_task())
sys.exit(exit_code)
"

TASK_EXIT_CODE=$?

# Clean up
echo
echo "Stopping agents..."
kill $GREEN_PID $WHITE_PID 2>/dev/null || true
sleep 2

echo "Agents stopped."
echo

if [ $TASK_EXIT_CODE -eq 0 ]; then
    echo "============================================================"
    echo "Test completed successfully!"
    echo "============================================================"
    echo "Check results in: $BASE_DIR/results/"
    ls -lht "$BASE_DIR/results/" | head -5
else
    echo "============================================================"
    echo "Test failed with exit code: $TASK_EXIT_CODE"
    echo "============================================================"
    echo "Check logs:"
    echo "  - Green Agent: /tmp/green_agent.log"
    echo "  - White Agent: /tmp/white_agent.log"
fi

exit $TASK_EXIT_CODE