#!/bin/bash
# DCU ENVIRONMENT Test Script
# Run with bash `curl -fsSL https://static.ksuser.cn/dcu/env.sh | bash -s --`
# By KsuserKqy
# 2025-10-12

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
BOLD_RED='\033[1;31m'
NC='\033[0m'

echo
echo "DCU ENVIRONMENT For Q3 Test Started ..."
echo -e "Powered by$YELLOW KsuserKqy(20251113620)$NC from SSPU"
echo "------------------------------------------"
# System
if [ -f /etc/os-release ]; then
    source /etc/os-release
    if [[ "$VERSION_ID" == "22.04" ]]; then
        echo -e "[$GREEN √ $NC] $NAME $VERSION"
        if ! command -v snap &> /dev/null; then
            echo -e "[$GREEN √ $NC] $NAME 22.04 (No Snaped)"
        else
            echo -e "[$YELLOW ⚠ $NC] $NAME 22.04 (Snaped)"
        fi
    else
        echo -e "[$RED ⚠ $NC] $NAME $VERSION_ID (Need 22.04)"
    fi
else
    echo -e "[$RED x $NC] Invalid System (Not Ubuntu)"
fi

# Docker
if grep -qE "/docker|/lxc" /proc/1/cgroup 2>/dev/null; then
    echo -e "[$GREEN √ $NC] Docker Installed"
else
    echo -e "[$YELLOW ⚠ $NC] Docker Not Installed"
fi

# DTK
if command -v rocm-smi &> /dev/null; then
    version=$(rocm-smi --version 2>/dev/null | grep -oE '[0-9]+\.[0-9]+\.[0-9]+')
    echo -e "[$GREEN √ $NC] DTK Installed"
    # if [[ "$version" == "25.04.1" ]]; then
    #     echo "[$GREEN √ $NC] 检测到 DTK 版本 $version"
    # else
    #     echo "[$YELLOW ⚠ $NC] 检测到 DTK 版本 $version，非 25.04.1"
    # fi
else
    echo -e "[$RED x $NC] DTK Not Installed"
fi

# GCC
if command -v gcc &> /dev/null; then
    gccver=$(gcc -dumpfullversion -dumpversion 2>/dev/null)
    if [[ "$gccver" == "11.4.0" ]]; then
        echo -e "[$GREEN √ $NC] gcc $gccver"
    else
        echo -e "[$YELLOW ⚠ $NC] gcc $gccver (Need 11.4.0)"
    fi
else
    echo -e "[$RED x $NC] gcc Not Installed"
fi

# Clang
if command -v clang &> /dev/null; then
    clangver=$(clang --version 2>/dev/null | head -n1 | grep -oE '[0-9]+\.[0-9]+\.[0-9]+')
    if [[ "$clangver" == "15.0.0" ]]; then
        echo -e "[$GREEN √ $NC] clang $clangver"
    else
        echo -e "[$YELLOW ⚠ $NC] clang $clangver (15.0.0)"
    fi
else
    echo -e "[$RED x $NC] clang Not Installed"
fi

# Python
if command -v python3 &> /dev/null; then
    pyver=$(python3 -V 2>&1 | awk '{print $2}')
    if [[ "$pyver" == "3.10.18" ]]; then
        echo -e "[$GREEN √ $NC] Python $pyver"
    else
        echo -e "[$YELLOW ⚠ $NC] Python $pyver (Need 3.10.18)"
    fi
else
    echo -e "[$RED x $NC] Python3 Not Installed"
fi

# Pytorch
if python3 -c "import torch" &> /dev/null; then
    torchver=$(python3 -c "import torch; print(torch.__version__)" 2>/dev/null)
    if [[ "$torchver" == "2.4.1" ]]; then
        echo -e "[$GREEN √ $NC] PyTorch $torchver"
    else
        echo -e "[$YELLOW ⚠ $NC] PyTorch $torchver (Need 2.4.1)"
    fi
else
    echo -e "[$RED x $NC] PyTorch Not Installed"
fi

# Conda
if command -v conda &>/dev/null; then
    CONDA_BASE=$(conda info --base 2>/dev/null)
    if [ -f "$CONDA_BASE/etc/profile.d/conda.sh" ]; then
        source "$CONDA_BASE/etc/profile.d/conda.sh"
        echo -e "[$GREEN √ $NC] Conda Initialized"
    else
        echo -e "[$RED x $NC] Could Not Find conda.sh For Initialization"
        exit 1
    fi
    if conda env list | grep -qE "^dcu-q3\s"; then
        echo -e "[$GREEN √ $NC] Conda environment 'dcu-q3' already exists"
    else
        echo -e "[$BLUE * $NC] Creating conda environment 'dcu-q3' with Python 3.10.18..."
        conda create -n dcu-q3 python=3.10.18 -y &>/dev/null
        echo -e "[$GREEN √ $NC] Environment 'dcu-q3' created successfully"
    fi
    mkdir -p ./venv
    ACTIVATE_URL="https://static.ksuser.cn/dcu/activate-q3.sh"
    echo -e "[$BLUE * $NC] Downloading activation script from: $ACTIVATE_URL"
    if command -v curl &>/dev/null; then
        curl -fsSL "$ACTIVATE_URL" -o ./venv/activate-q3.sh
    elif command -v wget &>/dev/null; then
        wget -q "$ACTIVATE_URL" -O ./venv/activate-q3.sh
    else
        echo -e "[$RED x $NC] Neither curl nor wget is available to download activate-q3.sh."
        exit 1
    fi
    chmod +x ./venv/activate-q3.sh
    echo -e "[$GREEN √ $NC] Activation script saved to ./venv/activate-q3.sh"
    echo -e "[$YELLOW * $NC] 第三题环境已初始化完成，输入以下命令以启用虚拟环境并安装相关依赖:"
    echo -e "    $YELLOW source ./venv/activate-q3.sh$NC"
else
    echo -e "[$RED x $NC] Conda Not Found."
    exit 1
fi
echo "------------------------------------------"
echo "DCU ENVIRONMENT Test Completed."
echo