#!/bin/bash

# === API Configuration ===
export ASSEMBLYAI_API_KEY="4a81251e347f4e888f9c72723849e900"

# GLM (Zhipu AI) - Currently active
export OPENAI_API_KEY="1477a6bbb684437ba56cdaef04e8ddf6.QlNZUfUcfTfIChSE"
export OPENAI_BASE_URL="https://open.bigmodel.cn/api/paas/v4"
export LLM_MODEL="glm-4.5-air"
# export LLM_MODEL="GLM-4.7-Flash"

# === Usage ===
# Run all steps:
#   ./run.sh
#   ./run.sh all
#
# Run individual step:
#   ./run.sh 1    # Step 1: Transcribe
#   ./run.sh 2    # Step 2: Format lines
#   ./run.sh 3    # Step 3: Infer speakers
#   ./run.sh 4    # Step 4: Translate to Chinese
#   ./run.sh 5    # Step 5: Assign colors
#
# Force reprocessing (disregard progress marks):
#   ./run.sh -f
#   ./run.sh -f 1
#   ./run.sh 1 -f

# Parse arguments
FORCE_FLAG=""
STEP="all"

for arg in "$@"; do
    if [ "$arg" = "-f" ] || [ "$arg" = "--force" ]; then
        FORCE_FLAG="-f"
    elif [ "$arg" = "1" ] || [ "$arg" = "2" ] || [ "$arg" = "3" ] || [ "$arg" = "4" ] || [ "$arg" = "5" ] || [ "$arg" = "all" ]; then
        STEP="$arg"
    fi
done

# Header
echo "========================================"
echo "Transcription Pipeline"
echo "========================================"
if [ -n "$FORCE_FLAG" ]; then
    echo "Force mode: ON (disregarding progress marks)"
fi
echo ""

FAILED=0

# Step 1: Transcribe
if [ "$STEP" = "all" ] || [ "$STEP" = "1" ]; then
    echo ""
    echo "========================================"
    echo "Step 1: Transcribing..."
    echo "========================================"
    uv run step1_transcribe.py $FORCE_FLAG || FAILED=1
fi

# Step 2: Format lines
if [ "$STEP" = "all" ] || [ "$STEP" = "2" ]; then
    echo ""
    echo "========================================"
    echo "Step 2: Formatting lines..."
    echo "========================================"
    uv run step2_format.py || FAILED=1
fi

# Step 3: Infer speakers
if [ "$STEP" = "all" ] || [ "$STEP" = "3" ]; then
    echo ""
    echo "========================================"
    echo "Step 3: Inferring speakers..."
    echo "========================================"
    uv run step3_infer_speakers.py || FAILED=1
fi

# Step 4: Translate to Chinese
if [ "$STEP" = "all" ] || [ "$STEP" = "4" ]; then
    echo ""
    echo "========================================"
    echo "Step 4: Translating to Chinese..."
    echo "========================================"
    uv run step4_translate.py || FAILED=1
fi

# Step 5: Assign colors
if [ "$STEP" = "all" ] || [ "$STEP" = "5" ]; then
    echo ""
    echo "========================================"
    echo "Step 5: Assigning colors..."
    echo "========================================"
    uv run step5_assign_colors.py || FAILED=1
fi

echo ""
echo "========================================"
if [ $FAILED -eq 0 ]; then
    echo "✅ All steps completed successfully!"
else
    echo "⚠️  Some steps had errors"
fi
echo "========================================"