#!/bin/bash

################################################################################
# Kimodo 环境初始化脚本
# 
# 用于在新机器上初始化环境
################################################################################

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WORKSPACE_DIR="$(dirname "$(dirname "${SCRIPT_DIR}"))")"
ENV_DIR="${WORKSPACE_DIR}/env"
KIMODO_DIR="${WORKSPACE_DIR}/kimodo"
PYTHON_BIN="${ENV_DIR}/conda_env/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"
}

print_warning() {
    echo -e "\033[33m[WARNING]\033[0m $1"
}

# 检查系统依赖
check_system_deps() {
    print_step "检查系统依赖..."
    
    local missing=()
    
    if ! command -v gcc &> /dev/null; then missing+=("gcc"); fi
    if ! command -v g++ &> /dev/null; then missing+=("g++"); fi
    if ! command -v cmake &> /dev/null; then missing+=("cmake"); fi
    if ! command -v tar &> /dev/null; then missing+=("tar"); fi
    if ! command -v gzip &> /dev/null; then missing+=("gzip"); fi
    
    if [ ${#missing[@]} -gt 0 ]; then
        print_warning "缺少系统工具: ${missing[*]}"
        print_warning "请参考 ${ENV_DIR}/configs/BINARY_DEPS.txt 安装这些工具"
        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_warning "未检测到 CUDA，将使用 CPU 模式（性能较低）"
    fi
}

# 检查环境是否已解压
check_env() {
    if [ ! -f "${PYTHON_BIN}" ]; then
        print_error "Python 环境未找到: ${PYTHON_BIN}"
        print_info "请先运行: ${ENV_DIR}/scripts/unpack_env.sh"
        exit 1
    fi
    
    print_info "Python 环境: ${PYTHON_BIN}"
    print_info "Python 版本: $(${PYTHON_BIN} --version)"
}

# 修复环境路径（如果移动了环境）
fix_env_paths() {
    print_step "修复环境路径..."
    
    # 激活 Python 并修复硬编码路径
    ${PYTHON_BIN} << PYTHON_EOF
import os
import sys
import re
from pathlib import Path

def fix_bin_shebangs(bin_dir):
    """修复二进制文件的 shebang"""
    for file in Path(bin_dir).iterdir():
        if file.is_file() and os.access(file, os.X_OK):
            try:
                with open(file, 'rb') as f:
                    content = f.read()
                
                # 检查是否是文本文件
                if b'\x00' in content[:1024]:
                    continue
                
                try:
                    text = content.decode('utf-8')
                    original = text
                    
                    # 替换旧的 shebang
                    if text.startswith('#!'):
                        lines = text.split('\n', 1)
                        shebang = lines[0]
                        
                        # 替换为当前环境路径
                        new_shebang = re.sub(
                            r'#!.*/(bin/python|envs/[^/]+/bin/python)',
                            f'#!{sys.executable}',
                            shebang
                        )
                        
                        if shebang != new_shebang:
                            lines[0] = new_shebang
                            text = '\n'.join(lines)
                            
                            if text != original:
                                with open(file, 'w') as f:
                                    f.write(text)
                                print(f"Fixed: {file.name}")
                except:
                    pass
            except:
                pass

# 修复 bin 目录
bin_dir = Path(sys.executable).parent
fix_bin_shebangs(bin_dir)
print("环境路径修复完成")
PYTHON_EOF
}

# 验证环境
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 available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"CUDA version: {torch.version.cuda}")
except ImportError as e:
    print(f"PyTorch import error: {e}")

try:
    import transformers
    print(f"Transformers: {transformers.__version__}")
except ImportError as e:
    print(f"Transformers import error: {e}")

try:
    import kimodo
    print(f"Kimodo: imported successfully")
except ImportError as e:
    print(f"Kimodo import error: {e}")

print("环境验证完成")
VERIFY_EOF
}

# 主函数
main() {
    echo ""
    echo "======================================================================"
    echo "                  Kimodo 环境初始化"
    echo "======================================================================"
    echo ""
    
    check_system_deps
    check_env
    fix_env_paths
    verify_env
    
    echo ""
    print_info "======================================================================"
    print_info "环境初始化完成！"
    print_info "运行启动命令: cd ${WORKSPACE_DIR} && ./start_local.sh"
    print_info "======================================================================"
    echo ""
}

main