#!/bin/bash

# ============================================
# 数据源 API 测试脚本
# ============================================
# 测试数据源的 CRUD 和取值功能
# 使用方法: ./test_datasource_api.sh [host] [port]
# 示例: ./test_datasource_api.sh localhost 5232
# ============================================

# 配置参数
HOST=${1:-localhost}
PORT=${2:-5232}
BASE_URL="http://${HOST}:${PORT}"
DATASOURCE_URL="${BASE_URL}/api/v1/datasource"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 输出函数
print_header() {
    echo -e "\n${BLUE}============================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}============================================${NC}"
}

print_step() {
    echo -e "\n${CYAN}>>> $1${NC}"
}

print_success() {
    echo -e "${GREEN}✓ $1${NC}"
}

print_error() {
    echo -e "${RED}✗ $1${NC}"
}

print_info() {
    echo -e "${YELLOW}ℹ $1${NC}"
}

# 变量存储
DATASOURCE_ID=""
DATASOURCE_ID_2=""

# ============================================
# 测试函数
# ============================================

test_create_text_datasource() {
    print_step "创建文本类型数据源"
    
    RESPONSE=$(curl -s -X POST "${DATASOURCE_URL}" \
        -H "Content-Type: application/json" \
        -H "X-User-Id: test-user-001" \
        -d '{
            "name": "项目名称",
            "type": "text",
            "documentId": "test-doc-001",
            "sourceType": "manual",
            "valueType": "text",
            "aggregateType": "first",
            "nodeIds": {
                "refs": [
                    {"type": "graph_node", "id": "node-001"}
                ]
            },
            "metadata": {"description": "测试数据源"}
        }')
    
    echo "Response: $RESPONSE"
    
    # 提取 ID
    DATASOURCE_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
    
    if [ -n "$DATASOURCE_ID" ]; then
        print_success "创建成功，ID: $DATASOURCE_ID"
    else
        print_error "创建失败"
        return 1
    fi
}

test_create_image_datasource() {
    print_step "创建图片类型数据源"
    
    RESPONSE=$(curl -s -X POST "${DATASOURCE_URL}" \
        -H "Content-Type: application/json" \
        -H "X-User-Id: test-user-001" \
        -d '{
            "name": "公司LOGO",
            "type": "image",
            "documentId": "test-doc-001",
            "sourceType": "file",
            "valueType": "image",
            "nodeIds": {
                "refs": [
                    {"type": "document_element", "id": "element-001"}
                ]
            }
        }')
    
    echo "Response: $RESPONSE"
    
    DATASOURCE_ID_2=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
    
    if [ -n "$DATASOURCE_ID_2" ]; then
        print_success "创建成功，ID: $DATASOURCE_ID_2"
    else
        print_error "创建失败"
        return 1
    fi
}

test_get_datasource() {
    print_step "获取数据源详情"
    
    if [ -z "$DATASOURCE_ID" ]; then
        print_error "没有可用的数据源ID"
        return 1
    fi
    
    RESPONSE=$(curl -s -X GET "${DATASOURCE_URL}/${DATASOURCE_ID}")
    echo "Response: $RESPONSE"
    
    if echo "$RESPONSE" | grep -q '"code":200'; then
        print_success "获取成功"
    else
        print_error "获取失败"
        return 1
    fi
}

test_get_by_document() {
    print_step "按文档ID查询数据源"
    
    RESPONSE=$(curl -s -X GET "${DATASOURCE_URL}/document/test-doc-001")
    echo "Response: $RESPONSE"
    
    if echo "$RESPONSE" | grep -q '"code":200'; then
        print_success "查询成功"
    else
        print_error "查询失败"
        return 1
    fi
}

test_get_by_user() {
    print_step "按用户ID查询数据源"
    
    RESPONSE=$(curl -s -X GET "${DATASOURCE_URL}/user/test-user-001")
    echo "Response: $RESPONSE"
    
    if echo "$RESPONSE" | grep -q '"code":200'; then
        print_success "查询成功"
    else
        print_error "查询失败"
        return 1
    fi
}

test_update_datasource() {
    print_step "更新数据源"
    
    if [ -z "$DATASOURCE_ID" ]; then
        print_error "没有可用的数据源ID"
        return 1
    fi
    
    RESPONSE=$(curl -s -X PUT "${DATASOURCE_URL}/${DATASOURCE_ID}" \
        -H "Content-Type: application/json" \
        -d '{
            "name": "项目名称（已更新）",
            "aggregateType": "concat",
            "separator": ", "
        }')
    
    echo "Response: $RESPONSE"
    
    if echo "$RESPONSE" | grep -q '"code":200'; then
        print_success "更新成功"
    else
        print_error "更新失败"
        return 1
    fi
}

