#!/bin/bash

# =============================================================================
# Media Manager - First-Time Setup Script
# =============================================================================
# This script guides you through the initial setup of the Media Manager
# application, including environment configuration and database initialization.
# =============================================================================

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# Helper functions
print_header() {
    echo ""
    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}${CYAN}  $1${NC}"
    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo ""
}

print_step() {
    echo -e "${GREEN}▶${NC} $1"
}

print_info() {
    echo -e "${CYAN}ℹ${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}⚠${NC} $1"
}

print_error() {
    echo -e "${RED}✗${NC} $1"
}

print_success() {
    echo -e "${GREEN}✓${NC} $1"
}

prompt_with_default() {
    local prompt="$1"
    local default="$2"
    local var_name="$3"
    local is_password="$4"

    if [ "$is_password" = "true" ]; then
        echo -en "${CYAN}?${NC} ${prompt} [${default}]: "
        read -s value
        echo ""
    else
        echo -en "${CYAN}?${NC} ${prompt} [${default}]: "
        read value
    fi

    if [ -z "$value" ]; then
        eval "$var_name=\"$default\""
    else
        eval "$var_name=\"$value\""
    fi
}

prompt_yes_no() {
    local prompt="$1"
    local default="$2"

    if [ "$default" = "y" ]; then
        echo -en "${CYAN}?${NC} ${prompt} [Y/n]: "
    else
        echo -en "${CYAN}?${NC} ${prompt} [y/N]: "
    fi

    read answer
    answer=${answer:-$default}

    case "$answer" in
        [Yy]* ) return 0;;
        * ) return 1;;
    esac
}

wait_for_container() {
    local container="$1"
    local max_attempts="${2:-30}"
    local attempt=1

    echo -n "  Waiting for $container"
    while [ $attempt -le $max_attempts ]; do
        if docker compose ps "$container" 2>/dev/null | grep -q "healthy\|running"; then
            echo -e " ${GREEN}ready${NC}"
            return 0
        fi
        echo -n "."
        sleep 2
        attempt=$((attempt + 1))
    done

    echo -e " ${RED}timeout${NC}"
    return 1
}

wait_for_postgres() {
    local max_attempts="${1:-30}"
    local attempt=1

    echo -n "  Waiting for PostgreSQL to be healthy"
    while [ $attempt -le $max_attempts ]; do
        if docker compose exec -T postgres pg_isready -U mediamanager -d library >/dev/null 2>&1; then
            echo -e " ${GREEN}ready${NC}"
            return 0
        fi
        echo -n "."
        sleep 2
        attempt=$((attempt + 1))
    done

    echo -e " ${RED}timeout${NC}"
    return 1
}

wait_for_api() {
    local max_attempts="${1:-30}"
    local attempt=1

    echo -n "  Waiting for API to be healthy"
    while [ $attempt -le $max_attempts ]; do
        if curl -s http://localhost:3001/health >/dev/null 2>&1; then
            echo -e " ${GREEN}ready${NC}"
            return 0
        fi
        echo -n "."
        sleep 2
        attempt=$((attempt + 1))
    done

    echo -e " ${RED}timeout${NC}"
    return 1
}

# =============================================================================
# Main Script
# =============================================================================

clear
echo ""
echo -e "${BOLD}${CYAN}"
echo "  ╔══════════════════════════════════════════════════════════════════╗"
echo "  ║                                                                  ║"
echo "  ║              Media Manager - First-Time Setup                    ║"
echo "  ║                                                                  ║"
echo "  ╚══════════════════════════════════════════════════════════════════╝"
echo -e "${NC}"
echo ""
echo "  This script will guide you through the initial setup process."
echo "  Press Ctrl+C at any time to cancel."
echo ""

# =============================================================================
# Prerequisites Check
# =============================================================================

print_header "Checking Prerequisites"

# Check Docker
if command -v docker &> /dev/null; then
    DOCKER_VERSION=$(docker --version | cut -d' ' -f3 | tr -d ',')
    print_success "Docker installed (v$DOCKER_VERSION)"
