#!/bin/bash
set -e

LOG_DIR=/data1/LMCache/logs/plugin_test
mkdir -p $LOG_DIR /tmp/test_fs_native_l2
rm -rf /tmp/test_fs_native_l2/*

echo "=== Step 1: Start LMCache MP Server (fs_native with read_ahead_size=4096) ==="
L2_JSON='{"type": "fs_native", "base_path": "/tmp/test_fs_native_l2", "num_workers": 4, "read_ahead_size": 4096}'
export LMCACHE_LOG_LEVEL=DEBUG
/usr/bin/python3 -m lmcache.v1.multiprocess.server \
    --host localhost --port 15556 \
    --chunk-size 256 --l1-size-gb 0.03 \
    --eviction-policy LRU --max-workers 1 \
    --l2-adapter "$L2_JSON" \
    > $LOG_DIR/lmcache_server.log 2>&1 &
LMCACHE_PID=$!
echo "LMCache PID: $LMCACHE_PID"
sleep 5

if ! kill -0 $LMCACHE_PID 2>/dev/null; then
    echo "FAIL: LMCache server died. Log:"
    cat $LOG_DIR/lmcache_server.log
    exit 1
fi
echo "LMCache server is running."

echo "=== Step 2: Start vLLM ==="
MODEL_PATH="/data1/model/DeepSeek-V2-Lite-Chat/"
KV_CFG='{"kv_connector":"LMCacheMPConnector","kv_role":"kv_both","kv_connector_extra_config":{"lmcache.mp.port":15556}}'
export CUDA_VISIBLE_DEVICES=0,1

/usr/bin/python3 -m vllm.entrypoints.cli.main serve $MODEL_PATH \
    -tp 2 \
    --load-format dummy \
    --trust-remote-code \
    --served-model-name vllm_cpu_offload \
    --gpu_memory_utilization 0.85 \
    --max-num-seqs 64 \
    --no-enable-prefix-caching --enforce-eager --max-model-len 8192 \
    --port 8001 \
    --disable-log-requests \
    --kv-transfer-config "$KV_CFG" \
    > $LOG_DIR/vllm_server.log 2>&1 &
VLLM_PID=$!
echo "vLLM PID: $VLLM_PID"
echo "Waiting 60s for vLLM to load model..."
sleep 60

if ! kill -0 $VLLM_PID 2>/dev/null; then
    echo "FAIL: vLLM server died. Last 20 lines:"
    tail -20 $LOG_DIR/vllm_server.log
    exit 1
fi
echo "vLLM server is running."

echo "=== Step 3: Send test requests ==="
LONG_PROMPT="The history of artificial intelligence (AI) began in antiquity, with myths, stories and rumors of artificial beings endowed with intelligence or consciousness by master craftsmen. The seeds of modern AI were planted by philosophers who attempted to describe the process of human thinking as the mechanical manipulation of symbols. This work culminated in the invention of the programmable digital computer in the 1940s, a machine based on the abstract essence of mathematical reasoning. This device and the ideas behind it inspired a handful of scientists to begin seriously discussing the possibility of building an electronic brain. The field of AI research was founded at a workshop held on the campus of Dartmouth College, USA during the summer of 1956. Those who attended would become the leaders of AI research for decades. Many of them predicted that a machine as intelligent as a human being would exist in no more than a generation, and they were given millions of dollars to make this vision come true. Eventually, it became obvious that commercial developers and researchers had grossly underestimated the difficulty of the project. In 1974, in response to the criticism from James Lighthill and ongoing pressure from congress, the U.S. and British Governments stopped funding undirected research into artificial intelligence, and the difficult years that followed became known as an AI winter. Seven years later, a visionary initiative by the Japanese Government inspired governments and industry to provide AI with billions of dollars, but by the late 1980s the investors became disillusioned and withdrew funding again. Investment and interest in AI boomed in the first decades of the 21st century when machine learning was successfully applied to many problems in academia and industry due to the availability of large amounts of data and fast computers. The achievements of deep learning in the 2010s, particularly breakthroughs in areas like image recognition, natural language processing, and game playing, led to a renewed surge of interest in AI technologies. Transformer architectures, introduced in 2017, revolutionized the field and enabled the development of large language models that could generate human-like text, answer questions, write code, and engage in sophisticated reasoning tasks that were previously thought to be exclusively human capabilities."

echo "--- Request 1 (store to L2) ---"
RESP1=$(curl -s http://localhost:8001/v1/completions \
    -H "Content-Type: application/json" \
    -d "{\"model\": \"vllm_cpu_offload\", \"prompt\": \"$LONG_PROMPT\", \"max_tokens\": 10, \"temperature\": 0}")
echo "$RESP1" | python3 -m json.tool
sleep 3

echo ""
echo "--- L2 files after request 1 ---"
ls -la /tmp/test_fs_native_l2/ 2>/dev/null || echo "(empty)"
FILE_COUNT=$(ls /tmp/test_fs_native_l2/*.data 2>/dev/null | wc -l)
echo "Data files: $FILE_COUNT"

echo ""
echo "--- Request 2 (should load from L2/L1) ---"
RESP2=$(curl -s http://localhost:8001/v1/completions \
    -H "Content-Type: application/json" \
    -d "{\"model\": \"vllm_cpu_offload\", \"prompt\": \"$LONG_PROMPT\", \"max_tokens\": 10, \"temperature\": 0}")
echo "$RESP2" | python3 -m json.tool
sleep 2

echo ""
echo "=== Step 4: Check LMCache logs ==="
echo "--- FS native adapter creation ---"
grep -i "Created FS native" $LOG_DIR/lmcache_server.log || echo "(not found)"
echo "--- read_ahead in logs ---"
grep -i "read_ahead" $LOG_DIR/lmcache_server.log || echo "(not found)"
echo "--- L2 store/load operations ---"
grep -iE "(stored|loaded|L2|l2_store|l2_load)" $LOG_DIR/lmcache_server.log | tail -20 || echo "(not found)"

echo ""
echo "=== Step 5: Cleanup ==="
kill $VLLM_PID 2>/dev/null || true
kill $LMCACHE_PID 2>/dev/null || true
sleep 2

echo ""
echo "=== TEST COMPLETE ==="
echo "Data files in /tmp/test_fs_native_l2:"
ls -lh /tmp/test_fs_native_l2/*.data 2>/dev/null || echo "(none)"