#!/bin/bash
# Aria NotebookLM Wrapper (KAR-69)
#
# Subcommands:
#   audio <URL>          Single-source notebook, gen audio + mind-map, download both
#   multi <URL1> <URL2>  Multi-source notebook, audio overview from all sources
#   raw <args...>        Pass-through to underlying notebooklm CLI
#
# Output dir: /tmp/aria-notebook-<hash>/
#   audio-overview.mp3   (typically 10-15 MB)
#   mindmap.json
#   metadata.json        (notebook_id, source_ids, artifact_ids, urls, generated_at)
#
# Env:
#   NLM_TIMEOUT=600      Max seconds to wait for audio generation
#   NLM_CLEANUP=0        If 1, delete the NotebookLM notebook after download
#
# Exit codes:
#   0  OK
#   2  Auth missing/invalid
#   3  Source-add failed
#   4  Audio generation timed out

set -euo pipefail

NLM=/root/aria/mcps/notebooklm/venv/bin/notebooklm
NLM_TIMEOUT="${NLM_TIMEOUT:-600}"
NLM_CLEANUP="${NLM_CLEANUP:-0}"
LOG=/var/log/aria-notebook.log

say()  { printf '[%s] %s\n' "$(date -Is)" "$*" | tee -a "$LOG" >&2; }
fail() { say "ERROR: $*"; exit "${2:-1}"; }

if [[ ! -x "$NLM" ]]; then
    fail "notebooklm CLI not found at $NLM. Run: cd /root/aria/mcps/notebooklm && python3 -m venv venv && ./venv/bin/pip install 'notebooklm-py[browser]'" 2
fi

# Quick auth check
if ! "$NLM" auth check 2>&1 | grep -q "Token fetch.*pass\|Token fetch.*skipped"; then
    if ! "$NLM" auth check --test 2>&1 | grep -q "Token fetch.*pass"; then
        fail "NotebookLM auth invalid. Run 'notebooklm login' on Mac, scp storage_state.json to /root/.notebooklm/profiles/default/" 2
    fi
fi

# Helper: compute work dir from input set
work_dir_for() {
    local key="$1"
    local hash
    hash=$(printf '%s' "$key" | sha1sum | cut -c1-12)
    echo "/tmp/aria-notebook-$hash"
}

cmd_audio() {
    local url="${1:?Usage: aria-notebook.sh audio <URL>}"
    local work
    work=$(work_dir_for "audio:$url")
    mkdir -p "$work"
    cd "$work"
    say "audio: $url -> $work"

    # 1. Create notebook
    local ts
    ts=$(date -Is)
    local nb_id
    nb_id=$("$NLM" create "Aria $ts" 2>&1 | grep -oE '[a-f0-9-]{36}' | head -1)
    [[ -n "$nb_id" ]] || fail "notebook creation failed" 3
    say "notebook: $nb_id"

    "$NLM" use "$nb_id" >/dev/null

    # 2. Add source
    local src_out
    src_out=$("$NLM" source add "$url" 2>&1)
    echo "$src_out" >>"$LOG"
    local src_id
    src_id=$(echo "$src_out" | grep -oE '[a-f0-9-]{36}' | tail -1)
    [[ -n "$src_id" ]] || fail "source add failed: $src_out" 3
    say "source: $src_id"

    # 3. Trigger mind-map (fast) + audio (slow) in parallel
    say "generating mind-map..."
    "$NLM" generate mind-map >mindmap.gen.log 2>&1 || say "WARN: mind-map gen failed (continuing)"

    say "generating audio overview (up to ${NLM_TIMEOUT}s)..."
    local audio_id
    audio_id=$("$NLM" generate audio 2>&1 | grep -oE 'Started: [a-f0-9-]{36}' | awk '{print $2}')
    [[ -n "$audio_id" ]] || fail "audio gen kickoff failed" 4
    say "audio task: $audio_id"

    # 4. Download mind-map (already done)
    "$NLM" download mind-map mindmap.json >>"$LOG" 2>&1 || say "WARN: mind-map download failed"

    # Render Mind-Map PNG (KAR-69 A)
    if [[ -f mindmap.json && -x /root/aria/scripts/aria-mindmap-to-png.py ]]; then
        python3 /root/aria/scripts/aria-mindmap-to-png.py mindmap.json mindmap.png >>"$LOG" 2>&1 \
            || say "WARN: mind-map PNG render failed"
    fi

    # 5. Wait for audio
    "$NLM" artifact wait "$audio_id" --timeout "$NLM_TIMEOUT" >audio.wait.log 2>&1 \
        || fail "audio generation timed out after ${NLM_TIMEOUT}s. Notebook: $nb_id, task: $audio_id" 4

    # 6. Download audio
    "$NLM" download audio audio-overview.mp3 >>"$LOG" 2>&1 \
        || fail "audio download failed" 4

    # 6b. Compress if >40 MB (KAR-69 D) — Telegram-Limit is 50 MB
    local audio_size_mb
    audio_size_mb=$(stat -c%s audio-overview.mp3 2>/dev/null | awk '{printf "%d", $1/1024/1024}')
    if (( audio_size_mb > 40 )); then
        say "audio ${audio_size_mb} MB > 40 MB, compressing to 64k mono..."
        if ffmpeg -y -i audio-overview.mp3 -ac 1 -b:a 64k audio-overview-compressed.mp3 \
           >>"$LOG" 2>&1; then
            local new_size_mb
            new_size_mb=$(stat -c%s audio-overview-compressed.mp3 2>/dev/null | awk '{printf "%d", $1/1024/1024}')
            say "compressed: ${audio_size_mb} MB -> ${new_size_mb} MB"
        else
            say "WARN: compression failed, keeping original"
        fi
    fi

    # 7. Metadata
    local audio_file_for_telegram="audio-overview.mp3"
    [[ -f audio-overview-compressed.mp3 ]] && audio_file_for_telegram="audio-overview-compressed.mp3"

    cat >metadata.json <<EOF
{
  "input": "$url",
  "mode": "audio",
  "notebook_id": "$nb_id",
  "source_ids": ["$src_id"],
  "audio_task_id": "$audio_id",
  "generated_at": "$ts",
  "files": {
    "audio": "audio-overview.mp3",
    "audio_for_telegram": "$audio_file_for_telegram",
    "mindmap": "mindmap.json",
    "mindmap_png": "mindmap.png"
  }
}
EOF

    # 8. Optional cleanup
    if [[ "$NLM_CLEANUP" = "1" ]]; then
        say "cleanup: deleting notebook $nb_id"
        "$NLM" delete "$nb_id" >>"$LOG" 2>&1 || say "WARN: cleanup failed"
    fi

    say "done: $work"
    echo "$work"
}

