# ============================================
# MindSpeed-Ops Docker Image
# NPU Type: Configurable (910b, a3, etc.)
# Supports: x86_64 and aarch64 (ARM)
# Supports: openEuler and Ubuntu
# ============================================

# Global arguments for base image selection
ARG OS=openeuler24.03
ARG BASE_IMAGE_VERSION=9.0.0-beta.2
ARG NPU_TYPE=910b
ARG PYTHON_VERSION=3.11
ARG BASE_IMAGE=""

# ------------------------------
# Stage 1: Base Image
# ------------------------------
FROM ${BASE_IMAGE:-swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:${BASE_IMAGE_VERSION}-${NPU_TYPE}-${OS}-py${PYTHON_VERSION}} AS base

# Switch to root user
USER root

ARG OS_FAMILY=openeuler

# Set default shell to bash
SHELL ["/bin/bash", "-c"]

# Detect architecture (x86_64 or aarch64)
RUN ARCH=$(uname -m) && \
    echo "Detected CPU architecture: ${ARCH}" && \
    if [ "$ARCH" = "x86_64" ]; then \
        echo "Architecture type: x86"; \
    elif [ "$ARCH" = "aarch64" ]; then \
        echo "Architecture type: ARM"; \
    else \
        echo "ERROR: Unsupported architecture: $ARCH"; \
        exit 1; \
    fi

# Configure DNS for stable network connection (using Docker's DNS configuration)
RUN echo "Configuring DNS settings..." && \
    echo "Current DNS configuration:" && \
    if [ -f /etc/resolv.conf ]; then \
        cat /etc/resolv.conf; \
    else \
        echo "No resolv.conf found"; \
    fi && \
    echo "Note: DNS is managed by Docker daemon. To use custom DNS, run container with --dns flag."

# Configure repository based on OS_FAMILY (openeuler/ubuntu)
COPY configure_repo.sh /tmp/configure_repo.sh
RUN chmod +x /tmp/configure_repo.sh && \
    bash /tmp/configure_repo.sh && \
    rm /tmp/configure_repo.sh

