#!/bin/bash
# ============================================================================
# Diarization Pipeline - Unified CLI Tool (feature/visualization branch)
# ============================================================================
# Commands: setup, start, stop, status, run, logs, clean
# This is the ONLY script needed - handles everything end-to-end
# ============================================================================

set -e

# ============================================================================
# Configuration
# ============================================================================
PROJECT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_DIR="$PROJECT_DIR/venv"
BACKEND_DIR="$PROJECT_DIR/visualizer/backend"
FRONTEND_DIR="$PROJECT_DIR/visualizer/frontend"
DATA_DIR="$PROJECT_DIR/data/fast_output_v6"

# PID files (stored in project root)
BACKEND_PID_FILE="$PROJECT_DIR/backend.pid"
FRONTEND_PID_FILE="$PROJECT_DIR/frontend.pid"
NGROK_PID_FILE="$PROJECT_DIR/ngrok.pid"

# Log files
BACKEND_LOG="$PROJECT_DIR/backend.log"
FRONTEND_LOG="$PROJECT_DIR/frontend.log"
NGROK_LOG="$PROJECT_DIR/ngrok.log"

# Ports
BACKEND_PORT=8000
FRONTEND_PORT=3000

# Ngrok configuration
NGROK_AUTHTOKEN="2L2WOT68heu3mLZwaArdoEk8kMX_4pJqg2ufc5931uhRcQT1E"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

# ============================================================================
# Helper Functions
# ============================================================================

log_info() { echo -e "${BLUE}ℹ ${NC}$1"; }
log_success() { echo -e "${GREEN}✓ ${NC}$1"; }
log_warn() { echo -e "${YELLOW}⚠ ${NC}$1"; }
log_error() { echo -e "${RED}✗ ${NC}$1"; }
log_header() { 
    echo ""
    echo -e "${CYAN}══════════════════════════════════════════════════════════════${NC}"
    echo -e "${CYAN}  $1${NC}"
    echo -e "${CYAN}══════════════════════════════════════════════════════════════${NC}"
}

# Check if a port is in use (use ss since lsof may not be available)
port_in_use() {
    local port=$1
    ss -tlnp 2>/dev/null | grep -q ":$port " && return 0
    netstat -tlnp 2>/dev/null | grep -q ":$port " && return 0
    return 1
}

# Get PID of process using a port
get_pid_on_port() {
    local port=$1
    # Try ss first (more reliable)
    local pid=$(ss -tlnp 2>/dev/null | grep ":$port " | sed -n 's/.*pid=\([0-9]*\).*/\1/p' | head -1)
    if [ -n "$pid" ]; then
        echo "$pid"
        return
    fi
    # Fallback to lsof
    lsof -t -i:$port 2>/dev/null | head -1
}

# Kill process on a port
kill_port() {
    local port=$1
    local pid=$(get_pid_on_port $port)
    if [ -n "$pid" ]; then
        kill $pid 2>/dev/null || true
        sleep 1
        # Force kill if still alive
        if ps -p $pid > /dev/null 2>&1; then
            kill -9 $pid 2>/dev/null || true
        fi
        return 0
    fi
    return 1
}

# Check if venv is active
check_venv() {
    if [ ! -d "$VENV_DIR" ]; then
        log_error "Virtual environment not found. Run: ./setup.sh setup"
        exit 1
    fi
}

# Activate venv
activate_venv() {
    source "$VENV_DIR/bin/activate"
}

# Check git branch
check_branch() {
    local branch=$(git branch --show-current 2>/dev/null || echo "unknown")
    if [ "$branch" != "feature/visualization" ]; then
        log_warn "Current branch: $branch (expected: feature/visualization)"
    else
        log_success "Branch: feature/visualization"
    fi
}

