#!/bin/bash
# =================================================================
# 织忆 (MemoryWeave) 端到端集成测试
# =================================================================
# API base: http://localhost:7821
# API Key:  zhiyi-dev-key-2026
# =================================================================

# ── 配置 ──────────────────────────────────────────────────────
API_BASE="http://localhost:7821"
API_KEY="zhiyi-dev-key-2026"
CLI="/home/muc/.local/bin/zhiyi-cli"
NAMESPACE="test-integration"
AGENT_ID="test-agent-$$"

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

# 测试计数器
TESTS_RUN=0
TESTS_PASSED=0
TESTS_FAILED=0

# ── 辅助函数 ──────────────────────────────────────────────────

log_info()  { echo -e "${BLUE}[INFO]${NC}  $*"; }
log_pass()  { echo -e "${GREEN}[PASS]${NC}  $*"; TESTS_PASSED=$((TESTS_PASSED+1)); }
log_fail()  { echo -e "${RED}[FAIL]${NC}  $*" >&2; TESTS_FAILED=$((TESTS_FAILED+1)); }
log_warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }

header()    { echo ""; echo "══════════════════════════════════════"; echo "  $*"; echo "══════════════════════════════════════"; }

# API 调用封装（不使用 -f，允许非 2xx 响应）
api_get() {
    local path="$1"
    curl -s -H "X-API-Key: ${API_KEY}" "${API_BASE}${path}" 2>&1
}

api_post() {
    local path="$1"
    local body="$2"
    curl -s -H "X-API-Key: ${API_KEY}" \
         -H "Content-Type: application/json" \
         -d "${body}" \
         "${API_BASE}${path}" 2>&1
}

# 断言函数
assert_eq() {
    local expected="$1"
    local actual="$2"
    local msg="$3"
    TESTS_RUN=$((TESTS_RUN+1))
    if [[ "$expected" == "$actual" ]]; then
        log_pass "${msg}"
        return 0
    else
        log_fail "${msg} - expected: '${expected}', got: '${actual}'"
        return 1
    fi
}

assert_contains() {
    local haystack="$1"
    local needle="$2"
    local msg="$3"
    TESTS_RUN=$((TESTS_RUN+1))
    if echo "$haystack" | grep -q "$needle"; then
        log_pass "${msg}"
        return 0
    else
        log_fail "${msg} - '$needle' not found in response"
        return 1
    fi
}

assert_not_empty() {
    local value="$1"
    local msg="$2"
    TESTS_RUN=$((TESTS_RUN+1))
    if [[ -n "$value" ]]; then
        log_pass "${msg}"
        return 0
    else
        log_fail "${msg} - empty value"
        return 1
    fi
}

assert_json_valid() {
    local json_str="$1"
    local msg="$2"
    TESTS_RUN=$((TESTS_RUN+1))
    if echo "$json_str" | python3 -c "import sys,json; json.load(sys.stdin); print('ok')" > /dev/null 2>&1; then
        log_pass "${msg}"
        return 0
    else
        log_fail "${msg} - invalid JSON"
        return 1
    fi
}

# ── SETUP ─────────────────────────────────────────────────────

setup() {
    log_info "执行测试前检查..."
    
    # 检查服务是否运行
    local health_resp
    health_resp=$(curl -s "${API_BASE}/health" 2>&1)
    if ! echo "$health_resp" | python3 -c "import sys,json; json.load(sys.stdin)" > /dev/null 2>&1; then
        log_fail "服务未运行，请先启动 zhiyid (health check failed)"
        exit 1
    fi
    
    # 检查 CLI 是否存在
    if [[ ! -x "$CLI" ]]; then
        log_warn "CLI 未找到或不可执行: $CLI"
    fi
    
    log_info "SETUP 完成"
}

# ── TEARDOWN ──────────────────────────────────────────────────

teardown() {
    log_info "清理测试数据..."
    log_info "TEARDOWN 完成"
}

# ── 测试用例 ──────────────────────────────────────────────────

test_health() {
    header "测试: 健康检查 /health"
    local resp
    resp=$(api_get "/health")
    
    assert_contains "$resp" '"status":"ok"' "health 端点返回 ok 状态"
    assert_contains "$resp" '"service":"zhiyid"' "health 端点包含服务名"
}

