#!/bin/bash
set -uo pipefail
set +H

export PIP_NO_CACHE_DIR=1
export PIP_DISABLE_PIP_VERSION_CHECK=1
export GIT_TERMINAL_PROMPT=0

source /venv/main/bin/activate
COMFYUI_DIR="${WORKSPACE:-/workspace}/ComfyUI"
if [ -f "/venv/main/bin/python" ]; then
    PYTHON_BIN="/venv/main/bin/python"
    PIP_BIN="/venv/main/bin/pip"
else
    PYTHON_BIN="python3"
    PIP_BIN="pip3"
fi

APT_PACKAGES=(
    "aria2"
)
PIP_PACKAGES=(
)
NODES=(
    "https://github.com/ltdrdata/ComfyUI-Impact-Pack.git"
    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack.git"
    "https://github.com/cubiq/ComfyUI_IPAdapter_plus.git"
    "https://github.com/ka-puna/comfyui-yanc.git"
)
CHECKPOINT_MODELS=(
    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/checkpoints/waiNSFWIllustrious_v150.safetensors"
)
CLIP_VISION_MODELS=(
    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/clip_vision/CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors"
)
IPADAPTER_MODELS=(
    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/ipadapter/ip-adapter-plus-face_sdxl_vit-h.safetensors"
)
ESRGAN_MODELS=(
    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/upscale_models/RealESRGAN_x4plus_anime_6B.pth"
)
WORKFLOWS=(
    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/Single_Posture_ko.json"
)

log_info() { echo "--> $1"; }
log_warn() { echo "    ⚠️  $1"; }
log_error() { echo "    ❌ $1"; }
log_success() { echo "    ✅ $1"; }
log_step() { echo ""; echo "=== [Step $1] $2 ==="; }

package_installed() {
    $PYTHON_BIN -c "import $1" 2>/dev/null && return 0 || return 1
}

