#!/usr/bin/env bash
# ----------------------------------------------------
# 互動式 Demucs 分離器（只問兩題）- URL 會先用 yt-dlp 下載 -> 再分離
# - 問題1：2軌(人聲/伴奏) 或 4軌(人聲/鼓/貝斯/其它)
# - 問題2：輸入來源（單檔/資料夾/URL，可直接拖曳到終端機）
# 預設：
#   venv=/home/bengju/demucs_env/venv
#   out=~/Music/demucs_out
#   model=htdemucs, jobs=1, 僅用GPU(-d cuda)
#   URL 下載暫存：~/Music/demucs_cache
# 特色：
#   - 不產生日誌檔（包含每檔 .log 與 url_download.log）
#   - 支援拖曳路徑（自動去除引號/尾端空白、展開 ~、還原跳脫空白）
#   - 橫向進度條（#條 + mm:ss）
# ----------------------------------------------------

# 只允許在 bash 下執行，避免 sh/dash 造成 set/陣列/[[ ]] 相容性問題
[ -n "$BASH_VERSION" ] || { echo "請用 bash 執行：chmod +x run_split.sh && ./run_split.sh"; exit 1; }
set -Eeuo pipefail

# === 可調預設值 ===
DEFAULT_VENV="/home/bengju/demucs_env/venv"
DEFAULT_OUT="$HOME/Music/demucs_out"
DEFAULT_MODEL="htdemucs"
DEFAULT_JOBS=1
CACHE="$HOME/Music/demucs_cache"

# === 輸出樣式與輔助 ===
BOLD="$(tput bold 2>/dev/null || true)"; RESET="$(tput sgr0 2>/dev/null || true)"
log(){ printf "%s\n" "$*"; }
err(){ printf "❌ %s\n" "$*" >&2; }
ok(){  printf "✅ %s\n" "$*"; }

# === 橫向進度條（不定長；右側 mm:ss），不產生日誌 ===
progress_bar() {
  local pid=$1 msg="$2" width=30
  local start_ts=$(date +%s) now elapsed pos=0 dir=1 step=1
  printf "%s\n" "$msg"
  while kill -0 "$pid" 2>/dev/null; do
    now=$(date +%s); elapsed=$(( now - start_ts ))
    pos=$(( pos + dir*step ))
    (( pos >= width )) && { dir=-1; pos=$((width-1)); }
    (( pos < 0 ))      && { dir=1;  pos=0; }
    printf "\r[%s%s]  (%02d:%02d)" \
      "$(printf '%*s' $((pos+1)) '' | tr ' ' '#')" \
      "$(printf '%*s' $((width-pos-1)) '')" \
      $((elapsed/60)) $((elapsed%60))
    sleep 0.2
  done
  now=$(date +%s); elapsed=$(( now - start_ts ))
  printf "\r[%s]  (%02d:%02d)\n" "$(printf '%*s' $width '' | tr ' ' '#')" \
    $((elapsed/60)) $((elapsed%60))
}

