#!/bin/bash

################################################################################
# Kimodo 环境初始化（完全独立版本）
################################################################################

set -e

WORKSPACE_DIR="/workspace"
ENV_DIR="${WORKSPACE_DIR}/env"
PYTHON_BIN="${ENV_DIR}/kimodo/bin/python"

# 颜色输出
print_info() {
    echo -e "\033[32m[INFO]\033[0m $1"
}

print_error() {
    echo -e "\033[31m[ERROR]\033[0m $1"
}

print_step() {
    echo -e "\033[36m[STEP]\033[0m $1"
}

# 检查环境
check_env() {
    if [ ! -f "${PYTHON_BIN}" ]; then
        print_error "Python 环境未找到: ${PYTHON_BIN}"
        exit 1
    fi

    print_info "Python 环境: ${PYTHON_BIN}"
    print_info "Python 版本: $(${PYTHON_BIN} --version)"
}

# 检查系统工具
check_system_tools() {
    print_step "检查系统工具..."

    local missing=()

    for tool in gcc g++ make cmake tar gzip bash; do
        if ! command -v $tool &> /dev/null; then
            missing+=($tool)
        fi
    done

    if [ ${#missing[@]} -gt 0 ]; then
        print_info "缺少系统工具: ${missing[*]}"
        print_info "安装命令: sudo apt-get install -y build-essential cmake tar gzip"
        echo ""
        print_info "注意：这些工具需要目标机器已安装"
    else
        print_info "所有系统工具已安装"
    fi

    # 检查 CUDA
    if command -v nvidia-smi &> /dev/null; then
        print_info "检测到 CUDA: $(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
    else
        print_info "未检测到 CUDA，将使用 CPU 模式"
    fi
}

# 显示目录结构
show_structure() {
    print_step "目录结构:"

    echo ""
    echo "${WORKSPACE_DIR}/"
    echo "├── env/kimodo/          # Python 环境 (~9GB)"
    echo "│   ├── bin/             # Python 和相关命令"
    echo "│   ├── lib/             # Python 库"
    echo "│   └── ..."
    echo "├── kimodo/              # 项目代码"
    echo "├── deps/                # 依赖信息"
    echo "└── start_standalone.sh  # 启动脚本"
    echo ""
}

# 验证环境
verify_env() {
    print_step "验证环境..."

    ${PYTHON_BIN} << VERIFY_EOF
import sys
print(f"Python: {sys.version}")

try:
    import torch
    print(f"✓ PyTorch: {torch.__version__}")
    print(f"✓ CUDA: {'可用' if torch.cuda.is_available() else '不可用'}")
except ImportError as e:
    print(f"✗ PyTorch: {e}")

try:
    import kimodo
    print(f"✓ Kimodo: 导入成功")
except ImportError as e:
    print(f"✗ Kimodo: {e}")

print("\n验证完成")
VERIFY_EOF
}

# 主函数
main() {
    echo ""
    echo "======================================================================"
    echo "              Kimodo 环境初始化"
    echo "======================================================================"
    echo ""

    check_env
    check_system_tools
    show_structure
    verify_env

    echo ""
    print_info "======================================================================"
    print_info "环境初始化完成！"
    print_info ""
    print_info "启动项目: ./start_standalone.sh"
    print_info "======================================================================"
    echo ""
}

main