# ============================================================================
# Command: setup
# ============================================================================
cmd_setup() {
    log_header "Setting up Diarization Pipeline"
    
    echo ""
    check_branch
    
    # Check Python version
    local python_version=$(python3 --version 2>&1 | awk '{print $2}')
    log_success "Python version: $python_version"
    
    # Create virtual environment (single venv for everything)
    if [ ! -d "$VENV_DIR" ]; then
        log_info "Creating virtual environment..."
        python3 -m venv "$VENV_DIR"
        log_success "Virtual environment created"
    else
        log_success "Virtual environment exists"
    fi
    
    # Activate and install dependencies
    activate_venv
    log_info "Installing Python dependencies..."
    pip install -q --upgrade pip
    
    # Install PyTorch (GPU if available)
    if command -v nvidia-smi &> /dev/null; then
        log_info "CUDA detected - installing GPU PyTorch..."
        pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    else
        log_warn "No CUDA - installing CPU PyTorch..."
        pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
    fi
    
    pip install -q -r "$PROJECT_DIR/requirements.txt"
    log_success "Python dependencies installed"
    
    # Check ffmpeg
    if ! command -v ffmpeg &> /dev/null; then
        log_error "ffmpeg not found! Install with: sudo apt-get install ffmpeg"
    else
        log_success "ffmpeg found"
    fi
    
    # Setup frontend (Node.js)
    if [ -d "$FRONTEND_DIR" ]; then
        log_info "Setting up frontend..."
        cd "$FRONTEND_DIR"
        if [ ! -d "node_modules" ]; then
            npm install --silent
        fi
        log_success "Frontend dependencies installed"
        cd "$PROJECT_DIR"
    fi
    
    # Create data directories
    mkdir -p "$DATA_DIR"
    mkdir -p "$PROJECT_DIR/logs"
    log_success "Data directories created"
    
    log_header "Setup Complete!"
    echo ""
    echo -e "  ${BOLD}Next steps:${NC}"
    echo -e "    ${CYAN}./setup.sh start${NC}     - Start visualizer (backend + frontend)"
    echo -e "    ${CYAN}./setup.sh run URL${NC}   - Process a YouTube video"
    echo -e "    ${CYAN}./setup.sh --help${NC}    - Show all commands"
    echo ""
}