test_stats() {
    header "测试: 统计接口 /api/v1/stats"
    local resp
    resp=$(api_get "/api/v1/stats")
    
    assert_contains "$resp" '"total_memories"' "stats 包含记忆总数字段"
    assert_contains "$resp" '"total_episodes"' "stats 包含 episode 总数字段"
    assert_contains "$resp" '"backend"' "stats 包含后端类型字段"
    
    log_info "stats 响应: $(echo $resp | head -c 200)..."
}

test_graph_stats() {
    header "测试: 图谱统计 /api/v1/graph/stats"
    local resp
    resp=$(api_get "/api/v1/graph/stats")
    
    assert_contains "$resp" '"node_count"' "graph/stats 包含节点数"
    assert_contains "$resp" '"edge_count"' "graph/stats 包含边数"
}

test_commit_recall_feedback() {
    header "测试: commit -> recall -> feedback 完整链路"
    
    # Step 1: Commit
    log_info "步骤1: 提交记忆 (commit)"
    local unique_content="集成测试记忆 $(date +%s) - 随机内容: $RANDOM"
    local commit_resp
    commit_resp=$(api_post "/api/v1/commit" "{
        \"agent_id\": \"${AGENT_ID}\",
        \"namespace\": \"${NAMESPACE}\",
        \"content\": \"${unique_content}\",
        \"category\": \"test\"
    }")
    
    log_info "commit 响应: $commit_resp"
    assert_contains "$commit_resp" '"status"' "commit 包含状态字段"
    
    # 提取 memory_id（如果存在）
    local memory_id=""
    if echo "$commit_resp" | grep -q '"memory_id"'; then
        memory_id=$(echo "$commit_resp" | grep -o '"memory_id":"[^"]*"' | cut -d'"' -f4 | head -1)
    fi
    
    local episode_id=""
    if echo "$commit_resp" | grep -q '"episode_id"'; then
        episode_id=$(echo "$commit_resp" | grep -o '"episode_id":"[^"]*"' | cut -d'"' -f4 | head -1)
    fi
    
    log_info "获取到 episode_id: ${episode_id}, memory_id: ${memory_id}"
    assert_not_empty "$episode_id" "commit 返回 episode_id"
    
    # 等待索引更新
    sleep 1
    
    # Step 2: Recall
    log_info "步骤2: 检索记忆 (recall)"
    local recall_resp
    recall_resp=$(api_post "/api/v1/recall" "{
        \"query\": \"集成测试记忆\",
        \"namespace\": \"${NAMESPACE}\",
        \"top_k\": 10
    }")
    
    log_info "recall 响应前200字符: $(echo $recall_resp | head -c 200)..."
    assert_contains "$recall_resp" '"results"' "recall 包含 results 字段"
    assert_contains "$recall_resp" '"count"' "recall 包含 count 字段"
    
    # 尝试从 recall 结果中提取 memory_id
    if [[ -z "$memory_id" ]]; then
        memory_id=$(echo "$recall_resp" | grep -o '"id":"mem_[^"]*"' | head -1 | cut -d'"' -f4)
        log_info "从 recall 结果提取 memory_id: ${memory_id}"
    fi
    
    # Step 3: Feedback
    if [[ -n "$memory_id" ]]; then
        log_info "步骤3: 反馈 (feedback)"
        local feedback_resp
        feedback_resp=$(api_post "/api/v1/feedback" "{
            \"memory_id\": \"${memory_id}\",
            \"useful\": true
        }")
        
        log_info "feedback 响应: $feedback_resp"
        assert_contains "$feedback_resp" '"status"' "feedback 包含状态字段"
    else
        log_warn "无法获取 memory_id，跳过 feedback 测试"
    fi
    
    log_info "commit -> recall -> feedback 链路测试完成"
}

