#!/usr/bin/env bash
#
# Path: setup_voice_assistant.sh
#
# Purpose and usage:
#     Sets up the voice assistant server environment on Heimdall
#     - Creates conda environment
#     - Installs dependencies (Whisper, Flask, Piper TTS)
#     - Downloads and configures TTS models
#     - Sets up systemd service (optional)
#     - Configures environment variables
#
# Requirements:
#     - conda/miniconda installed
#     - Internet connection for downloads
#     - Sudo access (for systemd service setup)
#
# Usage:
#     ./setup_voice_assistant.sh [--no-service] [--env-name NAME]
#
# Author: PRbL Library
# Created: $(date +"%Y-%m-%d")

# ----- PRbL Color and output functions -----
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

print_status() {
    local level="$1"
    shift
    case "$level" in
        "info")     echo -e "${BLUE}[INFO]${NC} $*" >&2 ;;
        "success")  echo -e "${GREEN}[SUCCESS]${NC} $*" >&2 ;;
        "warning")  echo -e "${YELLOW}[WARNING]${NC} $*" >&2 ;;
        "error")    echo -e "${RED}[ERROR]${NC} $*" >&2 ;;
        "debug")    [[ "$VERBOSE" == "true" ]] && echo -e "${PURPLE}[DEBUG]${NC} $*" >&2 ;;
        *)          echo -e "$*" >&2 ;;
    esac
}

# ----- Configuration -----
CONDA_ENV_NAME="voice-assistant"
PROJECT_DIR="$HOME/voice-assistant"
INSTALL_SYSTEMD=true
VERBOSE=false

# ----- Dependency checking -----
command_exists() {
    command -v "$1" &> /dev/null
}

check_conda() {
    if ! command_exists conda; then
        print_status error "conda not found. Please install miniconda first."
        print_status info "Install with: wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh"
        print_status info "             bash Miniconda3-latest-Linux-x86_64.sh"
        return 1
    fi
    return 0
}

# ----- Parse arguments -----
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --no-service)
                INSTALL_SYSTEMD=false
                shift
                ;;
            --env-name)
                CONDA_ENV_NAME="$2"
                shift 2
                ;;
            -v|--verbose)
                VERBOSE=true
                shift
                ;;
            -h|--help)
                cat << EOF
Usage: $(basename "$0") [OPTIONS]

Options:
    --no-service        Don't install systemd service
    --env-name NAME     Custom conda environment name (default: voice-assistant)
    -v, --verbose       Enable verbose output
    -h, --help          Show this help message

EOF
                exit 0
                ;;
            *)
                print_status error "Unknown option: $1"
                exit 1
                ;;
        esac
    done
}

# ----- Setup functions -----

create_project_directory() {
    print_status info "Creating project directory: $PROJECT_DIR"
    
    if [[ ! -d "$PROJECT_DIR" ]]; then
        mkdir -p "$PROJECT_DIR" || {
            print_status error "Failed to create project directory"
            return 1
        }
    fi
    
    # Create subdirectories
    mkdir -p "$PROJECT_DIR"/{logs,models,config}
    
    print_status success "Project directory created"
    return 0
}

create_conda_environment() {
    print_status info "Creating conda environment: $CONDA_ENV_NAME"
    
    # Check if environment already exists
    if conda env list | grep -q "^${CONDA_ENV_NAME}\s"; then
        print_status warning "Environment $CONDA_ENV_NAME already exists"
        read -p "Remove and recreate? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            print_status info "Removing existing environment..."
            conda env remove -n "$CONDA_ENV_NAME" -y
        else
            print_status info "Using existing environment"
            return 0
        fi
    fi
    
    # Create new environment
    print_status info "Creating Python 3.10 environment..."
    conda create -n "$CONDA_ENV_NAME" python=3.10 -y || {
        print_status error "Failed to create conda environment"
        return 1
    }
    
    print_status success "Conda environment created"
    return 0
}