else
    print_error "Docker is not installed"
    echo "  Please install Docker: https://docs.docker.com/get-docker/"
    exit 1
fi

# Check Docker Compose
if docker compose version &> /dev/null; then
    COMPOSE_VERSION=$(docker compose version --short)
    print_success "Docker Compose installed (v$COMPOSE_VERSION)"
else
    print_error "Docker Compose is not installed"
    echo "  Please install Docker Compose: https://docs.docker.com/compose/install/"
    exit 1
fi

# Check if Docker daemon is running
if docker info &> /dev/null; then
    print_success "Docker daemon is running"
else
    print_error "Docker daemon is not running"
    echo "  Please start Docker and try again"
    exit 1
fi

# Check for NVIDIA GPU (optional)
if command -v nvidia-smi &> /dev/null; then
    if nvidia-smi &> /dev/null; then
        GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -n1)
        print_success "NVIDIA GPU detected: $GPU_NAME"
        HAS_GPU=true
    else
        print_warning "NVIDIA driver not responding"
        HAS_GPU=false
    fi
else
    print_info "No NVIDIA GPU detected (optional - CPU encoding will be used)"
    HAS_GPU=false
fi

# =============================================================================
# Check Existing Configuration
# =============================================================================

if [ -f ".env" ]; then
    echo ""
    print_warning "An existing .env file was found"
    if prompt_yes_no "Do you want to overwrite it?" "n"; then
        print_info "Backing up existing .env to .env.backup"
        cp .env .env.backup
    else
        print_info "Keeping existing .env file"
        SKIP_ENV_CREATION=true
    fi
fi

# =============================================================================
# Environment Configuration
# =============================================================================