# ============================================================================
# Command: start
# ============================================================================
cmd_start() {
    log_header "Starting Visualizer Services"
    
    check_venv
    check_branch
    
    # Clean up any existing processes first
    cmd_stop_quiet
    
    # Start Backend
    echo ""
    log_info "Starting backend (FastAPI) on port $BACKEND_PORT..."
    
    cd "$BACKEND_DIR"
    activate_venv
    
    # Start uvicorn with nohup
    nohup "$VENV_DIR/bin/uvicorn" main:app \
        --host 0.0.0.0 \
        --port $BACKEND_PORT \
        > "$BACKEND_LOG" 2>&1 &
    
    local backend_pid=$!
    echo $backend_pid > "$BACKEND_PID_FILE"
    cd "$PROJECT_DIR"
    
    # Wait for backend to be ready
    local max_wait=30
    local waited=0
    echo -n "  Waiting for backend"
    while [ $waited -lt $max_wait ]; do
        # Check if health endpoint returns valid JSON
        local response=$(curl -s --max-time 2 "http://localhost:$BACKEND_PORT/api/health" 2>/dev/null || echo "")
        if [[ "$response" == *"healthy"* ]]; then
            echo ""
            log_success "Backend started (PID: $backend_pid)"
            break
        fi
        echo -n "."
        sleep 1
        waited=$((waited + 1))
    done
    
    if [ $waited -ge $max_wait ]; then
        echo ""
        log_error "Backend failed to start! Check: tail -f $BACKEND_LOG"
        exit 1
    fi
    
    # Start Frontend
    log_info "Starting frontend (Next.js) on port $FRONTEND_PORT..."
    
    cd "$FRONTEND_DIR"
    nohup npm run dev > "$FRONTEND_LOG" 2>&1 &
    local frontend_pid=$!
    echo $frontend_pid > "$FRONTEND_PID_FILE"
    cd "$PROJECT_DIR"
    
    # Wait for frontend to be ready
    waited=0
    echo -n "  Waiting for frontend"
    while [ $waited -lt $max_wait ]; do
        # Check if frontend returns any response
        local http_code=$(curl -s --max-time 2 -o /dev/null -w "%{http_code}" "http://localhost:$FRONTEND_PORT" 2>/dev/null || echo "000")
        if [[ "$http_code" == "200" ]] || [[ "$http_code" == "304" ]]; then
            echo ""
            log_success "Frontend started (PID: $frontend_pid)"
            break
        fi
        echo -n "."
        sleep 1
        waited=$((waited + 1))
    done
    
    if [ $waited -ge $max_wait ]; then
        echo ""
        log_error "Frontend failed to start! Check: tail -f $FRONTEND_LOG"
        exit 1
    fi
    
    # Start Ngrok tunnel for public access
    if command -v ngrok &> /dev/null; then
        log_info "Starting ngrok tunnel for public access..."
        
        # Configure ngrok authtoken if not already done
        ngrok config add-authtoken "$NGROK_AUTHTOKEN" 2>/dev/null || true
        
        # Kill any existing ngrok
        pkill -f ngrok 2>/dev/null || true
        sleep 1
        
        # Start ngrok in background
        nohup ngrok http $FRONTEND_PORT > "$NGROK_LOG" 2>&1 &
        local ngrok_pid=$!
        echo $ngrok_pid > "$NGROK_PID_FILE"
        
        # Wait for ngrok to start and get URL
        sleep 5
        local ngrok_url=$(curl -s http://localhost:4040/api/tunnels 2>/dev/null | python3 -c "import sys,json; t=json.load(sys.stdin).get('tunnels',[]); print(t[0]['public_url'] if t else '')" 2>/dev/null || echo "")
        
        if [ -n "$ngrok_url" ]; then
            log_success "Ngrok tunnel active"
        else
            log_warn "Ngrok may still be starting - check ngrok.log"
            ngrok_url="(starting...)"
        fi
    else
        log_warn "Ngrok not installed - public access unavailable"
        ngrok_url=""
    fi
    
    # Show status
    log_header "Visualizer is Running!"
    echo ""
    echo -e "  ${BOLD}Local URLs:${NC}"
    echo -e "    Frontend:  ${CYAN}http://localhost:$FRONTEND_PORT${NC}"
    echo -e "    Backend:   ${CYAN}http://localhost:$BACKEND_PORT${NC}"
    echo -e "    API Docs:  ${CYAN}http://localhost:$BACKEND_PORT/docs${NC}"
    
    if [ -n "$ngrok_url" ] && [ "$ngrok_url" != "(starting...)" ]; then
        echo ""
        echo -e "  ${BOLD}Public URL (ngrok):${NC}"
        echo -e "    ${GREEN}$ngrok_url${NC}"
    fi
    
    echo ""
    echo -e "  ${BOLD}Logs:${NC}"
    echo -e "    ${CYAN}tail -f backend.log${NC}"
    echo -e "    ${CYAN}tail -f frontend.log${NC}"
    echo -e "    ${CYAN}tail -f ngrok.log${NC}"
    echo ""
    echo -e "  ${BOLD}Stop:${NC}"
    echo -e "    ${CYAN}./setup.sh stop${NC}"
    echo ""
}

# ============================================================================
# Command: stop (quiet version for internal use)
# ============================================================================
cmd_stop_quiet() {
    # Kill by PID file
    for pidfile in "$BACKEND_PID_FILE" "$FRONTEND_PID_FILE" "$NGROK_PID_FILE"; do
        if [ -f "$pidfile" ]; then
            local pid=$(cat "$pidfile")
            if [ -n "$pid" ] && ps -p $pid > /dev/null 2>&1; then
                kill $pid 2>/dev/null || true
            fi
            rm -f "$pidfile"
        fi
    done
    
    # Kill by port (backup)
    kill_port $BACKEND_PORT 2>/dev/null || true
    kill_port $FRONTEND_PORT 2>/dev/null || true
    
    # Kill any lingering processes (use -9 for stubborn ones)
    pkill -9 -f "uvicorn" 2>/dev/null || true
    pkill -9 -f "next" 2>/dev/null || true
    pkill -9 -f "ngrok" 2>/dev/null || true
    
    sleep 1
}

# ============================================================================
# Command: stop
# ============================================================================
cmd_stop() {
    log_header "Stopping Visualizer Services"
    echo ""
    
    local stopped=false
    
    # Stop backend
    if [ -f "$BACKEND_PID_FILE" ]; then
        local pid=$(cat "$BACKEND_PID_FILE")
        if [ -n "$pid" ] && ps -p $pid > /dev/null 2>&1; then
            kill $pid 2>/dev/null
            log_success "Backend stopped (PID: $pid)"
            stopped=true
        fi
        rm -f "$BACKEND_PID_FILE"
    fi
    
    # Stop frontend
    if [ -f "$FRONTEND_PID_FILE" ]; then
        local pid=$(cat "$FRONTEND_PID_FILE")
        if [ -n "$pid" ] && ps -p $pid > /dev/null 2>&1; then
            kill $pid 2>/dev/null
            log_success "Frontend stopped (PID: $pid)"
            stopped=true
        fi
        rm -f "$FRONTEND_PID_FILE"
    fi
    
    # Stop ngrok
    if [ -f "$NGROK_PID_FILE" ]; then
        local pid=$(cat "$NGROK_PID_FILE")
        if [ -n "$pid" ] && ps -p $pid > /dev/null 2>&1; then
            kill $pid 2>/dev/null
            log_success "Ngrok stopped (PID: $pid)"
            stopped=true
        fi
        rm -f "$NGROK_PID_FILE"
    fi
    
    # Kill by port (cleanup any orphans)
    if port_in_use $BACKEND_PORT; then
        kill_port $BACKEND_PORT
        log_success "Killed process on port $BACKEND_PORT"
        stopped=true
    fi
    
    if port_in_use $FRONTEND_PORT; then
        kill_port $FRONTEND_PORT
        log_success "Killed process on port $FRONTEND_PORT"
        stopped=true
    fi
    
    # Final cleanup - use SIGKILL for stubborn processes
    pkill -9 -f "uvicorn" 2>/dev/null && log_success "Killed lingering uvicorn" || true
    pkill -9 -f "next" 2>/dev/null && log_success "Killed lingering next" || true
    pkill -9 -f "ngrok" 2>/dev/null && log_success "Killed lingering ngrok" || true
    
    if [ "$stopped" = false ]; then
        log_info "No services were running"
    fi
    
    echo ""
    log_success "All services stopped"
    echo ""
}

# ============================================================================
# Command: status
# ============================================================================
cmd_status() {
    log_header "Service Status"
    echo ""
    
    # Check branch
    check_branch
    echo ""
    
    # Backend status
    echo -e "${BOLD}Backend (port $BACKEND_PORT):${NC}"
    if port_in_use $BACKEND_PORT; then
        local pid=$(get_pid_on_port $BACKEND_PORT)
        local health=$(curl -s "http://localhost:$BACKEND_PORT/api/health" 2>/dev/null || echo "")
        if [ -n "$health" ]; then
            log_success "Running (PID: $pid)"
            echo -e "    Health: ${GREEN}$health${NC}"
        else
            log_warn "Port in use but not responding (PID: $pid)"
        fi
    else
        log_error "Not running"
    fi
    echo ""
    
    # Frontend status
    echo -e "${BOLD}Frontend (port $FRONTEND_PORT):${NC}"
    if port_in_use $FRONTEND_PORT; then
        local pid=$(get_pid_on_port $FRONTEND_PORT)
        if curl -s "http://localhost:$FRONTEND_PORT" > /dev/null 2>&1; then
            log_success "Running (PID: $pid)"
        else
            log_warn "Port in use but not responding (PID: $pid)"
        fi
    else
        log_error "Not running"
    fi
    echo ""
    
    # Ngrok status
    echo -e "${BOLD}Ngrok (public access):${NC}"
    if pgrep -f ngrok > /dev/null 2>&1; then
        local ngrok_url=$(curl -s http://localhost:4040/api/tunnels 2>/dev/null | python3 -c "import sys,json; t=json.load(sys.stdin).get('tunnels',[]); print(t[0]['public_url'] if t else '')" 2>/dev/null || echo "")
        if [ -n "$ngrok_url" ]; then
            log_success "Running"
            echo -e "    Public URL: ${GREEN}$ngrok_url${NC}"
        else
            log_warn "Running but tunnel not ready"
        fi
    else
        log_warn "Not running (no public access)"
    fi
    echo ""
    
    # Data status
    echo -e "${BOLD}Processed Videos:${NC}"
    if [ -d "$DATA_DIR" ]; then
        local count=$(find "$DATA_DIR" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | wc -l)
        log_info "$count videos in $DATA_DIR"
    else
        log_warn "Data directory not found"
    fi
    echo ""
    
    # URLs
    if port_in_use $BACKEND_PORT && port_in_use $FRONTEND_PORT; then
        echo -e "${BOLD}Local Access:${NC}"
        echo -e "    Frontend:  ${CYAN}http://localhost:$FRONTEND_PORT${NC}"
        echo -e "    Backend:   ${CYAN}http://localhost:$BACKEND_PORT${NC}"
        echo -e "    API Docs:  ${CYAN}http://localhost:$BACKEND_PORT/docs${NC}"
    fi
    echo ""
}

# ============================================================================
# Command: run
# ============================================================================
cmd_run() {
    check_venv
    activate_venv
    
    local videos=()
    local download_video=false
    
    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --download-video)
                download_video=true
                shift
                ;;
            --video)
                videos+=("$2")
                shift 2
                ;;
            --videos)
                shift
                while [[ $# -gt 0 ]] && [[ ! "$1" =~ ^-- ]]; do
                    videos+=("$1")
                    shift
                done
                ;;
            -*)
                log_error "Unknown option: $1"
                exit 1
                ;;
            *)
                videos+=("$1")
                shift
                ;;
        esac
    done
    
    if [ ${#videos[@]} -eq 0 ]; then
        log_error "No video URL provided"
        echo ""
        echo "Usage:"
        echo "  ./setup.sh run <URL>"
        echo "  ./setup.sh run --video <URL>"
        echo "  ./setup.sh run --videos <URL1> <URL2>"
        exit 1
    fi
    
    log_header "Processing ${#videos[@]} video(s)"
    echo ""
    
    for url in "${videos[@]}"; do
        log_info "Processing: $url"
        
        if [ "$download_video" = true ]; then
            python "$PROJECT_DIR/main.py" "$url" --download-video
        else
            python "$PROJECT_DIR/main.py" "$url"
        fi
        
        log_success "Complete: $url"
        echo ""
    done
    
    log_header "All videos processed!"
    echo ""
    echo -e "  View results: ${CYAN}http://localhost:$FRONTEND_PORT${NC}"
    echo ""
}

# ============================================================================
# Command: logs
# ============================================================================
cmd_logs() {
    local target=${1:-all}
    
    case $target in
        backend)
            tail -f "$BACKEND_LOG"
            ;;
        frontend)
            tail -f "$FRONTEND_LOG"
            ;;
        all|*)
            # Use tmux or show both
            if command -v tmux &> /dev/null; then
                tmux new-session "tail -f $BACKEND_LOG" \; split-window "tail -f $FRONTEND_LOG" \; select-layout even-vertical
            else
                log_info "Backend log: $BACKEND_LOG"
                log_info "Frontend log: $FRONTEND_LOG"
                echo ""
                echo "Use:"
                echo "  ./setup.sh logs backend"
                echo "  ./setup.sh logs frontend"
            fi
            ;;
    esac
}

# ============================================================================
# Command: clean
# ============================================================================
cmd_clean() {
    log_header "Cleaning up"
    echo ""
    
    # Stop services first
    cmd_stop_quiet
    
    # Remove PID files
    rm -f "$BACKEND_PID_FILE" "$FRONTEND_PID_FILE"
    log_success "Removed PID files"
    
    # Remove log files
    rm -f "$BACKEND_LOG" "$FRONTEND_LOG"
    log_success "Removed log files"
    
    # Remove Next.js cache (optional)
    if [ -d "$FRONTEND_DIR/.next" ]; then
        rm -rf "$FRONTEND_DIR/.next"
        log_success "Removed Next.js cache"
    fi
    
    echo ""
    log_success "Cleanup complete"
    echo ""
}

# ============================================================================
# Command: help
# ============================================================================
cmd_help() {
    echo ""
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║       Diarization Pipeline - CLI Tool                        ║${NC}"
    echo -e "${CYAN}║       Branch: feature/visualization                          ║${NC}"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "${BOLD}COMMANDS:${NC}"
    echo ""
    echo -e "  ${GREEN}setup${NC}        Install dependencies and configure environment"
    echo -e "  ${GREEN}start${NC}        Start visualizer (backend + frontend)"
    echo -e "  ${GREEN}stop${NC}         Stop all visualizer services"
    echo -e "  ${GREEN}status${NC}       Show service status and health"
    echo -e "  ${GREEN}run${NC}          Process video(s) through the pipeline"
    echo -e "  ${GREEN}logs${NC}         View service logs (backend|frontend|all)"
    echo -e "  ${GREEN}clean${NC}        Stop services and clean up temp files"
    echo ""
    echo -e "${BOLD}EXAMPLES:${NC}"
    echo ""
    echo -e "  ${YELLOW}# Initial setup${NC}"
    echo -e "  ./setup.sh setup"
    echo ""
    echo -e "  ${YELLOW}# Start visualizer${NC}"
    echo -e "  ./setup.sh start"
    echo ""
    echo -e "  ${YELLOW}# Check status${NC}"
    echo -e "  ./setup.sh status"
    echo ""
    echo -e "  ${YELLOW}# Process a video${NC}"
    echo -e "  ./setup.sh run https://www.youtube.com/watch?v=VIDEO_ID"
    echo ""
    echo -e "  ${YELLOW}# Process multiple videos${NC}"
    echo -e "  ./setup.sh run --videos URL1 URL2 URL3"
    echo ""
    echo -e "  ${YELLOW}# View logs${NC}"
    echo -e "  ./setup.sh logs backend"
    echo ""
    echo -e "  ${YELLOW}# Stop everything${NC}"
    echo -e "  ./setup.sh stop"
    echo ""
    echo -e "${BOLD}URLS (when running):${NC}"
    echo ""
    echo -e "  Frontend:  ${CYAN}http://localhost:$FRONTEND_PORT${NC}"
    echo -e "  Backend:   ${CYAN}http://localhost:$BACKEND_PORT${NC}"
    echo -e "  API Docs:  ${CYAN}http://localhost:$BACKEND_PORT/docs${NC}"
    echo ""
}

# ============================================================================
# Main
# ============================================================================
case "${1:-}" in
    setup)
        cmd_setup
        ;;
    start|start-viz)
        cmd_start
        ;;
    stop|stop-viz)
        cmd_stop
        ;;
    status)
        cmd_status
        ;;
    run)
        shift
        cmd_run "$@"
        ;;
    logs)
        shift
        cmd_logs "$@"
        ;;
    clean)
        cmd_clean
        ;;
    --help|-h|help)
        cmd_help
        ;;
    "")
        cmd_help
        ;;
    *)
        log_error "Unknown command: $1"
        echo ""
        cmd_help
        exit 1
        ;;
esac