#!/bin/bash

# HF Spaces Screenshot API 测试脚本
# 使用方法: ./test-hf-api.sh [SPACE_URL] [HF_TOKEN]

# 配置
SPACE_URL="${1:-https://your-username-your-space-name.hf.space}"
HF_TOKEN="${2:-}"
TEST_URL="https://www.baidu.com"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

echo -e "${BLUE}🚀 HF Spaces Screenshot API 测试工具${NC}"
echo "======================================"
echo "Space URL: $SPACE_URL"
echo "Token: ${HF_TOKEN:+已设置}${HF_TOKEN:-未设置}"
echo ""

# 构建认证头
AUTH_HEADER=""
if [ ! -z "$HF_TOKEN" ]; then
    if [[ $HF_TOKEN == hf_* ]]; then
        AUTH_HEADER="Authorization: Bearer $HF_TOKEN"
        echo -e "${GREEN}✅ 使用 HF Token 认证${NC}"
    else
        AUTH_HEADER="X-API-Key: $HF_TOKEN"
        echo -e "${GREEN}✅ 使用自定义 API Key 认证${NC}"
    fi
else
    echo -e "${YELLOW}⚠️  无认证信息，使用公开访问${NC}"
fi

echo ""

# 测试 1: 健康检查
echo -e "${BLUE}📋 测试 1: 健康检查${NC}"
echo "GET $SPACE_URL/"
echo ""

if [ ! -z "$AUTH_HEADER" ]; then
    HEALTH_RESPONSE=$(curl -s -w "HTTP_CODE:%{http_code}" -H "$AUTH_HEADER" "$SPACE_URL/")
else
    HEALTH_RESPONSE=$(curl -s -w "HTTP_CODE:%{http_code}" "$SPACE_URL/")
fi

HTTP_CODE=$(echo "$HEALTH_RESPONSE" | grep -o "HTTP_CODE:[0-9]*" | cut -d: -f2)
RESPONSE_BODY=$(echo "$HEALTH_RESPONSE" | sed 's/HTTP_CODE:[0-9]*$//')

if [ "$HTTP_CODE" -eq 200 ]; then
    echo -e "${GREEN}✅ 健康检查成功 (HTTP $HTTP_CODE)${NC}"
    echo "$RESPONSE_BODY" | python3 -m json.tool 2>/dev/null || echo "$RESPONSE_BODY"
else
    echo -e "${RED}❌ 健康检查失败 (HTTP $HTTP_CODE)${NC}"
    echo "$RESPONSE_BODY"
fi

echo ""
echo "======================================"

# 测试 2: 状态检查
echo -e "${BLUE}📊 测试 2: 状态检查${NC}"
echo "GET $SPACE_URL/status"
echo ""

if [ ! -z "$AUTH_HEADER" ]; then
    STATUS_RESPONSE=$(curl -s -w "HTTP_CODE:%{http_code}" -H "$AUTH_HEADER" "$SPACE_URL/status")
else
    STATUS_RESPONSE=$(curl -s -w "HTTP_CODE:%{http_code}" "$SPACE_URL/status")
fi

HTTP_CODE=$(echo "$STATUS_RESPONSE" | grep -o "HTTP_CODE:[0-9]*" | cut -d: -f2)
RESPONSE_BODY=$(echo "$STATUS_RESPONSE" | sed 's/HTTP_CODE:[0-9]*$//')

if [ "$HTTP_CODE" -eq 200 ]; then
    echo -e "${GREEN}✅ 状态检查成功 (HTTP $HTTP_CODE)${NC}"
    echo "$RESPONSE_BODY" | python3 -m json.tool 2>/dev/null || echo "$RESPONSE_BODY"
else
    echo -e "${RED}❌ 状态检查失败 (HTTP $HTTP_CODE)${NC}"
    echo "$RESPONSE_BODY"
fi

echo ""
echo "======================================"

# 测试 3: 截图测试
echo -e "${BLUE}📸 测试 3: 截图测试${NC}"
echo "POST $SPACE_URL/screenshot"
echo "目标URL: $TEST_URL"
echo ""

# 准备请求数据
REQUEST_DATA='{
  "url": "'$TEST_URL'",
  "width": 1280,
  "height": 720,
  "quality": 75
}'

echo "请求数据:"
echo "$REQUEST_DATA" | python3 -m json.tool 2>/dev/null || echo "$REQUEST_DATA"
echo ""

# 执行截图请求
START_TIME=$(date +%s.%N)

if [ ! -z "$AUTH_HEADER" ]; then
    curl -X POST "$SPACE_URL/screenshot" \
         -H "Content-Type: application/json" \
         -H "$AUTH_HEADER" \
         -d "$REQUEST_DATA" \
         -w "HTTP_CODE:%{http_code}\nTIME_TOTAL:%{time_total}s\nSIZE_DOWNLOAD:%{size_download} bytes\n" \
         -o "test_screenshot_$(date +%s).jpg" \
         -s > curl_output.tmp
else
    curl -X POST "$SPACE_URL/screenshot" \
         -H "Content-Type: application/json" \
         -d "$REQUEST_DATA" \
         -w "HTTP_CODE:%{http_code}\nTIME_TOTAL:%{time_total}s\nSIZE_DOWNLOAD:%{size_download} bytes\n" \
         -o "test_screenshot_$(date +%s).jpg" \
         -s > curl_output.tmp
fi

# 解析结果
HTTP_CODE=$(grep "HTTP_CODE:" curl_output.tmp | cut -d: -f2)
TIME_TOTAL=$(grep "TIME_TOTAL:" curl_output.tmp | cut -d: -f2)
SIZE_DOWNLOAD=$(grep "SIZE_DOWNLOAD:" curl_output.tmp | cut -d: -f2 | cut -d' ' -f1)

if [ "$HTTP_CODE" -eq 200 ]; then
    echo -e "${GREEN}✅ 截图成功 (HTTP $HTTP_CODE)${NC}"
    echo "⏱️  响应时间: $TIME_TOTAL"
    echo "📦 文件大小: $SIZE_DOWNLOAD bytes ($(echo "scale=1; $SIZE_DOWNLOAD/1024" | bc)KB)"
    echo "💾 保存位置: test_screenshot_$(date +%s).jpg"
    
    # 检查文件是否存在且大小合理
    LATEST_FILE=$(ls -t test_screenshot_*.jpg 2>/dev/null | head -1)
    if [ -f "$LATEST_FILE" ] && [ -s "$LATEST_FILE" ]; then
        FILE_SIZE=$(stat -f%z "$LATEST_FILE" 2>/dev/null || stat -c%s "$LATEST_FILE" 2>/dev/null)
        echo "✅ 文件验证通过，大小: ${FILE_SIZE} bytes"
    else
        echo -e "${YELLOW}⚠️  文件可能损坏或为空${NC}"
    fi
else
    echo -e "${RED}❌ 截图失败 (HTTP $HTTP_CODE)${NC}"
    # 显示错误信息
    SCREENSHOT_FILE=$(ls -t test_screenshot_*.jpg 2>/dev/null | head -1)
    if [ -f "$SCREENSHOT_FILE" ]; then
        echo "错误信息:"
        cat "$SCREENSHOT_FILE"
    fi
fi

# 清理临时文件
rm -f curl_output.tmp

echo ""
echo "======================================"

# 测试总结
echo -e "${BLUE}📋 测试总结${NC}"
echo ""

if [ "$HTTP_CODE" -eq 200 ]; then
    echo -e "${GREEN}🎉 所有测试通过！API 工作正常${NC}"
    echo ""
    echo "💡 使用建议:"
    echo "  - 可以在生产环境中使用此 API"
    echo "  - 建议设置适当的速率限制"
    echo "  - 监控服务器资源使用情况"
    
    if [ -z "$HF_TOKEN" ]; then
        echo ""
        echo -e "${YELLOW}🔐 安全建议: 考虑将 Space 设为私有并使用 HF Token${NC}"
    fi
else
    echo -e "${RED}❌ 测试失败，请检查以下问题:${NC}"
    echo "  1. Space URL 是否正确"
    echo "  2. Space 是否正在运行"
    echo "  3. 是否需要认证 token"
    echo "  4. 网络连接是否正常"
fi

echo ""
echo "🔗 更多信息:"
echo "  - Demo界面: $SPACE_URL/demo"
echo "  - API文档: 查看 README.md"
echo "  - 健康检查: $SPACE_URL/"