if [ "$SKIP_ENV_CREATION" != "true" ]; then

    print_header "Required Configuration"

    # Generate JWT secret
    JWT_SECRET=$(openssl rand -hex 32 2>/dev/null || cat /dev/urandom | tr -dc 'a-f0-9' | fold -w 64 | head -n 1)
    print_info "Generated secure JWT secret"

    # Generate PostgreSQL password
    POSTGRES_PASSWORD=$(openssl rand -base64 24 2>/dev/null || cat /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 32 | head -n 1)
    print_info "Generated secure PostgreSQL password"

    echo ""
    echo "  Enter the initial admin account credentials."
    echo "  You will use these to log into the admin dashboard."
    echo ""

    prompt_with_default "Admin email" "admin@localhost" ADMIN_EMAIL
    prompt_with_default "Admin password" "changeme" ADMIN_PASSWORD "true"

    # =============================================================================
    # Media Paths
    # =============================================================================

    print_header "Media Library Paths"

    echo "  Configure where your media files are stored."
    echo "  Use absolute paths. Leave blank to use defaults."
    echo ""

    # Try to detect current directory structure
    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    DEFAULT_MEDIA_ROOT="$(dirname "$SCRIPT_DIR")/media"

    prompt_with_default "Media root directory" "$DEFAULT_MEDIA_ROOT" MEDIA_ROOT
    prompt_with_default "Local media directory" "$MEDIA_ROOT/local" MEDIA_LOCAL
    prompt_with_default "Public media directory" "$MEDIA_ROOT/public" MEDIA_PUBLIC
    prompt_with_default "Studios directory" "$MEDIA_LOCAL/studios" STUDIOS_PATH
    prompt_with_default "GIFs directory" "$MEDIA_LOCAL/gifs" GIFS_PATH
    prompt_with_default "Playback output directory" "$MEDIA_PUBLIC/playback" PLAYBACK_PATH
    prompt_with_default "Compilations directory" "$MEDIA_PUBLIC/compilations" COMPILATIONS_PATH
    prompt_with_default "Curated content directory" "$MEDIA_PUBLIC/curated" PUBLIC_CURATED_PATH

    # =============================================================================
    # Optional: SMTP Configuration
    # =============================================================================

    print_header "Email Configuration (Optional)"

    echo "  Email is used for user verification and password reset."
    echo "  Skip this section to use admin-approval workflow instead."
    echo ""

    if prompt_yes_no "Configure SMTP email?" "n"; then
        prompt_with_default "SMTP host" "smtp.gmail.com" SMTP_HOST
        prompt_with_default "SMTP port" "587" SMTP_PORT
        prompt_with_default "SMTP username" "" SMTP_USER
        prompt_with_default "SMTP password" "" SMTP_PASS "true"
        prompt_with_default "From email address" "noreply@localhost" SMTP_FROM
        prompt_with_default "From name" "Media Manager" SMTP_FROM_NAME
        SMTP_SECURE="false"
    else
        SMTP_HOST=""
        SMTP_PORT=""
        SMTP_USER=""
        SMTP_PASS=""
        SMTP_FROM=""
        SMTP_FROM_NAME=""
        SMTP_SECURE=""
        print_info "Skipping SMTP configuration"
    fi

    # =============================================================================
    # Optional: AI Features
    # =============================================================================

    print_header "AI Features Configuration"

    if [ "$HAS_GPU" = "true" ]; then
        echo "  All AI services will start automatically with the application."
        echo "  These use your NVIDIA GPU for accelerated processing."
        echo ""
        echo "  Services included:"
        echo "    - Ollama (vision + text models for video analysis)"
        echo "    - Whisper (speech transcription)"
        echo "    - JoyCaption (high-quality image captioning)"
        echo "    - Face Recognition (performer detection)"
        echo "    - YOLO Detection (person counting)"
        echo "    - JoyTag (visual tagging)"
        echo ""

        # Content Safety
        echo -e "  ${BOLD}Content Safety Moderation${NC}"
        echo "  AI-powered comment moderation using Llama Guard."
        echo "  Screens user comments for harmful content."
        echo ""

        if prompt_yes_no "Enable content safety moderation?" "y"; then
            SAFETY_CHECK_ENABLED="true"
            print_info "Content safety will be enabled"
        else
            SAFETY_CHECK_ENABLED="false"
            print_info "Content safety disabled (can be enabled later in .env)"
        fi

        # JoyCaption model check
        echo ""
        echo -e "  ${BOLD}JoyCaption Vision Backend${NC}"
        echo "  Higher quality image captioning using JoyCaption GGUF model."
        echo "  Requires model files in ./models directory (~4.6GB download)"
        echo ""

        if prompt_yes_no "Will you download JoyCaption model files?" "y"; then
            DIGEST_VISION_BACKEND="llama-server"
            print_info "JoyCaption will be the primary vision backend"
            print_warning "Download model files to ./models/ before using"
            echo "  Download from: https://huggingface.co/Mungert/llama-joycaption-beta-one-hf-llava-GGUF"
        else
            DIGEST_VISION_BACKEND="ollama"
            print_info "Using Ollama for vision (default)"
        fi
    else
        echo "  AI features require an NVIDIA GPU."
        echo "  AI services will start but may not function properly without GPU."
        echo ""
        SAFETY_CHECK_ENABLED="false"
        DIGEST_VISION_BACKEND="ollama"
        print_warning "No GPU detected - AI features may be limited"
    fi

    # =============================================================================
    # Application Settings
    # =============================================================================

    print_header "Application Settings"

    prompt_with_default "Public gallery URL" "http://localhost:3080" APP_BASE_URL
    prompt_with_default "Timezone" "UTC" TZ

    # =============================================================================
    # Generate .env File
    # =============================================================================

    print_header "Generating Configuration"

    print_step "Creating .env file..."

    cat > .env << EOF
# =============================================================================
# Media Manager Configuration
# Generated by build.sh on $(date)
# =============================================================================

# PostgreSQL Database (REQUIRED)
POSTGRES_PASSWORD=$POSTGRES_PASSWORD

# JWT Authentication
JWT_SECRET=$JWT_SECRET