install_python_dependencies() {
    print_status info "Installing Python dependencies..."
    
    # Activate conda environment
    eval "$(conda shell.bash hook)"
    conda activate "$CONDA_ENV_NAME" || {
        print_status error "Failed to activate conda environment"
        return 1
    }
    
    # Install base dependencies
    print_status info "Installing base packages..."
    pip install --upgrade pip --break-system-packages || true
    
    # Install Whisper (OpenAI)
    print_status info "Installing OpenAI Whisper..."
    pip install -U openai-whisper --break-system-packages || {
        print_status error "Failed to install Whisper"
        return 1
    }
    
    # Install Flask
    print_status info "Installing Flask..."
    pip install flask --break-system-packages || {
        print_status error "Failed to install Flask"
        return 1
    }
    
    # Install requests
    print_status info "Installing requests..."
    pip install requests --break-system-packages || {
        print_status error "Failed to install requests"
        return 1
    }
    
    # Install python-dotenv
    print_status info "Installing python-dotenv..."
    pip install python-dotenv --break-system-packages || {
        print_status warning "Failed to install python-dotenv (optional)"
    }
    
    # Install Piper TTS
    print_status info "Installing Piper TTS..."
    # Note: Piper TTS installation method varies, adjust as needed
    # For now, we'll install the Python package if available
    pip install piper-tts --break-system-packages || {
        print_status warning "Piper TTS pip package not found"
        print_status info "You may need to install Piper manually from: https://github.com/rhasspy/piper"
    }
    
    # Install PyAudio for audio handling
    print_status info "Installing PyAudio dependencies..."
    if command_exists apt-get; then
        sudo apt-get install -y portaudio19-dev python3-pyaudio || {
            print_status warning "Failed to install portaudio dev packages"
        }
    fi
    
    pip install pyaudio --break-system-packages || {
        print_status warning "Failed to install PyAudio (may need manual installation)"
    }
    
    print_status success "Python dependencies installed"
    return 0
}

download_piper_models() {
    print_status info "Downloading Piper TTS models..."
    
    local models_dir="$PROJECT_DIR/models/piper"
    mkdir -p "$models_dir"
    
    # Download a default voice model
    # Example: en_US-lessac-medium
    local model_url="https://huggingface.co/rhasspy/piper-voices/resolve/main/en/en_US/lessac/medium/en_US-lessac-medium.onnx"
    local config_url="https://huggingface.co/rhasspy/piper-voices/resolve/main/en/en_US/lessac/medium/en_US-lessac-medium.onnx.json"
    
    if [[ ! -f "$models_dir/en_US-lessac-medium.onnx" ]]; then
        print_status info "Downloading voice model..."
        wget -q --show-progress -O "$models_dir/en_US-lessac-medium.onnx" "$model_url" || {
            print_status warning "Failed to download Piper model (manual download may be needed)"
        }
        
        wget -q --show-progress -O "$models_dir/en_US-lessac-medium.onnx.json" "$config_url" || {
            print_status warning "Failed to download Piper config"
        }
    else
        print_status info "Piper model already downloaded"
    fi
    
    print_status success "Piper models ready"
    return 0
}

create_config_file() {
    print_status info "Creating configuration file..."
    
    local config_file="$PROJECT_DIR/config/.env"
    
    if [[ -f "$config_file" ]]; then
        print_status warning "Config file already exists: $config_file"
        return 0
    fi
    
    cat > "$config_file" << 'EOF'
# Voice Assistant Configuration
# Path: ~/voice-assistant/config/.env

# Home Assistant Configuration
HA_URL=http://homeassistant.local:8123
HA_TOKEN=your_long_lived_access_token_here

# Server Configuration
SERVER_HOST=0.0.0.0
SERVER_PORT=5000

# Whisper Configuration
WHISPER_MODEL=medium

# Piper TTS Configuration
PIPER_MODEL=/path/to/piper/model.onnx
PIPER_CONFIG=/path/to/piper/model.onnx.json

# Logging
LOG_LEVEL=INFO
LOG_FILE=/home/$USER/voice-assistant/logs/voice_assistant.log
EOF
    
    # Update paths in config
    sed -i "s|/path/to/piper/model.onnx|$PROJECT_DIR/models/piper/en_US-lessac-medium.onnx|g" "$config_file"
    sed -i "s|/path/to/piper/model.onnx.json|$PROJECT_DIR/models/piper/en_US-lessac-medium.onnx.json|g" "$config_file"
    sed -i "s|/home/\$USER|$HOME|g" "$config_file"
    
    chmod 600 "$config_file"
    
    print_status success "Config file created: $config_file"
    print_status warning "Please edit $config_file and add your Home Assistant token"
    
    return 0
}