test_update_refs() {
    print_step "更新节点引用"
    
    if [ -z "$DATASOURCE_ID" ]; then
        print_error "没有可用的数据源ID"
        return 1
    fi
    
    # 追加模式
    RESPONSE=$(curl -s -X PUT "${DATASOURCE_URL}/${DATASOURCE_ID}/refs" \
        -H "Content-Type: application/json" \
        -d '{
            "mode": "append",
            "refs": [
                {"type": "graph_node", "id": "node-002"},
                {"type": "document_element", "id": "element-002"}
            ]
        }')
    
    echo "Response: $RESPONSE"
    
    if echo "$RESPONSE" | grep -q '"code":200'; then
        print_success "追加引用成功"
    else
        print_error "追加引用失败"
        return 1
    fi
}

test_get_value() {
    print_step "获取数据源值"
    
    if [ -z "$DATASOURCE_ID" ]; then
        print_error "没有可用的数据源ID"
        return 1
    fi
    
    RESPONSE=$(curl -s -X GET "${DATASOURCE_URL}/${DATASOURCE_ID}/value")
    echo "Response: $RESPONSE"
    
    if echo "$RESPONSE" | grep -q '"code":200'; then
        print_success "获取值成功"
    else
        print_error "获取值失败"
        return 1
    fi
}

test_batch_get_value() {
    print_step "批量获取数据源值"
    
    if [ -z "$DATASOURCE_ID" ] || [ -z "$DATASOURCE_ID_2" ]; then
        print_error "没有足够的数据源ID"
        return 1
    fi
    
    RESPONSE=$(curl -s -X POST "${DATASOURCE_URL}/batch-value" \
        -H "Content-Type: application/json" \
        -d "{
            \"dataSourceIds\": [\"${DATASOURCE_ID}\", \"${DATASOURCE_ID_2}\"]
        }")
    
    echo "Response: $RESPONSE"
    
    if echo "$RESPONSE" | grep -q '"code":200'; then
        print_success "批量获取值成功"
    else
        print_error "批量获取值失败"
        return 1
    fi
}

test_delete_datasource() {
    print_step "删除数据源"
    
    if [ -z "$DATASOURCE_ID" ]; then
        print_error "没有可用的数据源ID"
        return 1
    fi
    
    RESPONSE=$(curl -s -X DELETE "${DATASOURCE_URL}/${DATASOURCE_ID}")
    echo "Response: $RESPONSE"
    
    if echo "$RESPONSE" | grep -q '"code":200'; then
        print_success "删除成功"
    else
        print_error "删除失败"
        return 1
    fi
    
    # 删除第二个
    if [ -n "$DATASOURCE_ID_2" ]; then
        curl -s -X DELETE "${DATASOURCE_URL}/${DATASOURCE_ID_2}" > /dev/null
        print_success "删除第二个数据源"
    fi
}

# ============================================
# 主流程
# ============================================

show_help() {
    echo "数据源 API 测试脚本"
    echo ""
    echo "用法: $0 [选项] [host] [port]"
    echo ""
    echo "选项:"
    echo "  -h, --help     显示帮助信息"
    echo "  -a, --all      运行所有测试（默认）"
    echo "  -c, --create   仅测试创建"
    echo "  -r, --read     仅测试查询"
    echo "  -u, --update   仅测试更新"
    echo "  -v, --value    仅测试取值"
    echo "  -d, --delete   仅测试删除"
    echo ""
    echo "示例:"
    echo "  $0                    # 默认测试 localhost:5232"
    echo "  $0 192.168.1.100 8080 # 指定 host 和 port"
    echo "  $0 -c                 # 仅测试创建"
}

run_all_tests() {
    print_header "数据源 API 测试开始"
    print_info "目标地址: ${BASE_URL}"
    
    echo ""
    print_header "1. 创建测试"
    test_create_text_datasource
    test_create_image_datasource
    
    echo ""
    print_header "2. 查询测试"
    test_get_datasource
    test_get_by_document
    test_get_by_user
    
    echo ""
    print_header "3. 更新测试"
    test_update_datasource
    test_update_refs
    
    echo ""
    print_header "4. 取值测试"
    test_get_value
    test_batch_get_value
    
    echo ""
    print_header "5. 删除测试"
    test_delete_datasource
    
    echo ""
    print_header "测试完成"
}

# 解析命令行参数
case "$1" in
    -h|--help)
        show_help
        exit 0
        ;;
    -a|--all)
        shift
        HOST=${1:-localhost}
        PORT=${2:-5232}
        BASE_URL="http://${HOST}:${PORT}"
        DATASOURCE_URL="${BASE_URL}/api/v1/datasource"
        run_all_tests
        ;;
    -c|--create)
        shift
        HOST=${1:-localhost}
        PORT=${2:-5232}
        BASE_URL="http://${HOST}:${PORT}"
        DATASOURCE_URL="${BASE_URL}/api/v1/datasource"
        print_header "创建测试"
        test_create_text_datasource
        test_create_image_datasource
        ;;
    -v|--value)
        shift
        HOST=${1:-localhost}
        PORT=${2:-5232}
        BASE_URL="http://${HOST}:${PORT}"
        DATASOURCE_URL="${BASE_URL}/api/v1/datasource"
        print_header "取值测试"
        test_create_text_datasource
        test_get_value
        ;;
    *)
        run_all_tests
        ;;
esac