# Initial Admin Account
INITIAL_ADMIN_EMAIL=$ADMIN_EMAIL
INITIAL_ADMIN_PASSWORD=$ADMIN_PASSWORD

# SMTP Configuration
SMTP_HOST=$SMTP_HOST
SMTP_PORT=$SMTP_PORT
SMTP_SECURE=${SMTP_SECURE:-false}
SMTP_USER=$SMTP_USER
SMTP_PASS=$SMTP_PASS
SMTP_FROM=$SMTP_FROM
SMTP_FROM_NAME=$SMTP_FROM_NAME

# Email Settings
EMAIL_VERIFICATION_EXPIRY_HOURS=24
PASSWORD_RESET_EXPIRY_HOURS=1

# Application Settings
APP_BASE_URL=$APP_BASE_URL
ALLOWED_ORIGINS=http://localhost:3080,http://localhost:8080
TZ=$TZ
VITE_DEFAULT_LOCALE=

# Content Safety
SAFETY_CHECK_ENABLED=$SAFETY_CHECK_ENABLED
OLLAMA_BASE_URL=http://ollama:11434
OLLAMA_MODEL=llama-guard3:1b

# Digest Feature Models (for AI video analysis)
DIGEST_VISION_MODEL=huihui_ai/qwen3-vl-abliterated:2b
DIGEST_TEXT_MODEL=huihui_ai/qwen3-abliterated:4b
DIGEST_OLLAMA_URL=http://ollama:11434
DIGEST_WHISPER_URL=http://whisper:5000
DIGEST_FRAME_INTERVAL=30

# JoyCaption (alternative vision backend)
DIGEST_VISION_BACKEND=$DIGEST_VISION_BACKEND
JOYCAPTION_URL=http://joycaption:8080
JOYCAPTION_MODEL=llama-joycaption-beta-one-hf-llava
JOYCAPTION_FALLBACK_TO_OLLAMA=true
JOYCAPTION_MODEL_PATH=/models/llama-joycaption-beta-one-hf-llava.Q2_K.gguf
JOYCAPTION_MMPROJ_PATH=/models/llama-joycaption-beta-one-llava-mmproj-model-f16.gguf
JOYCAPTION_CTX_SIZE=4096
JOYCAPTION_GPU_LAYERS=999
JOYCAPTION_PARALLEL=1

# Whisper Transcription
WHISPER_MODEL=base
WHISPER_DEVICE=cuda
WHISPER_COMPUTE=float16

# Ollama Auto-Pull Models
OLLAMA_MODELS=llama-guard3:1b,huihui_ai/qwen3-vl-abliterated:2b,huihui_ai/qwen3-abliterated:4b

# Media Library Paths
MEDIA_ROOT=$MEDIA_ROOT
MEDIA_LOCAL=$MEDIA_LOCAL
MEDIA_PUBLIC=$MEDIA_PUBLIC
STUDIOS_PATH=$STUDIOS_PATH
ONLYFANS_PATH=$MEDIA_LOCAL/OnlyFans
GIFS_PATH=$GIFS_PATH
PLAYBACK_PATH=$PLAYBACK_PATH
COMPILATIONS_PATH=$COMPILATIONS_PATH
PUBLIC_CURATED_PATH=$PUBLIC_CURATED_PATH
QUICKIES_PATH=$MEDIA_PUBLIC/quickies
EOF

    print_success ".env file created"
fi

# =============================================================================
# Create Docker Network
# =============================================================================

print_header "Setting Up Docker"

print_step "Checking Docker network..."

if docker network inspect media-network &> /dev/null; then
    print_success "Docker network 'media-network' already exists"
else
    print_step "Creating Docker network 'media-network'..."
    docker network create media-network
    print_success "Docker network created"
fi

# =============================================================================
# Create Required Directories
# =============================================================================

print_step "Creating data directories..."

mkdir -p data/thumbnails
mkdir -p data/digest_frames
mkdir -p data/ads
mkdir -p data/gallery
mkdir -p models