# Clean YUM/APT cache and temporary files
RUN if [ "$OS_FAMILY" = "openeuler" ]; then \
        yum clean all && rm -rf /var/cache/yum; \
    elif [ "$OS_FAMILY" = "ubuntu" ]; then \
        apt-get clean && rm -rf /var/lib/apt/lists/*; \
    fi && \
    rm -rf /tmp/* /var/tmp/* /var/log/*

# Install system dependencies (including C++ build toolchain for CMake/pybind11 extensions)
RUN echo "Installing system dependencies..." && \
    if [ "$OS_FAMILY" = "openeuler" ]; then \
        yum install -y \
            git \
            iproute \
            wget \
            curl \
            gcc \
            gcc-c++ \
            make \
            cmake \
            ninja-build \
            autoconf \
            automake \
            libtool \
            pkgconfig \
        && yum clean all; \
    elif [ "$OS_FAMILY" = "ubuntu" ]; then \
        apt-get update && \
        apt-get install -y \
            git \
            iproute2 \
            wget \
            curl \
            gcc \
            g++ \
            make \
            cmake \
            ninja-build \
            autoconf \
            automake \
            libtool \
            pkg-config \
            build-essential \
        && apt-get clean && rm -rf /var/lib/apt/lists/*; \
    else \
        echo "ERROR: Unsupported OS: $OS_FAMILY"; \
        exit 1; \
    fi

# ------------------------------
# Stage 2: Builder (Python env + PyTorch + triton-ascend)
# ------------------------------
FROM base AS builder

ARG TORCH_VERSION=2.7.1
ARG TORCH_NPU_VERSION=2.7.1
ARG TRITON_ASCEND_VERSION=3.2.1

WORKDIR /tmp

# Install Python 3.11 and pip directly (no Conda)
RUN echo "Installing Python 3.11 and pip..." && \
    if [ "$OS_FAMILY" = "openeuler" ]; then \
        # Check available Python packages in openEuler 24.03
        echo "Available Python packages:" && \
        yum list available python3* | grep -E "python3|pip" || true && \
        # Install Python 3.11 and pip (package names may vary)
        yum install -y python3 python3-pip python3-devel && \
        # Verify Python version
        python3 --version && \
        # Create symlinks
        ln -sf /usr/bin/python3 /usr/bin/python && \
        ln -sf /usr/bin/pip3 /usr/bin/pip; \
    elif [ "$OS_FAMILY" = "ubuntu" ]; then \
        apt-get update && \
        apt-get install -y python3.11 python3-pip python3.11-dev && \
        ln -sf /usr/bin/python3.11 /usr/bin/python && \
        ln -sf /usr/bin/pip3 /usr/bin/pip; \
    fi

# Configure pip source
RUN pip config set global.index-url https://repo.huaweicloud.com/repository/pypi/simple && \
    pip config set global.trusted-host "repo.huaweicloud.com"

# Install PyTorch and torch_npu (ONLINE INSTALL ONLY)
RUN echo "Installing PyTorch and torch_npu from PyPI..." && \
        ARCH=$(uname -m) && \
        MAX_RETRIES=3 && \
        for retry in $(seq 1 $MAX_RETRIES); do \
            echo ">>> Attempt $retry of $MAX_RETRIES"; \
            pip cache purge 2>/dev/null || true; \
            rm -rf /root/.cache/pip; \
            if [ "$ARCH" = "x86_64" ]; then \
                echo "Installing PyTorch for x86_64..." && \
                pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu && break; \
            elif [ "$ARCH" = "aarch64" ]; then \
                echo "Installing PyTorch for aarch64..." && \
                pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio && break; \
            fi; \
        done && \
        for retry in $(seq 1 $MAX_RETRIES); do \
            echo ">>> Installing torch-npu, attempt $retry of $MAX_RETRIES"; \
            pip cache purge 2>/dev/null || true; \
            rm -rf /root/.cache/pip; \
            pip install --no-cache-dir torch-npu==${TORCH_NPU_VERSION} && break; \
        done;

# Install triton-ascend (MindSpeed-Ops specific dependency)
RUN echo "Installing triton-ascend..." && \
    MAX_RETRIES=3 && \
    for retry in $(seq 1 $MAX_RETRIES); do \
        echo ">>> Installing triton-ascend, attempt $retry of $MAX_RETRIES"; \
        pip cache purge 2>/dev/null || true; \
        rm -rf /root/.cache/pip; \
        pip install --no-cache-dir triton-ascend==${TRITON_ASCEND_VERSION} --extra-index-url https://triton-ascend.osinfra.cn/pypi/simple && break; \
    done;

# Install Python build dependencies required by MindSpeed-Ops
RUN pip install --no-cache-dir \
    pybind11 \
    pyyaml \
    numpy \
    wheel \
    regex \
    setuptools>=64

# ------------------------------
# Stage 3: Final Runtime Image
# ------------------------------
FROM builder AS final

# Arguments for MindSpeed-Ops
ARG MINDSPEED_OPS_BRANCH=master
ARG SOC_VERSION=""

WORKDIR /workspace

# Disable Git SSL verification (before cloning)
RUN git config --global http.sslVerify false

# Clone MindSpeed-Ops
RUN git clone https://gitcode.com/Ascend/MindSpeed-Ops.git && \
    cd MindSpeed-Ops && \
    if [ -n "${MINDSPEED_OPS_BRANCH}" ] && [ "${MINDSPEED_OPS_BRANCH}" != "master" ]; then \
        git checkout ${MINDSPEED_OPS_BRANCH}; \
    fi

# Install MindSpeed-Ops
# - If SOC_VERSION is set, ACLNN operators will be compiled (requires NPU driver in container)
# - If SOC_VERSION is not set, only Python/Triton operators will be installed (C++ ext skipped)
RUN cd /workspace/MindSpeed-Ops && \
    if [ -n "${SOC_VERSION}" ]; then \
        export SOC_VERSION="${SOC_VERSION}" && \
        echo "Building MindSpeed-Ops with SOC_VERSION=${SOC_VERSION} (ACLNN + C++ extensions enabled)..." && \
        pip install -e . --extra-index-url https://triton-ascend.osinfra.cn/pypi/simple; \
    else \
        echo "Building MindSpeed-Ops without SOC_VERSION (Triton/Python operators only, C++ ext skipped)..." && \
        pip install -e . --extra-index-url https://triton-ascend.osinfra.cn/pypi/simple; \
    fi

# Set working directory
WORKDIR /workspace/MindSpeed-Ops

# Auto-run on login
RUN echo 'cd /workspace/MindSpeed-Ops' >> /root/.bashrc

# Environment variables
ENV ASCEND_TOOLKIT_HOME=/usr/local/Ascend/ascend-toolkit
ENV ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest