#!/bin/bash
# Interest Tree Extender 打包脚本
# 用法: bash pack.sh [选项]

set -e

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

# 基础路径
BASE_DIR="/data/interest_tag/兴趣构建"
PROJECT="interest_tree_extender"
DATE=$(date +%Y%m%d_%H%M%S)

echo -e "${GREEN}========================================${NC}"
echo -e "${GREEN}Interest Tree Extender 打包工具${NC}"
echo -e "${GREEN}========================================${NC}"
echo ""

# 显示菜单
show_menu() {
    echo "请选择打包方案:"
    echo ""
    echo "  1) 最小代码包 (推荐) - 只打包核心代码和配置"
    echo "     大小: < 5MB"
    echo "     包含: src/ config/ prompts/ docs/ scripts/ *.md"
    echo ""
    echo "  2) 代码 + 小型结果 - 包含 JSON 统计文件"
    echo "     大小: < 10MB"
    echo "     包含: 方案1 + output/**/*.json (排除 .jsonl .pkl)"
    echo ""
    echo "  3) 只打包 src/ 源代码"
    echo "     大小: ~1.1GB"
    echo "     包含: 只有 src/ 目录"
    echo ""
    echo "  4) 完整打包 (慎用) - 包含所有数据"
    echo "     大小: ~13GB"
    echo "     包含: 整个项目（包括 output/ cache/ 等）"
    echo ""
    echo "  5) 自定义 - 手动指定要打包的目录"
    echo ""
    echo "  0) 退出"
    echo ""
    echo -n "请输入选项 [0-5]: "
}

# 方案1：最小代码包
pack_minimal() {
    echo -e "\n${YELLOW}正在执行方案1：最小代码包...${NC}\n"
    
    OUTPUT_FILE="${BASE_DIR}/${PROJECT}_code_${DATE}.tar.gz"
    
    cd "$BASE_DIR"
    tar -czf "$OUTPUT_FILE" \
        ${PROJECT}/src \
        ${PROJECT}/config \
        ${PROJECT}/prompts \
        ${PROJECT}/docs \
        ${PROJECT}/scripts \
        ${PROJECT}/USAGE_GUIDE.md \
        ${PROJECT}/PACKAGING_GUIDE.md \
        ${PROJECT}/FILE_LIST.md \
        ${PROJECT}/analyze_pipeline.py \
        ${PROJECT}/test_step5_merged_trees.py \
        2>/dev/null || true
    
    echo -e "${GREEN}✓ 打包完成！${NC}"
    echo -e "文件: ${OUTPUT_FILE}"
    ls -lh "$OUTPUT_FILE"
}

# 方案2：代码 + 小型结果
pack_with_results() {
    echo -e "\n${YELLOW}正在执行方案2：代码 + 小型结果...${NC}\n"
    
    OUTPUT_FILE="${BASE_DIR}/${PROJECT}_with_results_${DATE}.tar.gz"
    
    cd "$BASE_DIR"
    tar -czf "$OUTPUT_FILE" \
        ${PROJECT}/src \
        ${PROJECT}/config \
        ${PROJECT}/prompts \
        ${PROJECT}/docs \
        ${PROJECT}/scripts \
        ${PROJECT}/USAGE_GUIDE.md \
        ${PROJECT}/PACKAGING_GUIDE.md \
        ${PROJECT}/FILE_LIST.md \
        ${PROJECT}/analyze_pipeline.py \
        ${PROJECT}/test_step5_merged_trees.py \
        --exclude='*.jsonl' \
        --exclude='*.pkl' \
        --exclude='*.log' \
        --exclude='__pycache__' \
        ${PROJECT}/output \
        2>/dev/null || true
    
    echo -e "${GREEN}✓ 打包完成！${NC}"
    echo -e "文件: ${OUTPUT_FILE}"
    ls -lh "$OUTPUT_FILE"
}