print_success "Data directories created"

# =============================================================================
# Start Services
# =============================================================================

print_header "Starting Services"

print_step "Pulling Docker images (this may take a while)..."

docker compose pull

# Check for JoyCaption model files
print_step "Checking JoyCaption model files..."
if [ -f "models/llama-joycaption-beta-one-hf-llava-Q3_K_M.gguf" ] && \
   [ -f "models/llama-joycaption-beta-one-llava-mmproj-model-f16.gguf" ]; then
    print_success "JoyCaption model files found"
else
    print_warning "JoyCaption model files not found in ./models/"
    echo ""
    echo "  Download these files to enable JoyCaption vision backend:"
    echo "  1. llama-joycaption-beta-one-hf-llava-Q3_K_M.gguf (~4GB)"
    echo "  2. llama-joycaption-beta-one-llava-mmproj-model-f16.gguf (~600MB)"
    echo ""
    echo "  From: https://huggingface.co/Mungert/llama-joycaption-beta-one-hf-llava-GGUF"
    echo ""
    print_info "JoyCaption container will start but won't work until models are downloaded"
fi

print_step "Starting containers..."

# Start core services first
docker compose up -d

# Ask about AI services
echo ""
if [ "$HAS_GPU" = "true" ]; then
    if prompt_yes_no "Start AI services now? (can be started later with: docker compose --profile ai up -d)" "y"; then
        print_step "Starting AI services..."
        docker compose --profile ai up -d
        START_AI_SERVICES=true
    else
        START_AI_SERVICES=false
        print_info "AI services not started. Start later with: docker compose --profile ai up -d"
    fi
else
    START_AI_SERVICES=false
    print_info "Skipping AI services (no GPU detected)"
fi

# =============================================================================
# Wait for Services
# =============================================================================

print_header "Waiting for Services"

print_step "Checking service health..."

if ! wait_for_postgres 60; then
    print_error "PostgreSQL failed to start"
    echo ""
    echo "  Check the logs with: docker compose logs postgres"
    exit 1
fi

# Wait for Ollama to start and pull models (only if AI services were started)
if [ "$START_AI_SERVICES" = "true" ]; then
    print_step "Waiting for Ollama to start and pull models..."
    echo "  (This may take several minutes on first run as models are downloaded)"
    echo ""

    # Wait for Ollama to be responsive (not necessarily fully ready with models)
    MAX_OLLAMA_WAIT=30
    OLLAMA_ATTEMPT=1
    echo -n "  Waiting for Ollama"
    while [ $OLLAMA_ATTEMPT -le $MAX_OLLAMA_WAIT ]; do
        if curl -s http://localhost:11435/ >/dev/null 2>&1; then
            echo -e " ${GREEN}ready${NC}"
            break
        fi
        echo -n "."
        sleep 2
        OLLAMA_ATTEMPT=$((OLLAMA_ATTEMPT + 1))
    done

    if [ $OLLAMA_ATTEMPT -gt $MAX_OLLAMA_WAIT ]; then
        print_warning "Ollama is still starting (models may still be downloading)"
        echo "  Check progress with: docker compose --profile ai logs -f ollama"
    else
        print_info "Ollama is running. Models will be pulled in the background."
        echo "  Monitor model downloads with: docker compose --profile ai logs -f ollama"
    fi
fi

# =============================================================================
# Initialize Database
# =============================================================================

print_header "Initializing Database"

print_step "Pushing database schema..."

if docker compose exec -T api npx drizzle-kit push 2>&1 | tail -5; then
    print_success "Database schema initialized"
else
    print_error "Failed to initialize database schema"
    echo ""
    echo "  Check the logs with: docker compose logs api"
    exit 1
fi

# =============================================================================
# Restart API to Create Admin User
# =============================================================================

print_step "Restarting API to create admin user..."

docker compose restart api

sleep 3

if ! wait_for_api 30; then
    print_warning "API health check timed out, but it may still be starting"
fi

# Check if admin was created
print_step "Verifying admin user creation..."

