#!/bin/bash
# 打包 nebula step3_4_merge agent 相关代码和数据

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
PACKAGE_NAME="nebula_step3_4_merge_agent_$(date +%Y%m%d_%H%M%S)"
PACKAGE_DIR="/tmp/$PACKAGE_NAME"
OUTPUT_FILE="$PROJECT_ROOT/${PACKAGE_NAME}.tar.gz"

echo "=========================================="
echo "打包 Nebula Step3-4-Merge Agent"
echo "=========================================="
echo ""

# 创建临时目录
mkdir -p "$PACKAGE_DIR"
cd "$PACKAGE_DIR"

# 创建目录结构
mkdir -p interest_tree_extender/src/agents/nebula
mkdir -p interest_tree_extender/src/common
mkdir -p interest_tree_extender/src/orchestrator
mkdir -p interest_tree_extender/prompts
mkdir -p interest_tree_extender/config
mkdir -p interest_tree_extender/output/nebula/step2_content_search
mkdir -p interest_tree_extender/scripts
mkdir -p scripts

echo "1. 复制核心代码文件..."

# 核心 agent 代码
cp "$PROJECT_ROOT/interest_tree_extender/src/agents/nebula/step3_4_merge_interest_builder.py" \
   "$PACKAGE_DIR/interest_tree_extender/src/agents/nebula/"

cp "$PROJECT_ROOT/interest_tree_extender/src/agents/nebula/llm_client_wrapper.py" \
   "$PACKAGE_DIR/interest_tree_extender/src/agents/nebula/"

cp "$PROJECT_ROOT/interest_tree_extender/src/agents/nebula/__init__.py" \
   "$PACKAGE_DIR/interest_tree_extender/src/agents/nebula/" 2>/dev/null || true

# 依赖的通用模块
cp "$PROJECT_ROOT/interest_tree_extender/src/common/storage.py" \
   "$PACKAGE_DIR/interest_tree_extender/src/common/"

cp "$PROJECT_ROOT/interest_tree_extender/src/common/types.py" \
   "$PACKAGE_DIR/interest_tree_extender/src/common/" 2>/dev/null || true

cp "$PROJECT_ROOT/interest_tree_extender/src/common/__init__.py" \
   "$PACKAGE_DIR/interest_tree_extender/src/common/" 2>/dev/null || true

# CLI 入口（只复制相关部分，或者整个文件）
cp "$PROJECT_ROOT/interest_tree_extender/src/cli.py" \
   "$PACKAGE_DIR/interest_tree_extender/src/"

cp "$PROJECT_ROOT/interest_tree_extender/src/prompt_loader.py" \
   "$PACKAGE_DIR/interest_tree_extender/src/" 2>/dev/null || true

# 创建 src/__init__.py
touch "$PACKAGE_DIR/interest_tree_extender/src/__init__.py"

echo "2. 复制 Prompt 模板文件..."

# Prompt 文件
cp "$PROJECT_ROOT/interest_tree_extender/prompts/step3_4_merge_final_merger.md" \
   "$PACKAGE_DIR/interest_tree_extender/prompts/" 2>/dev/null || true

cp "$PROJECT_ROOT/interest_tree_extender/prompts/step3_4_merge_batch_extractor_delete.md" \
   "$PACKAGE_DIR/interest_tree_extender/prompts/" 2>/dev/null || true

echo "3. 复制配置文件..."

# 配置文件
if [ -f "$PROJECT_ROOT/interest_tree_extender/config/nebula_config.json" ]; then
    cp "$PROJECT_ROOT/interest_tree_extender/config/nebula_config.json" \
       "$PACKAGE_DIR/interest_tree_extender/config/"
fi

if [ -f "$PROJECT_ROOT/interest_tree_extender/config/nebula_config.example.json" ]; then
    cp "$PROJECT_ROOT/interest_tree_extender/config/nebula_config.example.json" \
       "$PACKAGE_DIR/interest_tree_extender/config/"
fi

echo "4. 复制所有 filtered.jsonl 数据文件..."

# 复制所有 filtered.jsonl 文件
DATA_DIR="$PROJECT_ROOT/interest_tree_extender/output/nebula/step2_content_search"
if [ -d "$DATA_DIR" ]; then
    FILTERED_FILES=$(find "$DATA_DIR" -name "*filtered*.jsonl" -type f)
    if [ -n "$FILTERED_FILES" ]; then
        FILE_COUNT=0
        TOTAL_SIZE=0
        for FILE in $FILTERED_FILES; do
            if [ -f "$FILE" ]; then
                cp "$FILE" "$PACKAGE_DIR/interest_tree_extender/output/nebula/step2_content_search/"
                FILE_SIZE=$(stat -f%z "$FILE" 2>/dev/null || stat -c%s "$FILE" 2>/dev/null || echo "0")
                TOTAL_SIZE=$((TOTAL_SIZE + FILE_SIZE))
                FILE_COUNT=$((FILE_COUNT + 1))
                echo "  已复制: $(basename "$FILE") ($(du -h "$FILE" | cut -f1))"
            fi
        done
        echo "  总计: $FILE_COUNT 个文件，总大小: $(echo "scale=2; $TOTAL_SIZE/1024/1024" | bc)MB"
    else
        echo "  警告: 未找到 filtered.jsonl 文件"
    fi
else
    echo "  警告: 数据目录不存在: $DATA_DIR"
fi

echo "5. 复制相关脚本..."

# 复制 qw.py（LLM客户端）
cp "$PROJECT_ROOT/scripts/qw.py" \
   "$PACKAGE_DIR/scripts/" 2>/dev/null || true