test_navigate() {
    header "测试: 图谱导航 /api/v1/graph/navigate"
    
    # 先创建一条可导航的记忆
    local nav_entity="测试实体_$$"
    api_post "/api/v1/commit" "{
        \"agent_id\": \"${AGENT_ID}\",
        \"namespace\": \"${NAMESPACE}\",
        \"content\": \"${nav_entity} 是一个测试实体，用于图谱导航测试\",
        \"category\": \"test\"
    }" > /dev/null 2>&1
    
    sleep 1
    
    # 执行导航查询
    local nav_resp
    nav_resp=$(api_post "/api/v1/graph/navigate" "{
        \"entity\": \"${nav_entity}\",
        \"max_hops\": 2,
        \"namespace\": \"${NAMESPACE}\"
    }")
    
    log_info "navigate 响应: $(echo $nav_resp | head -c 200)..."
    assert_json_valid "$nav_resp" "navigate 返回有效 JSON"
    
    # 双向导航测试
    local nav_bi_resp
    nav_bi_resp=$(api_post "/api/v1/graph/navigate" "{
        \"source\": \"${nav_entity}\",
        \"target\": \"另一个实体\",
        \"max_hops\": 2,
        \"namespace\": \"${NAMESPACE}\"
    }")
    
    assert_json_valid "$nav_bi_resp" "双向 navigate 返回有效 JSON"
}

test_concurrent() {
    header "测试: 并发请求"
    
    local pids=()
    local outputs=()
    
    # 同时发起 5 个并发请求
    log_info "发起 5 个并发 commit 请求..."
    
    for i in {1..5}; do
        (
            api_post "/api/v1/commit" "{
                \"agent_id\": \"${AGENT_ID}\",
                \"namespace\": \"${NAMESPACE}\",
                \"content\": \"并发测试记忆 ${i} - $(date +%s%N)\",
                \"category\": \"concurrent-test\"
            }" 2>&1
        ) &
        pids+=($!)
    done
    
    # 等待所有请求完成
    local all_ok=true
    for pid in "${pids[@]}"; do
        if ! wait "$pid"; then
            all_ok=false
        fi
    done
    
    TESTS_RUN=$((TESTS_RUN+1))
    if $all_ok; then
        log_pass "5 个并发 commit 请求全部成功"
    else
        log_fail "部分并发请求失败"
    fi
    
    # 测试并发 recall
    log_info "发起 5 个并发 recall 请求..."
    pids=()
    for i in {1..5}; do
        (
            api_post "/api/v1/recall" "{
                \"query\": \"并发测试\",
                \"namespace\": \"${NAMESPACE}\",
                \"top_k\": 5
            }" 2>&1
        ) &
        pids+=($!)
    done
    
    all_ok=true
    for pid in "${pids[@]}"; do
        if ! wait "$pid"; then
            all_ok=false
        fi
    done
    
    TESTS_RUN=$((TESTS_RUN+1))
    if $all_ok; then
        log_pass "5 个并发 recall 请求全部成功"
    else
        log_fail "部分并发 recall 请求失败"
    fi
}