cmd_multi() {
    [[ $# -lt 2 ]] && fail "Usage: aria-notebook.sh multi <URL1> <URL2> [URL3...]"
    local urls=("$@")
    local key
    key=$(printf '%s\n' "${urls[@]}" | sort | tr '\n' '|')
    local work
    work=$(work_dir_for "multi:$key")
    mkdir -p "$work"
    cd "$work"
    say "multi (${#urls[@]} sources) -> $work"

    local ts
    ts=$(date -Is)
    local nb_id
    nb_id=$("$NLM" create "Aria Multi $ts" 2>&1 | grep -oE '[a-f0-9-]{36}' | head -1)
    [[ -n "$nb_id" ]] || fail "notebook creation failed" 3
    say "notebook: $nb_id"

    "$NLM" use "$nb_id" >/dev/null

    local src_ids=()
    for url in "${urls[@]}"; do
        local src_out
        src_out=$("$NLM" source add "$url" 2>&1)
        local src_id
        src_id=$(echo "$src_out" | grep -oE '[a-f0-9-]{36}' | tail -1)
        if [[ -n "$src_id" ]]; then
            src_ids+=("$src_id")
            say "  +source: $src_id ($url)"
        else
            say "  WARN: source add failed for $url"
        fi
    done
    [[ ${#src_ids[@]} -gt 0 ]] || fail "no sources added" 3

    say "generating mind-map (multi-source)..."
    "$NLM" generate mind-map >mindmap.gen.log 2>&1 || say "WARN: mind-map gen failed"

    say "generating audio overview (multi-source, up to ${NLM_TIMEOUT}s)..."
    local audio_id
    audio_id=$("$NLM" generate audio 2>&1 | grep -oE 'Started: [a-f0-9-]{36}' | awk '{print $2}')
    [[ -n "$audio_id" ]] || fail "audio gen kickoff failed" 4

    "$NLM" download mind-map mindmap.json >>"$LOG" 2>&1 || true
    "$NLM" artifact wait "$audio_id" --timeout "$NLM_TIMEOUT" >audio.wait.log 2>&1 \
        || fail "audio generation timed out" 4
    "$NLM" download audio audio-overview.mp3 >>"$LOG" 2>&1 || fail "audio download failed" 4

    cat >metadata.json <<EOF
{
  "input": $(printf '%s\n' "${urls[@]}" | python3 -c 'import json,sys; print(json.dumps([l.strip() for l in sys.stdin if l.strip()]))'),
  "mode": "multi",
  "notebook_id": "$nb_id",
  "source_ids": $(printf '"%s"\n' "${src_ids[@]}" | python3 -c 'import json,sys; print(json.dumps([l.strip().strip("\"") for l in sys.stdin if l.strip()]))'),
  "audio_task_id": "$audio_id",
  "generated_at": "$ts",
  "files": {"audio": "audio-overview.mp3", "mindmap": "mindmap.json"}
}
EOF

    [[ "$NLM_CLEANUP" = "1" ]] && "$NLM" delete "$nb_id" >>"$LOG" 2>&1 || true

    say "done: $work"
    echo "$work"
}

case "${1:-}" in
    audio)     shift; cmd_audio "$@" ;;
    multi)     shift; cmd_multi "$@" ;;
    raw)       shift; "$NLM" "$@" ;;
    *)         echo "Usage: $0 {audio <URL> | multi <URL1> <URL2>... | raw <args...>}" >&2
               exit 1 ;;
esac