#!/bin/bash

set -euo pipefail

# Function to print usage
usage() {
    echo "Usage: $0 --client CLIENT_ID --host-name HOST_NAME --api-server API_SERVER --api-token API_TOKEN"
    exit 1
}

# Parse args
CLIENT_ID=""
HOST_NAME=""
API_SERVER=""
API_TOKEN=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --client) CLIENT_ID="$2"; shift 2 ;;
        --host-name) HOST_NAME="$2"; shift 2 ;;
        --api-server) API_SERVER="$2"; shift 2 ;;
        --api-token) API_TOKEN="$2"; shift 2 ;;
        *) echo "Unknown option: $1"; usage ;;
    esac
done

[[ -z "$CLIENT_ID" || -z "$HOST_NAME" || -z "$API_SERVER" || -z "$API_TOKEN" ]] && usage

AUTH_HEADER="Authorization: Bearer $API_TOKEN"

echo "🔍 Finding machine with host_name='$HOST_NAME' under client='$CLIENT_ID'..."
MACHINE_JSON=$(curl -s -H "$AUTH_HEADER" "$API_SERVER/api/v1/machine/?client_id=$CLIENT_ID&host_name=$HOST_NAME")
MACHINE_COUNT=$(echo "$MACHINE_JSON" | jq '.results | length')

if [[ "$MACHINE_COUNT" -eq 0 ]]; then
    echo "❌ No matching machine found. Exiting."
    exit 1
elif [[ "$MACHINE_COUNT" -gt 1 ]]; then
    echo "❌ Multiple machines found with name '$HOST_NAME'. Exiting."
    exit 1
fi

MACHINE_PK=$(echo "$MACHINE_JSON" | jq -r '.results[0].pk')
echo "✅ Found machine with PK: $MACHINE_PK"

echo "📥 Requesting script task for machine..."
TASK_ID=$(curl -s -H "$AUTH_HEADER" "$API_SERVER/api/v1/machine/$MACHINE_PK/script/" | jq -r '.')

echo "⏳ Waiting for script task $TASK_ID to complete..."
while true; do
    TASK_STATUS=$(curl -s -H "$AUTH_HEADER" "$API_SERVER/api/v1/tasks/$TASK_ID/")
    COMPLETE=$(echo "$TASK_STATUS" | jq -r '.complete')
    FAILED=$(echo "$TASK_STATUS" | jq -r '.failed')
    if [[ "$FAILED" == "true" ]]; then
        echo "❌ Task failed: $(echo "$TASK_STATUS" | jq -r '.traceback')"
        exit 1
    elif [[ "$COMPLETE" == "true" ]]; then
        echo "✅ Task completed."
        break
    fi
    sleep 2
done

RELATIVE_PATH=$(echo "$TASK_STATUS" | jq -r '.download')
DOWNLOAD_URL="$API_SERVER$RELATIVE_PATH"
#DOWNLOAD_URL=$(echo "$TASK_STATUS" | jq -r '.download')
echo "📦 Downloading script from $DOWNLOAD_URL..."
curl -s -H "$AUTH_HEADER" "$DOWNLOAD_URL" -o script.zip

echo "📂 Extracting script.zip..."
rm -rf extracted && mkdir extracted
unzip -q script.zip -d extracted

SCRIPT_PATH="extracted/xylok-collect.sh"
echo "🚀 Running script $SCRIPT_PATH..."
chmod +x "$SCRIPT_PATH"
bash "$SCRIPT_PATH"
echo "✅ Script complete"

RESULT_FILE=$(find extracted -name "*.xylok")
if [[ -z "$RESULT_FILE" ]]; then
    echo "❌ .xylok result file not found."
    exit 1
fi

echo "📤 Uploading result: $RESULT_FILE"
UPLOAD_RESP=$(curl -s -H "$AUTH_HEADER" \
  --form "file=@$RESULT_FILE" \
  "$API_SERVER/api/v1/upload/")

TASK_ID=$(echo "$UPLOAD_RESP" | jq -r '.[0]')

echo "⏳ Waiting for upload task $TASK_ID to complete..."
while true; do
    TASK_STATUS=$(curl -s -H "$AUTH_HEADER" "$API_SERVER/api/v1/tasks/$TASK_ID/")
    COMPLETE=$(echo "$TASK_STATUS" | jq -r '.complete')
    FAILED=$(echo "$TASK_STATUS" | jq -r '.failed')
    if [[ "$FAILED" == "true" ]]; then
        echo "❌ Upload task failed: $(echo "$TASK_STATUS" | jq -r '.traceback')"
        exit 1
    elif [[ "$COMPLETE" == "true" ]]; then
        echo "✅ Upload complete."
        break
    fi
    sleep 2
done

echo "📦 Fetching scan IDs for machine..."
SCAN_RESP=$(curl -s -H "$AUTH_HEADER" "$API_SERVER/api/v1/machine/$MACHINE_PK/scans/?limit=2")
SCAN_COUNT=$(echo "$SCAN_RESP" | jq '.results | length')

if [[ "$SCAN_COUNT" -eq 0 ]]; then
    echo "❌ No scans found."
    exit 1
fi

NEW_SCAN_PK=$(echo "$SCAN_RESP" | jq -r '.results[0].pk')
echo "📈 New scan ID: $NEW_SCAN_PK"

if [[ "$SCAN_COUNT" -gt 1 ]]; then
    OLD_SCAN_PK=$(echo "$SCAN_RESP" | jq -r '.results[1].pk')
    echo "🔁 Copying answers from $OLD_SCAN_PK to $NEW_SCAN_PK..."
    curl -s -X POST -H "$AUTH_HEADER" "$API_SERVER/api/v1/scans/$OLD_SCAN_PK/copy-answers-to/$NEW_SCAN_PK/" > /dev/null
fi

echo "🧠 Running automated analysis on scan $NEW_SCAN_PK..."
AA_RESP=$(curl -s -X POST -H "$AUTH_HEADER" "$API_SERVER/api/v1/scans/$NEW_SCAN_PK/aa/execute/")
AA_TASK_ID=$(echo "$AA_RESP" | jq -r '.')

echo "⏳ Waiting for AA task $AA_TASK_ID to complete..."
while true; do
    TASK_STATUS=$(curl -s -H "$AUTH_HEADER" "$API_SERVER/api/v1/tasks/$AA_TASK_ID/")
    COMPLETE=$(echo "$TASK_STATUS" | jq -r '.complete')
    SUCCESS=$(echo "$TASK_STATUS" | jq -r '.success')
    if [[ "$COMPLETE" == "true" ]]; then
        if [[ "$SUCCESS" == "true" ]]; then
            echo "✅ Scan run, uploaded, and analyzed successfully"
            exit 0
        else
            echo "❌ Scan AA failed:"
            echo "$TASK_STATUS" | jq -r '.traceback'
            exit 1
        fi
    fi
    sleep 2
done