create_systemd_service() {
    if [[ "$INSTALL_SYSTEMD" != "true" ]]; then
        print_status info "Skipping systemd service installation"
        return 0
    fi
    
    print_status info "Creating systemd service..."
    
    local service_file="/etc/systemd/system/voice-assistant.service"
    
    # Create service file
    sudo tee "$service_file" > /dev/null << EOF
[Unit]
Description=Voice Assistant Server
After=network.target

[Service]
Type=simple
User=$USER
WorkingDirectory=$PROJECT_DIR
Environment="PATH=$HOME/miniconda3/envs/$CONDA_ENV_NAME/bin:/usr/local/bin:/usr/bin:/bin"
EnvironmentFile=$PROJECT_DIR/config/.env
ExecStart=$HOME/miniconda3/envs/$CONDA_ENV_NAME/bin/python $PROJECT_DIR/voice_server.py
Restart=on-failure
RestartSec=10
StandardOutput=append:$PROJECT_DIR/logs/voice_assistant.log
StandardError=append:$PROJECT_DIR/logs/voice_assistant_error.log

[Install]
WantedBy=multi-user.target
EOF
    
    # Reload systemd
    sudo systemctl daemon-reload
    
    print_status success "Systemd service created"
    print_status info "To enable and start the service:"
    print_status info "  sudo systemctl enable voice-assistant"
    print_status info "  sudo systemctl start voice-assistant"
    
    return 0
}

create_test_script() {
    print_status info "Creating test script..."
    
    local test_script="$PROJECT_DIR/test_server.sh"
    
    cat > "$test_script" << 'EOF'
#!/bin/bash
# Test script for voice assistant server

# Activate conda environment
eval "$(conda shell.bash hook)"
conda activate voice-assistant

# Load environment variables
if [[ -f ~/voice-assistant/config/.env ]]; then
    export $(grep -v '^#' ~/voice-assistant/config/.env | xargs)
fi

# Run server
cd ~/voice-assistant
python voice_server.py --verbose
EOF
    
    chmod +x "$test_script"
    
    print_status success "Test script created: $test_script"
    return 0
}

install_voice_server_script() {
    print_status info "Installing voice_server.py..."
    
    # Check if voice_server.py exists in outputs
    if [[ -f "$HOME/voice_server.py" ]]; then
        cp "$HOME/voice_server.py" "$PROJECT_DIR/voice_server.py"
        print_status success "voice_server.py installed"
    elif [[ -f "./voice_server.py" ]]; then
        cp "./voice_server.py" "$PROJECT_DIR/voice_server.py"
        print_status success "voice_server.py installed"
    else
        print_status warning "voice_server.py not found in current directory"
        print_status info "Please copy voice_server.py to $PROJECT_DIR manually"
    fi
    
    return 0
}

# ----- Main -----
main() {
    print_status info "Starting voice assistant setup..."
    
    # Parse arguments
    parse_args "$@"
    
    # Check dependencies
    check_conda || exit 1
    
    # Setup steps
    create_project_directory || exit 1
    create_conda_environment || exit 1
    install_python_dependencies || exit 1
    download_piper_models || exit 1
    create_config_file || exit 1
    install_voice_server_script || exit 1
    create_test_script || exit 1
    
    if [[ "$INSTALL_SYSTEMD" == "true" ]]; then
        create_systemd_service || exit 1
    fi
    
    # Final instructions
    print_status success "Setup complete!"
    echo
    print_status info "Next steps:"
    print_status info "1. Edit config file: vim $PROJECT_DIR/config/.env"
    print_status info "2. Add your Home Assistant long-lived access token"
    print_status info "3. Test the server: $PROJECT_DIR/test_server.sh"
    print_status info "4. Configure your Maix Duino device"
    
    if [[ "$INSTALL_SYSTEMD" == "true" ]]; then
        echo
        print_status info "To run as a service:"
        print_status info "  sudo systemctl enable voice-assistant"
        print_status info "  sudo systemctl start voice-assistant"
        print_status info "  sudo systemctl status voice-assistant"
    fi
    
    echo
    print_status info "Project directory: $PROJECT_DIR"
    print_status info "Conda environment: $CONDA_ENV_NAME"
    print_status info "Activate with: conda activate $CONDA_ENV_NAME"
}

# Run main
main "$@"