test_batch_commit() {
    header "测试: 批量提交 /api/v1/batch-commit"
    
    local batch_resp
    batch_resp=$(api_post "/api/v1/batch-commit" "{
        \"namespace\": \"${NAMESPACE}\",
        \"items\": [
            {\"content\": \"批量测试项1\", \"tags\": [\"test\", \"batch\"]},
            {\"content\": \"批量测试项2\", \"tags\": [\"test\", \"batch\"]},
            {\"content\": \"批量测试项3\", \"tags\": [\"test\", \"batch\"]}
        ]
    }")
    
    log_info "batch-commit 响应: $(echo $batch_resp | head -c 200)..."
    
    assert_json_valid "$batch_resp" "batch-commit 返回有效 JSON"
}

test_cli_consistency() {
    header "测试: CLI 与 API 一致性"
    
    if [[ ! -x "$CLI" ]]; then
        log_warn "CLI 不可用，跳过 CLI 一致性测试"
        return
    fi
    
    # 测试 CLI stats 与 API stats 一致性
    log_info "比较 CLI stats 与 API stats..."
    
    local cli_stats
    cli_stats=$("$CLI" -url "${API_BASE}" -key "${API_KEY}" -n "${NAMESPACE}" stats 2>&1)
    
    log_info "CLI stats 输出: $cli_stats"
    
    TESTS_RUN=$((TESTS_RUN+1))
    if [[ $? -eq 0 ]]; then
        log_pass "CLI stats 命令执行成功"
    else
        log_fail "CLI stats 命令执行失败"
    fi
    
    # 测试 CLI recall 与 API recall
    log_info "比较 CLI recall 与 API recall..."
    local cli_recall
    cli_recall=$("$CLI" -url "${API_BASE}" -key "${API_KEY}" -n "${NAMESPACE}" recall "测试" 2>&1)
    
    local api_recall
    api_recall=$(api_post "/api/v1/recall" "{
        \"query\": \"测试\",
        \"namespace\": \"${NAMESPACE}\",
        \"top_k\": 10
    }")
    
    log_info "CLI recall 输出前100字符: $(echo $cli_recall | head -c 100)..."
    log_info "API recall 输出前100字符: $(echo $api_recall | head -c 100)..."
    
    # API recall 应返回有效 JSON
    assert_json_valid "$api_recall" "CLI 和 API recall 都返回有效响应"
}

test_recall_debug() {
    header "测试: Recall 诊断接口 /api/v1/recall/debug"
    
    local debug_resp
    debug_resp=$(api_post "/api/v1/recall/debug" "{
        \"query\": \"集成测试\",
        \"namespace\": \"${NAMESPACE}\",
        \"top_k\": 5
    }")
    
    log_info "recall/debug 响应: $(echo $debug_resp | head -c 200)..."
    
    assert_json_valid "$debug_resp" "recall/debug 返回有效 JSON"
    assert_contains "$debug_resp" '"steps"' "recall/debug 包含诊断步骤"
    assert_contains "$debug_resp" '"total_ms"' "recall/debug 包含耗时信息"
}

test_pagerank() {
    header "测试: PageRank /api/v1/graph/pagerank"
    
    local pr_resp
    pr_resp=$(api_get "/api/v1/graph/pagerank")
    
    log_info "pagerank 响应: $(echo $pr_resp | head -c 200)..."
    
    assert_json_valid "$pr_resp" "pagerank 返回有效 JSON"
}

test_distill_quota() {
    header "测试: 蒸馏配额 /api/v1/distill/quota"
    
    local quota_resp
    quota_resp=$(api_get "/api/v1/distill/quota")
    
    log_info "distill/quota 响应: $quota_resp"
    
    assert_json_valid "$quota_resp" "distill/quota 返回有效 JSON"
}

test_distill_status() {
    header "测试: 蒸馏状态 /api/v1/distill/status"
    
    local status_resp
    status_resp=$(api_get "/api/v1/distill/status")
    
    log_info "distill/status 响应: $status_resp"
    
    assert_json_valid "$status_resp" "distill/status 返回有效 JSON"
}

# ── 运行所有测试 ──────────────────────────────────────────────

main() {
    echo ""
    echo "╔════════════════════════════════════════════════════════╗"
    echo "║  织忆 (MemoryWeave) 端到端集成测试                      ║"
    echo "║  API: ${API_BASE}                                  ║"
    echo "╚════════════════════════════════════════════════════════╝"
    
    setup
    
    # 执行所有测试
    test_health
    test_stats
    test_graph_stats
    test_commit_recall_feedback
    test_navigate
    test_concurrent
    test_batch_commit
    test_recall_debug
    test_pagerank
    test_distill_quota
    test_distill_status
    test_cli_consistency
    
    teardown
    
    # 输出测试总结
    echo ""
    echo "══════════════════════════════════════"
    echo "  测试总结"
    echo "══════════════════════════════════════"
    echo -e "  运行: ${TESTS_RUN}"
    echo -e "  通过: ${GREEN}${TESTS_PASSED}${NC}"
    echo -e "  失败: ${RED}${TESTS_FAILED}${NC}"
    echo "══════════════════════════════════════"
    
    if [[ ${TESTS_FAILED} -gt 0 ]]; then
        echo -e "${RED}测试失败${NC}"
        exit 1
    else
        echo -e "${GREEN}全部测试通过${NC}"
        exit 0
    fi
}

# 捕获 EXIT 信号
trap 'teardown' EXIT

main "$@"