#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOF'
Usage:
  run_music_flow.sh --prompt "<text>" --target /abs/output/dir [--count N] [--session NAME] [--no-headed]

Example:
  run_music_flow.sh \
    --prompt "创作一段 90 BPM 的 lo-fi hiphop，温暖、夜晚、钢琴和刷镲，时长 30 秒。" \
    --target /Users/xd/java/xhs/output/gemini-music \
    --count 2
EOF
}

PROMPT=""
TARGET=""
COUNT=1
SESSION="gmw$(date +%s)"
HEADED=1

while [[ $# -gt 0 ]]; do
  case "$1" in
    --prompt)
      PROMPT="${2:-}"
      shift 2
      ;;
    --target)
      TARGET="${2:-}"
      shift 2
      ;;
    --count)
      COUNT="${2:-1}"
      shift 2
      ;;
    --session)
      SESSION="${2:-$SESSION}"
      shift 2
      ;;
    --no-headed)
      HEADED=0
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown arg: $1" >&2
      usage
      exit 1
      ;;
  esac
done

if [[ -z "$PROMPT" || -z "$TARGET" ]]; then
  echo "Both --prompt and --target are required." >&2
  usage
  exit 1
fi

if ! [[ "$COUNT" =~ ^[0-9]+$ ]] || [[ "$COUNT" -lt 1 ]]; then
  echo "--count must be a positive integer." >&2
  exit 1
fi

CODEX_HOME="${CODEX_HOME:-$HOME/.codex}"
PWCLI="${PWCLI:-$CODEX_HOME/skills/playwright/scripts/playwright_cli.sh}"
COLLECT_SCRIPT="$(cd "$(dirname "$0")" && pwd)/collect_downloads.py"

if ! command -v npx >/dev/null 2>&1; then
  echo "npx is required." >&2
  exit 1
fi
if [[ ! -x "$PWCLI" ]]; then
  echo "Playwright wrapper not found or not executable: $PWCLI" >&2
  exit 1
fi
if [[ ! -f "$COLLECT_SCRIPT" ]]; then
  echo "Collector script not found: $COLLECT_SCRIPT" >&2
  exit 1
fi

pw() {
  "$PWCLI" --session "$SESSION" "$@"
}

json_escape() {
  python3 - "$1" <<'PY'
import json
import sys
print(json.dumps(sys.argv[1]))
PY
}

is_login_required() {
  local out
  out="$(
    pw eval "() => {
      const hasAccount = !!document.querySelector('button[aria-label*=\\\"Google 账号\\\"], button[aria-label*=\\\"Google Account\\\"]');
      const hasService = !!document.querySelector('a[href*=\\\"ServiceLogin\\\"]');
      const hasLoginCtl = Array.from(document.querySelectorAll('a,button')).some(el => /登录|Sign in/i.test((el.textContent || '').trim()));
      return !hasAccount && (hasService || hasLoginCtl);
    }"
  )"
  echo "$out" | rg -q '^true$'
}

enter_music_tool() {
  local js
  js="$(cat <<'JS'
const labels = [/创作音乐/, /制作音乐/, /Create music/i, /Music/i];

const tryCardButtons = async () => {
  for (const re of labels) {
    const btn = page.getByRole('button', { name: re }).first();
    if (await btn.count()) {
      try {
        await btn.click({ timeout: 2000 });
        return true;
      } catch (_) {
        // Overlay may intercept pointer. Fall through to menu strategy.
      }
    }
  }
  return false;
};

const tryToolMenu = async () => {
  await page.getByRole('button', { name: '工具', exact: true }).click();
  for (const re of labels) {
    const itemCheck = page.getByRole('menuitemcheckbox', { name: re }).first();
    if (await itemCheck.count()) {
      await itemCheck.click();
      return true;
    }
    const itemPlain = page.getByRole('menuitem', { name: re }).first();
    if (await itemPlain.count()) {
      await itemPlain.click();
      return true;
    }
  }
  return false;
};

let ok = await tryCardButtons();
if (!ok) ok = await tryToolMenu();
if (!ok) {
  // Re-open the tool menu once and retry as a last attempt.
  ok = await tryToolMenu();
}
if (!ok) {
  throw new Error('Music tool entry not found');
}
JS
)"
  pw run-code "$js" >/dev/null
}

submit_and_download_one() {
  local track_prompt="$1"
  local escaped
  escaped="$(json_escape "$track_prompt")"
  local js
  js="$(cat <<JS
const prompt = $escaped;
const input = page.getByRole('textbox', { name: /为 Gemini 输入提示|Enter a prompt/i }).first();
await input.click();
await input.fill(prompt);
await input.press('Enter');

const stopBtn = page.getByRole('button', { name: /停止回答|Stop response/i }).first();
await stopBtn.waitFor({ state: 'visible', timeout: 15000 }).catch(() => {});
await stopBtn.waitFor({ state: 'hidden', timeout: 240000 });

const downloadBtn = page.getByRole('button', { name: /下载音乐作品|Download music/i }).last();
await downloadBtn.click();

const mp3Item = page.getByRole('menuitem', { name: /纯音频|MP3/i }).first();
if (await mp3Item.count()) {
  await mp3Item.click();
} else {
  const anyItem = page.getByRole('menuitem').first();
  if (await anyItem.count()) await anyItem.click();
}

await page.waitForTimeout(1200);
JS
)"
  pw run-code "$js" >/dev/null
}

mkdir -p "$TARGET"
start_ts="$(python3 - <<'PY'
import time
print(time.time())
PY
)"

if [[ "$HEADED" -eq 1 ]]; then
  pw open "https://gemini.google.com/app" --headed >/dev/null
else
  pw open "https://gemini.google.com/app" >/dev/null
fi
pw snapshot >/dev/null

if is_login_required; then
  echo "Gemini is not logged in. Please log in at https://gemini.google.com/app and rerun." >&2
  exit 2
fi

enter_music_tool

for ((i=1; i<=COUNT; i++)); do
  current_prompt="$PROMPT"
  if [[ "$COUNT" -gt 1 ]]; then
    current_prompt="$PROMPT
变体要求：这是第 $i / $COUNT 首。保持风格一致，但旋律和节奏细节需要变化。"
  fi
  submit_and_download_one "$current_prompt"
done

python3 "$COLLECT_SCRIPT" \
  --target "$TARGET" \
  --since "$start_ts" \
  --expected-count "$COUNT" \
  --limit "$COUNT" \
  --prefix "gemini-music" \
  --prompt "$PROMPT"