# === 路徑清理：處理拖曳路徑的引號/尾端空白/跳脫空白/家目錄波浪 ===
sanitize_input() {
  local s="$1"
  # 去前後空白
  s="$(printf '%s' "$s" | sed -E 's/^[[:space:]]+|[[:space:]]+$//g')"
  # 移除最外層的單/雙引號
  if [[ "$s" =~ ^\'(.*)\'$ ]]; then
    s="${BASH_REMATCH[1]}"
  elif [[ "$s" =~ ^\"(.*)\"$ ]]; then
    s="${BASH_REMATCH[1]}"
  fi
  # 將被跳脫的空白還原（e.g., path\ with\ space -> path with space）
  s="${s//\\ / }"
  # 展開家目錄波浪（若以 ~ 開頭）
  if [[ "$s" == "~"* ]]; then
    s="${s/#\~/$HOME}"
  fi
  printf '%s' "$s"
}

# === 互動：兩題 ===
read -r -p "選擇模式 (2=人聲+伴奏；4=人聲+鼓+貝斯+其它) [4]: " MODE
MODE="${MODE:-4}"
while [[ "$MODE" != "2" && "$MODE" != "4" ]]; do
  echo "輸入無效，請輸入 2 或 4"
  read -r -p "選擇模式 (2 或 4): " MODE
done

read -r -p "輸入來源（檔案/資料夾/URL，可直接拖曳）: " INPUT
while [[ -z "$INPUT" ]]; do
  echo "輸入來源不可空白"
  read -r -p "輸入來源（檔案/資料夾/URL，可直接拖曳）: " INPUT
done

# 針對拖曳過來的路徑/手動輸入做清理
INPUT="$(sanitize_input "$INPUT")"

# === 環境檢查 ===
[[ -d "$DEFAULT_VENV" ]] || { err "找不到虛擬環境：$DEFAULT_VENV"; exit 1; }
# shellcheck disable=SC1091
source "$DEFAULT_VENV/bin/activate" || { err "啟用 venv 失敗"; exit 1; }
command -v demucs      >/dev/null 2>&1 || { err "venv 內找不到 demucs"; exit 1; }
command -v nvidia-smi  >/dev/null 2>&1 || { err "找不到 nvidia-smi（請先安裝驅動）"; exit 1; }
mkdir -p "$DEFAULT_OUT" "$CACHE"

# === 蒐集輸入檔（含 URL 下載）— 全程不寫日誌檔 ===
declare -a files
if [[ "$INPUT" =~ ^https?:// ]]; then
  # 需要 yt-dlp + ffmpeg
  command -v yt-dlp >/dev/null 2>&1 || { err "URL 模式需要 yt-dlp；請先安裝：sudo apt install yt-dlp 或 pipx install yt-dlp"; exit 1; }
  command -v ffmpeg >/dev/null 2>&1 || { err "URL 模式需要 ffmpeg；請先安裝：sudo apt install ffmpeg"; exit 1; }

  # 下載為 wav（最佳音質），檔名「標題-ID.wav」；不寫日誌檔
  outtpl="%(title)s-%(id)s.%(ext)s"
  ytcmd=(yt-dlp --no-playlist -f "bestaudio/best" --extract-audio --audio-format wav --audio-quality 0
         -o "$CACHE/$outtpl" "$INPUT")

  ( "${ytcmd[@]}" >/dev/null 2>&1 ) &   # 若想看 yt-dlp 輸出，移除 >/dev/null 2>&1
  pid=$!
  progress_bar "$pid" "下載音檔中（yt-dlp）"
  if ! wait "$pid"; then
    err "下載失敗（未記錄日誌，請重試或移除 >/dev/null 2>&1 以顯示詳情）"; exit 1
  fi

  # 找到最新的 wav 當作下載結果
  mapfile -t candidates < <(find "$CACHE" -maxdepth 1 -type f -iname "*.wav" -printf "%T@ %p\n" | sort -nr | awk '{ $1=""; sub(/^ /,""); print }')
  [[ ${#candidates[@]} -gt 0 ]] || { err "下載後找不到 WAV 檔"; exit 1; }
  dlfile="${candidates[0]}"
  files=("$dlfile")

elif [[ -d "$INPUT" ]]; then
  while IFS= read -r -d '' f; do files+=("$f"); done < <(
    find "$INPUT" -type f \( -iname '*.wav' -o -iname '*.flac' -o -iname '*.mp3' \
                             -o -iname '*.m4a' -o -iname '*.aac'  -o -iname '*.ogg' \
                             -o -iname '*.opus' -o -iname '*.wma' \) -print0 | sort -z
  )
  [[ ${#files[@]} -gt 0 ]] || { err "資料夾裡找不到支援的音訊檔。"; exit 1; }

elif [[ -f "$INPUT" ]]; then
  files=("$INPUT")

else
  err "找不到檔案或資料夾：$INPUT"; exit 1
fi

# === Demucs 參數（只用 GPU）===
BASE_CMD=(demucs -d cuda --jobs "$DEFAULT_JOBS" -n "$DEFAULT_MODEL" -o "$DEFAULT_OUT")
stem_args=()
[[ "$MODE" == "2" ]] && stem_args=(--two-stems=vocals)

# === 分離作業（逐檔、顯示進度、不寫任何 .log）===
total=${#files[@]}
echo "${BOLD}開始分離（GPU 模式）${RESET}"
echo "模式：$MODE 軌"
echo "輸入：$INPUT"
echo "發現檔案數：$total"
echo "輸出：$DEFAULT_OUT"
echo

idx=0
for f in "${files[@]}"; do
  idx=$((idx+1))
  base="$(basename "$f")"
  echo "——— [${idx}/${total}] $base ——"

  # 不重定向到檔案 → 不會產生每檔 .log
  ( "${BASE_CMD[@]}" "${stem_args[@]}" "$f" >/dev/null 2>&1 ) &  # 想看 demucs 詳細輸出，移除 >/dev/null 2>&1
  pid=$!
  progress_bar "$pid" "分離中：$base"

  if ! wait "$pid"; then
    err "[${idx}/${total}] 失敗：$base（未記錄日誌；必要時請移除 >/dev/null 2>&1 以顯示命令輸出）"
  else
    ok  "[${idx}/${total}] 完成：$base"
  fi
done

echo
ok "全部處理完成！"
echo "輸出根目錄：$DEFAULT_OUT"