# 复制网络监控脚本
cp "$PROJECT_ROOT/scripts/network_monitor.py" \
   "$PACKAGE_DIR/scripts/" 2>/dev/null || true

echo "6. 创建 README 和使用说明..."

cat > "$PACKAGE_DIR/README.md" << 'EOF'
# Nebula Step3-4-Merge Agent 部署包

## 文件结构

```
interest_tree_extender/
├── src/
│   ├── agents/
│   │   └── nebula/
│   │       ├── step3_4_merge_interest_builder.py  # 主代码
│   │       └── llm_client_wrapper.py               # LLM客户端包装
│   ├── common/
│   │   └── storage.py                              # 存储工具
│   └── cli.py                                      # CLI入口
├── prompts/
│   ├── step3_4_merge_final_merger.md              # 合并Prompt
│   └── step3_4_merge_batch_extractor_delete.md   # 批次抽取Prompt
├── config/
│   └── nebula_config.json                         # 配置文件
└── output/
    └── nebula/
        └── step2_content_search/
            └── external_corpus_sample.jsonl        # 示例数据
scripts/
├── qw.py                                          # LLM API客户端
└── network_monitor.py                             # 网络监控工具
```

## 使用方法

### 1. 安装依赖

```bash
pip install dashscope volcenginesdkarkruntime psutil
```

### 2. 配置 API Key

编辑 `scripts/qw.py`，配置你的 API Key：
- 千问 API Key
- 豆包 API Key（可选）

### 3. 准备数据

将你的 step2 输出数据文件放在：
```
interest_tree_extender/output/nebula/step2_content_search/external_corpus_*.jsonl
```

### 4. 运行

```bash
cd interest_tree_extender
python -m src.cli nebula-step3-4-merge \
    --corpus-path ./output/nebula/step2_content_search/external_corpus_sample.jsonl \
    --output ./output \
    --api-key YOUR_API_KEY \
    --llm-model qwen-plus \
    --num-workers 3 \
    --rate-limit 1.0
```

### 5. 查看输出

输出文件位于：
```
output/nebula/step3_4_merge/merged_interest_trees.jsonl
output/nebula/step3_4_merge/build_stats.json
```

## 参数说明

- `--corpus-path`: step2 输出的 JSONL 文件路径
- `--output`: 输出目录
- `--api-key`: LLM API Key
- `--llm-model`: 模型名称（qwen-plus 或 doubao）
- `--num-workers`: 并发线程数（默认3）
- `--rate-limit`: API调用频率限制（秒，默认1.0）
- `--max-chars-per-doc`: 每篇文档最大字符数（默认2000）
- `--max-chars-per-chunk`: 每个chunk最大字符数（默认30000）
- `--limit`: 限制处理的from_query数量（用于测试）

## 注意事项

1. 确保网络连接正常，API调用可能需要较长时间
2. 建议先用小数据集测试（使用 --limit 参数）
3. 监控网络和CPU使用情况，避免资源耗尽
4. 如果遇到连接错误，会自动重试，但建议检查网络状况
EOF

cat > "$PACKAGE_DIR/QUICK_START.md" << 'EOF'
# 快速开始

## 1. 解压

```bash
tar -xzf nebula_step3_4_merge_agent_*.tar.gz
cd nebula_step3_4_merge_agent_*
```

## 2. 安装依赖

```bash
pip install dashscope volcenginesdkarkruntime psutil
```

## 3. 配置 API Key

编辑 `scripts/qw.py`，填入你的 API Key。

## 4. 运行测试

```bash
cd interest_tree_extender
python -m src.cli nebula-step3-4-merge \
    --corpus-path ./output/nebula/step2_content_search/external_corpus_sample.jsonl \
    --output ./output \
    --api-key YOUR_API_KEY \
    --llm-model qwen-plus \
    --num-workers 2 \
    --limit 5
```

## 5. 查看结果

```bash
cat output/nebula/step3_4_merge/merged_interest_trees.jsonl | head -1 | jq .
cat output/nebula/step3_4_merge/build_stats.json | jq .
```
EOF

echo "7. 创建 requirements.txt..."

cat > "$PACKAGE_DIR/requirements.txt" << 'EOF'
dashscope>=1.10.0
volcenginesdkarkruntime>=1.0.0
psutil>=5.9.0
EOF

echo "8. 创建打包清单..."

cat > "$PACKAGE_DIR/MANIFEST.txt" << EOF
打包时间: $(date)
打包内容:
- Nebula Step3-4-Merge Agent 核心代码
- LLM 客户端包装
- Prompt 模板
- 配置文件
- 示例数据（100行）
- 相关工具脚本

文件列表:
EOF

find . -type f -not -path '*/\.*' | sort >> "$PACKAGE_DIR/MANIFEST.txt"

echo ""
echo "9. 创建压缩包..."

cd /tmp
tar -czf "$OUTPUT_FILE" "$PACKAGE_NAME"

# 显示文件大小
FILE_SIZE=$(du -h "$OUTPUT_FILE" | cut -f1)
echo ""
echo "=========================================="
echo "打包完成！"
echo "=========================================="
echo "压缩包: $OUTPUT_FILE"
echo "大小: $FILE_SIZE"
echo ""
echo "文件清单:"
tar -tzf "$OUTPUT_FILE" | head -20
echo "..."
echo ""
echo "解压命令:"
echo "  tar -xzf $(basename "$OUTPUT_FILE")"
echo ""

# 清理临时目录
rm -rf "$PACKAGE_DIR"

echo "临时文件已清理"