# 方案3：只打包 src/
pack_src_only() {
    echo -e "\n${YELLOW}正在执行方案3：只打包 src/ 源代码...${NC}\n"
    
    OUTPUT_FILE="${BASE_DIR}/${PROJECT}_src_only_${DATE}.tar.gz"
    
    cd "$BASE_DIR"
    tar -czf "$OUTPUT_FILE" \
        ${PROJECT}/src \
        2>/dev/null || true
    
    echo -e "${GREEN}✓ 打包完成！${NC}"
    echo -e "文件: ${OUTPUT_FILE}"
    ls -lh "$OUTPUT_FILE"
}

# 方案4：完整打包
pack_full() {
    echo -e "\n${RED}警告：完整打包将生成约 13GB 的文件！${NC}"
    echo -n "确认继续? (yes/no): "
    read confirm
    
    if [ "$confirm" != "yes" ]; then
        echo "已取消。"
        return
    fi
    
    echo -e "\n${YELLOW}正在执行方案4：完整打包（这可能需要几分钟）...${NC}\n"
    
    OUTPUT_FILE="${BASE_DIR}/${PROJECT}_full_${DATE}.tar.gz"
    
    cd "$BASE_DIR"
    tar -czf "$OUTPUT_FILE" \
        --exclude='${PROJECT}/**/__pycache__' \
        --exclude='${PROJECT}/**/*.pyc' \
        --exclude='${PROJECT}/.git' \
        --exclude='${PROJECT}/nohup.out' \
        ${PROJECT}/ \
        2>/dev/null || true
    
    echo -e "${GREEN}✓ 打包完成！${NC}"
    echo -e "文件: ${OUTPUT_FILE}"
    ls -lh "$OUTPUT_FILE"
}

# 方案5：自定义
pack_custom() {
    echo -e "\n${YELLOW}自定义打包${NC}\n"
    echo "请输入要打包的目录/文件（相对于 ${PROJECT}/），用空格分隔："
    echo "示例: src config prompts docs"
    echo ""
    echo -n "输入: "
    read -a custom_paths
    
    if [ ${#custom_paths[@]} -eq 0 ]; then
        echo -e "${RED}未输入任何路径，已取消。${NC}"
        return
    fi
    
    OUTPUT_FILE="${BASE_DIR}/${PROJECT}_custom_${DATE}.tar.gz"
    
    echo -e "\n${YELLOW}将打包以下内容:${NC}"
    for path in "${custom_paths[@]}"; do
        echo "  - ${PROJECT}/${path}"
    done
    echo ""
    echo -n "确认继续? (yes/no): "
    read confirm
    
    if [ "$confirm" != "yes" ]; then
        echo "已取消。"
        return
    fi
    
    cd "$BASE_DIR"
    
    # 构建 tar 命令
    tar_paths=""
    for path in "${custom_paths[@]}"; do
        tar_paths="$tar_paths ${PROJECT}/${path}"
    done
    
    tar -czf "$OUTPUT_FILE" $tar_paths 2>/dev/null || true
    
    echo -e "${GREEN}✓ 打包完成！${NC}"
    echo -e "文件: ${OUTPUT_FILE}"
    ls -lh "$OUTPUT_FILE"
}

# 主程序
main() {
    while true; do
        show_menu
        read choice
        
        case $choice in
            1)
                pack_minimal
                break
                ;;
            2)
                pack_with_results
                break
                ;;
            3)
                pack_src_only
                break
                ;;
            4)
                pack_full
                break
                ;;
            5)
                pack_custom
                break
                ;;
            0)
                echo "退出。"
                exit 0
                ;;
            *)
                echo -e "${RED}无效选项，请重试。${NC}\n"
                ;;
        esac
    done
    
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}打包完成！${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "生成的文件位于: ${BASE_DIR}/"
    echo ""
    echo "查看压缩包内容:"
    echo "  tar -tzf [文件名].tar.gz | less"
    echo ""
    echo "解压:"
    echo "  tar -xzf [文件名].tar.gz"
}

# 运行主程序
main