filter_requirements() {
    local req_file="$1"
    local tmp_file="${req_file}.filtered"
    > "$tmp_file"
    while IFS= read -r line; do
        [[ -z "$line" || "$line" =~ ^# ]] && continue
        local pkg_name
        pkg_name=$(echo "$line" | sed 's/[<>=!].*//' | xargs)
        if [[ -z "$pkg_name" ]]; then continue; fi
        local module_name="${pkg_name//-/_}"
        if package_installed "$module_name"; then
            log_info "설치된 패키지 건너뛰기: $pkg_name"
        else
            echo "$line" >> "$tmp_file"
        fi
    done < "$req_file"
    if [ -s "$tmp_file" ]; then
        mv "$tmp_file" "$req_file"
        return 0
    else
        rm -f "$tmp_file"
        log_info "모든 패키지가 이미 설치되었습니다"
        return 1
    fi
}

provisioning_has_valid_hf_token() {
    [[ -n "${HF_TOKEN:-}" ]] || return 1
    local response
    response=$(curl -o /dev/null -s -w "%{http_code}" -X GET "https://huggingface.co/api/whoami-v2" -H "Authorization: Bearer $HF_TOKEN" -H "Content-Type: application/json")
    [ "$response" -eq 200 ]
}

provisioning_has_valid_civitai_token() {
    [[ -n "${CIVITAI_TOKEN:-}" ]] || return 1
    local response
    response=$(curl -o /dev/null -s -w "%{http_code}" -X GET "https://civitai.com/api/v1/models?hidden=1&limit=1" -H "Authorization: Bearer $CIVITAI_TOKEN" -H "Content-Type: application/json")
    [ "$response" -eq 200 ]
}

install_node() {
    local repo_url="$1"
    local repo_name
    repo_name=$(basename "$repo_url" .git)
    local install_path="${COMFYUI_DIR}/custom_nodes/${repo_name}"

    if [ -d "$install_path" ]; then
        if [[ "${AUTO_UPDATE:-true}" != "false" ]]; then
            log_info "노드 업데이트: $repo_name"
            (cd "$install_path" && git pull -q 2>&1 | grep -v "Already up to date" || true)
        else
            log_info "'$repo_name'이(가) 이미 존재합니다. 건너뜁니다"
            return
        fi
    else
        log_info "노드 복제: $repo_name"
        git clone --depth 1 --single-branch "$repo_url" "$install_path" -q 2>&1 || true
    fi

    if [ -f "$install_path/requirements.txt" ]; then
        log_info "$repo_name의 종속성 처리 중..."
        sed -i -e '/^torch/d' -e '/^sam2/d' "$install_path/requirements.txt" 2>/dev/null || true
        if filter_requirements "$install_path/requirements.txt"; then
            log_info "$repo_name의 새로운 종속성 설치 중..."
            $PIP_BIN install -q --no-cache-dir -r "$install_path/requirements.txt" 2>&1 | grep -v "Requirement already satisfied" || log_warn "일부 종속성 설치 실패"
        fi
    fi

    if [ -f "$install_path/install.py" ]; then
        log_info "$repo_name의 설치 스크립트 실행 중..."
        $PYTHON_BIN "$install_path/install.py" 2>&1 || log_warn "설치 스크립트 실행 실패"
    fi
}

download_file() {
    local dest_path="$1"
    local url="$2"
    local filename
    filename=$(basename "$dest_path")
    local tmp_path="${dest_path}.tmp"

    mkdir -p "$(dirname "$dest_path")"
    if [ -s "$dest_path" ]; then
        log_info "파일 '$filename'이(가) 이미 존재하고 완전합니다. 다운로드를 건너뜁니다"
        return 0
    fi

    log_info "다운로드: $filename"
    local max_retries=3
    local attempt=0
    local auth_header=""
    local success=1

    if [[ "$url" =~ huggingface\.co ]] && provisioning_has_valid_hf_token; then
        auth_header="Authorization: Bearer $HF_TOKEN"
        log_info "HuggingFace 토큰 사용"
    elif [[ "$url" =~ civitai\.com ]] && provisioning_has_valid_civitai_token; then
        auth_header="Authorization: Bearer $CIVITAI_TOKEN"
        log_info "CivitAI 토큰 사용"
    fi

    while [ "$attempt" -lt "$max_retries" ]; do
        attempt=$((attempt + 1))
        [ "$attempt" -gt 1 ] && sleep 10
        if command -v aria2c >/dev/null 2>&1; then
            log_info "aria2c (3 스레드)로 다운로드: $filename (시도 $attempt/$max_retries)"
            local aria_opts=(--console-log-level=error -c -x 3 -s 3 -k 1M --max-connection-per-server=3 --max-tries=3 --retry-wait=5 --timeout=180 --file-allocation=falloc --auto-file-renaming=false -d "$(dirname "$dest_path")" -o "${filename}.tmp")
            [[ -n "$auth_header" ]] && aria_opts+=(--header="$auth_header")
            aria2c "${aria_opts[@]}" "$url"
            if [ $? -eq 0 ]; then success=0; break; fi
        else
            log_info "wget으로 다운로드: $filename (시도 $attempt/$max_retries)"
            local wget_opts=(-O "$tmp_path" -c --timeout=60 --tries=3 --content-disposition --show-progress)
            [[ -n "$auth_header" ]] && wget_opts+=(--header="$auth_header")
            wget "${wget_opts[@]}" "$url"
            if [ $? -eq 0 ]; then success=0; break; fi
        fi
    done

    if [ "$success" -eq 0 ] && [ -s "$tmp_path" ]; then
        mv "$tmp_path" "$dest_path"
        log_success "다운로드 완료: $filename"
        return 0
    else
        log_error "다운로드 실패: $filename"
        rm -f "$tmp_path"
        return 1
    fi
}

download_to_directory() {
    local dest_dir="$1"; shift; local urls=("$@")
    if [ ${#urls[@]} -eq 0 ]; then return 0; fi
    mkdir -p "$dest_dir"
    log_info "${#urls[@]}개의 파일을 $dest_dir(으)로 다운로드 중"
    local MAX_PARALLEL=3
    for url in "${urls[@]}"; do
        while [ $(jobs -r | wc -l) -ge $MAX_PARALLEL ]; do sleep 1; done
        local filename
        filename=$(basename "$url" | sed 's/?.*//')
        download_file "${dest_dir}/${filename}" "$url" &
    done
    wait
}

verify_and_retry_downloads() {
    local dest_dir="$1"; shift; local urls=("$@")
    if [ ${#urls[@]} -eq 0 ]; then return 0; fi
    log_info "$dest_dir의 파일 확인 중..."
    local missing_files=()
    for url in "${urls[@]}"; do
        local filename
        filename=$(basename "$url" | sed 's/?.*//')
        local dest_path="${dest_dir}/${filename}"
        if [ ! -s "$dest_path" ]; then
            log_warn "파일 누락 또는 불완전: $filename"
            missing_files+=("$url")
        fi
    done
    if [ ${#missing_files[@]} -gt 0 ]; then
        log_warn "${#missing_files[@]}개의 누락/불완전한 파일 발견, 재다운로드 중..."
        download_to_directory "$dest_dir" "${missing_files[@]}"
    else
        log_success "모든 파일이 존재하고 완전한 것으로 확인되었습니다"
    fi
}

provisioning_print_header() {
    echo ""; echo "##############################################"
    echo "#                                            #"
    echo "#             프로비저닝 컨테이너             #"
    echo "#                                            #"
    echo "#         시간이 좀 걸릴 수 있습니다         #"
    echo "#                                            #"
    echo "##############################################"; echo ""
}

provisioning_print_end() {
    echo ""; echo "##############################################"
    echo "#                                            #"
    echo "#              프로비저닝 완료!              #"
    echo "#                                            #"
    echo "#      Total time: $((END_TIME - START_TIME)) seconds                 #"
    echo "#                                            #"
    echo "##############################################"; echo ""
}

provisioning_step1_install_core_deps() {
    log_step "1" "시스템 및 ComfyUI 핵심 종속성 설치"
    if [ ${#APT_PACKAGES[@]} -gt 0 ]; then
        log_info "${#APT_PACKAGES[@]}개의 시스템 패키지 설치 준비 중..."
        local packages_to_install=()
        for pkg in "${APT_PACKAGES[@]}"; do
            if ! dpkg -s "$pkg" >/dev/null 2>&1; then
                packages_to_install+=("$pkg")
            else
                log_success "$pkg이(가) 이미 설치되어 있습니다. 건너뜁니다"
            fi
        done
        if [ ${#packages_to_install[@]} -gt 0 ]; then
            log_info "설치 중: ${packages_to_install[*]}"
            apt-get update -qq && apt-get install -y -qq "${packages_to_install[@]}" && log_success "시스템 패키지 설치 완료" || log_warn "일부 시스템 패키지 설치 실패"
        fi
    fi

    local comfyui_req_file="${COMFYUI_DIR}/requirements.txt"
    if [ -f "$comfyui_req_file" ]; then
        log_info "ComfyUI 핵심 종속성 처리 중: $comfyui_req_file"
        if filter_requirements "$comfyui_req_file"; then
            log_info "ComfyUI의 새로운 종속성 설치 중..."
            $PIP_BIN install -q --no-cache-dir -r "$comfyui_req_file" 2>&1 | grep -v "Requirement already satisfied" || log_warn "ComfyUI 일부 핵심 종속성 설치 실패"
        fi
    else
        log_warn "ComfyUI의 핵심 requirements.txt 파일을 찾을 수 없습니다!"
    fi
}

provisioning_step2_nodes() {
    log_step "2" "노드 및 추가 패키지 설치"
    if [ ${#NODES[@]} -gt 0 ]; then
        cd "${COMFYUI_DIR}/custom_nodes" || exit 1
        log_info "${#NODES[@]}개의 노드를 병렬 설치 중..."
        local MAX_PARALLEL=2
        for node in "${NODES[@]}"; do
            while [ $(jobs -r | wc -l) -ge $MAX_PARALLEL ]; do sleep 1; done
            install_node "$node" &
        done
        wait
        log_success "노드 설치 완료"
        cd "${COMFYUI_DIR}" || exit 1
    fi
}

provisioning_step3_downloads() {
    log_step "3" "모델 및 워크플로우 파일 다운로드"
    download_to_directory "${COMFYUI_DIR}/models/checkpoints" "${CHECKPOINT_MODELS[@]}"
    download_to_directory "${COMFYUI_DIR}/models/clip_vision" "${CLIP_VISION_MODELS[@]}"
    download_to_directory "${COMFYUI_DIR}/models/ipadapter" "${IPADAPTER_MODELS[@]}"
    download_to_directory "${COMFYUI_DIR}/models/upscale_models" "${ESRGAN_MODELS[@]}"
    if [ ${#WORKFLOWS[@]} -gt 0 ]; then
        log_info "워크플로우 파일 다운로드 중..."
        download_to_directory "${COMFYUI_DIR}/user/default/workflows" "${WORKFLOWS[@]}"
    fi
}

if [ -f "/workspace/.provision_complete" ]; then
    log_success "환경이 구성되었습니다. 반복 실행을 건너뜁니다"
    log_info "재구성하려면 /workspace/.provision_complete를 삭제하십시오"
    exit 0
fi
if [[ -f /.noprovisioning ]]; then
    log_warn "/.noprovisioning 파일 감지됨, 구성을 건너뜁니다"
    exit 0
fi

START_TIME=$(date +%s)
provisioning_print_header
cd "${COMFYUI_DIR}" || exit 1

provisioning_step1_install_core_deps

provisioning_step2_nodes &
NODE_PID=$!
provisioning_step3_downloads
wait $NODE_PID

log_step "4" "다운로드 무결성 확인"
verify_and_retry_downloads "${COMFYUI_DIR}/models/checkpoints" "${CHECKPOINT_MODELS[@]}"
verify_and_retry_downloads "${COMFYUI_DIR}/models/clip_vision" "${CLIP_VISION_MODELS[@]}"
verify_and_retry_downloads "${COMFYUI_DIR}/models/ipadapter" "${IPADAPTER_MODELS[@]}"
verify_and_retry_downloads "${COMFYUI_DIR}/models/upscale_models" "${ESRGAN_MODELS[@]}"
verify_and_retry_downloads "${COMFYUI_DIR}/user/default/workflows" "${WORKFLOWS[@]}"

log_step "5" "구성 완료"
touch "/workspace/.provision_complete"
log_success "구성 마커 파일이 생성되었습니다"

END_TIME=$(date +%s)
provisioning_print_end
log_success "모든 구성 단계가 완료되었습니다!"
exit 0