sleep 2

ADMIN_CHECK=$(docker compose logs api 2>&1 | grep -i "initial admin" | tail -1)
if echo "$ADMIN_CHECK" | grep -q "created"; then
    print_success "Admin user created successfully"
elif echo "$ADMIN_CHECK" | grep -q "exists"; then
    print_info "Admin user already exists"
else
    print_warning "Could not verify admin creation (check logs if login fails)"
fi

# =============================================================================
# Seed Tags
# =============================================================================

print_step "Checking tag seeding..."

TAG_CHECK=$(docker compose logs api 2>&1 | grep -i "seed" | tail -1)
if echo "$TAG_CHECK" | grep -q "Seeded"; then
    print_success "Default tags seeded"
elif echo "$TAG_CHECK" | grep -q "already"; then
    print_info "Tags already exist"
fi

# =============================================================================
# Final Status
# =============================================================================

print_header "Setup Complete!"

echo -e "  ${GREEN}${BOLD}Your Media Manager is ready!${NC}"
echo ""
echo "  Access the application at:"
echo ""
echo -e "    ${CYAN}Admin Dashboard:${NC}  http://localhost:8080"
echo -e "    ${CYAN}Public Gallery:${NC}   http://localhost:3080"
echo -e "    ${CYAN}API:${NC}              http://localhost:3001"

# Show AI features status
if [ "$HAS_GPU" = "true" ]; then
    echo ""
    if [ "$START_AI_SERVICES" = "true" ]; then
        echo "  AI Services (running):"
        echo -e "    ${GREEN}✓${NC} Ollama (Vision + Text models)"
        echo -e "    ${GREEN}✓${NC} Whisper (Speech transcription)"
        echo -e "    ${GREEN}✓${NC} JoyCaption (Image captioning)"
        echo -e "    ${GREEN}✓${NC} Face Recognition"
        echo -e "    ${GREEN}✓${NC} YOLO Detection"
        echo -e "    ${GREEN}✓${NC} JoyTag (Visual tagging)"
        if [ "$SAFETY_CHECK_ENABLED" = "true" ]; then
            echo -e "    ${GREEN}✓${NC} Content Safety Moderation"
        fi
        echo ""
        echo -e "    ${CYAN}Ollama:${NC}           http://localhost:11435"
        echo -e "    ${CYAN}JoyCaption:${NC}       http://localhost:11436"
    else
        echo "  AI Services (not started):"
        echo -e "    ${YELLOW}○${NC} Ollama, Whisper, JoyCaption, Face Recognition, YOLO, JoyTag"
        echo ""
        echo -e "    Start with: ${YELLOW}docker compose --profile ai up -d${NC}"
    fi
fi

echo ""
echo "  Login credentials:"
echo ""
echo -e "    ${CYAN}Email:${NC}     $ADMIN_EMAIL"
echo -e "    ${CYAN}Password:${NC}  (the password you entered during setup)"
echo ""
echo "  Useful commands:"
echo ""
echo -e "    ${YELLOW}docker compose logs -f${NC}                    # View core service logs"
echo -e "    ${YELLOW}docker compose down${NC}                       # Stop core services"
echo -e "    ${YELLOW}docker compose up -d${NC}                      # Start core services"
echo -e "    ${YELLOW}docker compose restart api${NC}                # Restart API"
echo ""
echo -e "    ${YELLOW}docker compose --profile ai up -d${NC}         # Start AI services"
echo -e "    ${YELLOW}docker compose --profile ai down${NC}          # Stop all (including AI)"
echo -e "    ${YELLOW}docker compose --profile ai logs -f ollama${NC} # View Ollama logs"


echo ""
echo "  Next steps:"
echo ""
echo "    1. Log into the admin dashboard"
echo "    2. Click 'Scan Library' to index your videos"
echo "    3. Use 'Fractal' section to analyze videos with AI"
echo "    4. Explore the Library, Jobs, and other sections"
echo ""
echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo ""