#!/bin/bash

# LLM Proxy Gateway Deployment Script
# This script automates the deployment of the LLM Proxy Gateway on a Linux server

set -e  # Exit on error
set -u  # Exit on undefined variable

# Configuration
APP_NAME="llm-proxy"
APP_USER="llmproxy"
APP_GROUP="llmproxy"
GIT_REPO="ssh://git.dustin.coffee:2222/hobokenchicken/llm-proxy.git"
INSTALL_DIR="/opt/$APP_NAME"
CONFIG_DIR="/etc/$APP_NAME"
DATA_DIR="/var/lib/$APP_NAME"
LOG_DIR="/var/log/$APP_NAME"
SERVICE_FILE="/etc/systemd/system/$APP_NAME.service"
ENV_FILE="$CONFIG_DIR/.env"

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

# Logging functions
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root"
        exit 1
    fi
}

# Install system dependencies
install_dependencies() {
    log_info "Installing system dependencies..."
    
    # Detect package manager
    if command -v apt-get &> /dev/null; then
        # Debian/Ubuntu
        apt-get update
        apt-get install -y \
            build-essential \
            pkg-config \
            libssl-dev \
            sqlite3 \
            curl \
            git
    elif command -v yum &> /dev/null; then
        # RHEL/CentOS
        yum groupinstall -y "Development Tools"
        yum install -y \
            openssl-devel \
            sqlite \
            curl \
            git
    elif command -v dnf &> /dev/null; then
        # Fedora
        dnf groupinstall -y "Development Tools"
        dnf install -y \
            openssl-devel \
            sqlite \
            curl \
            git
    elif command -v pacman &> /dev/null; then
        # Arch Linux
        pacman -Syu --noconfirm \
            base-devel \
            openssl \
            sqlite \
            curl \
            git
    else
        log_warn "Could not detect package manager. Please install dependencies manually."
    fi
}

# Install Rust if not present
install_rust() {
    log_info "Checking for Rust installation..."
    
    if ! command -v rustc &> /dev/null; then
        log_info "Installing Rust..."
        curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
        source "$HOME/.cargo/env"
    else
        log_info "Rust is already installed"
    fi
    
    # Verify installation
    rustc --version
    cargo --version
}

# Create system user and directories
setup_directories() {
    log_info "Creating system user and directories..."
    
    # Create user and group if they don't exist
    if ! id "$APP_USER" &>/dev/null; then
        # Arch uses /usr/bin/nologin, Debian/Ubuntu use /usr/sbin/nologin
        NOLOGIN=$(command -v nologin 2>/dev/null || echo "/usr/bin/nologin")
        useradd -r -s "$NOLOGIN" -M "$APP_USER"
    fi
    
    # Create directories
    mkdir -p "$INSTALL_DIR"
    mkdir -p "$CONFIG_DIR"
    mkdir -p "$DATA_DIR"
    mkdir -p "$LOG_DIR"
    
    # Set permissions
    chown -R "$APP_USER:$APP_GROUP" "$INSTALL_DIR"
    chown -R "$APP_USER:$APP_GROUP" "$CONFIG_DIR"
    chown -R "$APP_USER:$APP_GROUP" "$DATA_DIR"
    chown -R "$APP_USER:$APP_GROUP" "$LOG_DIR"
    
    chmod 750 "$INSTALL_DIR"
    chmod 750 "$CONFIG_DIR"
    chmod 750 "$DATA_DIR"
    chmod 750 "$LOG_DIR"
}

# Build the application
build_application() {
    log_info "Building the application..."
    
    # Clone or update repository
    if [[ ! -d "$INSTALL_DIR/.git" ]]; then
        log_info "Cloning repository..."
        git clone "$GIT_REPO" "$INSTALL_DIR"
    else
        log_info "Updating repository..."
        cd "$INSTALL_DIR"
        git pull
    fi
    
    # Build in release mode
    cd "$INSTALL_DIR"
    log_info "Building release binary..."
    cargo build --release
    
    # Verify build
    if [[ -f "target/release/$APP_NAME" ]]; then
        log_info "Build successful"
    else
        log_error "Build failed"
        exit 1
    fi
}

# Create configuration files
create_configuration() {
    log_info "Creating configuration files..."
    
    # Create .env file with API keys
    cat > "$ENV_FILE" << EOF
# LLM Proxy Gateway Environment Variables
# Add your API keys here

# OpenAI API Key
# OPENAI_API_KEY=sk-your-key-here

# Google Gemini API Key
# GEMINI_API_KEY=AIza-your-key-here

# DeepSeek API Key
# DEEPSEEK_API_KEY=sk-your-key-here

# xAI Grok API Key
# GROK_API_KEY=gk-your-key-here

# Authentication tokens (comma-separated)
# LLM_PROXY__SERVER__AUTH_TOKENS=token1,token2,token3
EOF
    
    # Create config.toml
    cat > "$CONFIG_DIR/config.toml" << EOF
# LLM Proxy Gateway Configuration

[server]
port = 8080
host = "0.0.0.0"
# auth_tokens = ["token1", "token2", "token3"]  # Uncomment to enable authentication

[database]
path = "$DATA_DIR/llm_proxy.db"
max_connections = 5

[providers.openai]
enabled = true
api_key_env = "OPENAI_API_KEY"
base_url = "https://api.openai.com/v1"
default_model = "gpt-4o"

[providers.gemini]
enabled = true
api_key_env = "GEMINI_API_KEY"
base_url = "https://generativelanguage.googleapis.com/v1"
default_model = "gemini-2.0-flash"

[providers.deepseek]
enabled = true
api_key_env = "DEEPSEEK_API_KEY"
base_url = "https://api.deepseek.com"
default_model = "deepseek-reasoner"

[providers.grok]
enabled = false  # Disabled by default until API is researched
api_key_env = "GROK_API_KEY"
base_url = "https://api.x.ai/v1"
default_model = "grok-beta"

[model_mapping]
"gpt-*" = "openai"
"gemini-*" = "gemini"
"deepseek-*" = "deepseek"
"grok-*" = "grok"

[pricing]
openai = { input = 0.01, output = 0.03 }
gemini = { input = 0.0005, output = 0.0015 }
deepseek = { input = 0.00014, output = 0.00028 }
grok = { input = 0.001, output = 0.003 }
EOF
    
    # Set permissions
    chown "$APP_USER:$APP_GROUP" "$ENV_FILE"
    chown "$APP_USER:$APP_GROUP" "$CONFIG_DIR/config.toml"
    chmod 640 "$ENV_FILE"
    chmod 640 "$CONFIG_DIR/config.toml"
}

# Create systemd service
create_systemd_service() {
    log_info "Creating systemd service..."
    
    cat > "$SERVICE_FILE" << EOF
[Unit]
Description=LLM Proxy Gateway
Documentation=https://git.dustin.coffee/hobokenchicken/llm-proxy
After=network.target
Wants=network.target

[Service]
Type=simple
User=$APP_USER
Group=$APP_GROUP
WorkingDirectory=$INSTALL_DIR
EnvironmentFile=$ENV_FILE
Environment="RUST_LOG=info"
Environment="LLM_PROXY__CONFIG_PATH=$CONFIG_DIR/config.toml"
Environment="LLM_PROXY__DATABASE__PATH=$DATA_DIR/llm_proxy.db"
ExecStart=$INSTALL_DIR/target/release/$APP_NAME
Restart=on-failure
RestartSec=5

# Security hardening
NoNewPrivileges=true
PrivateTmp=true
ProtectSystem=strict
ProtectHome=true
ReadWritePaths=$DATA_DIR $LOG_DIR

# Resource limits (adjust based on your server)
MemoryMax=400M
MemorySwapMax=100M
CPUQuota=50%
LimitNOFILE=65536

# Logging
StandardOutput=journal
StandardError=journal
SyslogIdentifier=$APP_NAME

[Install]
WantedBy=multi-user.target
EOF
    
    # Reload systemd
    systemctl daemon-reload
}

# Setup nginx reverse proxy (optional)
setup_nginx_proxy() {
    if ! command -v nginx &> /dev/null; then
        log_warn "nginx not installed. Skipping reverse proxy setup."
        return
    fi
    
    log_info "Setting up nginx reverse proxy..."
    
    cat > "/etc/nginx/sites-available/$APP_NAME" << EOF
server {
    listen 80;
    server_name your-domain.com;  # Change to your domain
    
    # Redirect to HTTPS (recommended)
    return 301 https://\$server_name\$request_uri;
}

server {
    listen 443 ssl http2;
    server_name your-domain.com;  # Change to your domain
    
    # SSL certificates (adjust paths)
    ssl_certificate /etc/letsencrypt/live/your-domain.com/fullchain.pem;
    ssl_certificate_key /etc/letsencrypt/live/your-domain.com/privkey.pem;
    
    # SSL configuration
    ssl_protocols TLSv1.2 TLSv1.3;
    ssl_ciphers ECDHE-RSA-AES256-GCM-SHA512:DHE-RSA-AES256-GCM-SHA512:ECDHE-RSA-AES256-GCM-SHA384:DHE-RSA-AES256-GCM-SHA384;
    ssl_prefer_server_ciphers off;
    
    # Proxy to LLM Proxy Gateway
    location / {
        proxy_pass http://127.0.0.1:8080;
        proxy_http_version 1.1;
        proxy_set_header Upgrade \$http_upgrade;
        proxy_set_header Connection "upgrade";
        proxy_set_header Host \$host;
        proxy_set_header X-Real-IP \$remote_addr;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto \$scheme;
        
        # Timeouts
        proxy_connect_timeout 60s;
        proxy_send_timeout 60s;
        proxy_read_timeout 60s;
    }
    
    # Health check endpoint
    location /health {
        proxy_pass http://127.0.0.1:8080/health;
        access_log off;
    }
    
    # Dashboard
    location /dashboard {
        proxy_pass http://127.0.0.1:8080/dashboard;
    }
}
EOF
    
    # Enable site
    ln -sf "/etc/nginx/sites-available/$APP_NAME" "/etc/nginx/sites-enabled/"
    
    # Test nginx configuration
    nginx -t
    
    log_info "nginx configuration created. Please update the domain and SSL certificate paths."
}

# Setup firewall
setup_firewall() {
    log_info "Configuring firewall..."
    
    # Check for ufw (Ubuntu)
    if command -v ufw &> /dev/null; then
        ufw allow 22/tcp  # SSH
        ufw allow 80/tcp  # HTTP
        ufw allow 443/tcp # HTTPS
        ufw --force enable
        log_info "UFW firewall configured"
    fi
    
    # Check for firewalld (RHEL/CentOS)
    if command -v firewall-cmd &> /dev/null; then
        firewall-cmd --permanent --add-service=ssh
        firewall-cmd --permanent --add-service=http
        firewall-cmd --permanent --add-service=https
        firewall-cmd --reload
        log_info "Firewalld configured"
    fi
}

# Initialize database
initialize_database() {
    log_info "Initializing database..."
    
    # Run the application once to create database
    sudo -u "$APP_USER" "$INSTALL_DIR/target/release/$APP_NAME" --help &> /dev/null || true
    
    log_info "Database initialized at $DATA_DIR/llm_proxy.db"
}

# Start and enable service
start_service() {
    log_info "Starting $APP_NAME service..."
    
    systemctl enable "$APP_NAME"
    systemctl start "$APP_NAME"
    
    # Check status
    sleep 2
    systemctl status "$APP_NAME" --no-pager
}

# Verify installation
verify_installation() {
    log_info "Verifying installation..."
    
    # Check if service is running
    if systemctl is-active --quiet "$APP_NAME"; then
        log_info "Service is running"
    else
        log_error "Service is not running"
        journalctl -u "$APP_NAME" -n 20 --no-pager
        exit 1
    fi
    
    # Test health endpoint
    if curl -s http://localhost:8080/health | grep -q "OK"; then
        log_info "Health check passed"
    else
        log_error "Health check failed"
        exit 1
    fi
    
    # Test dashboard
    if curl -s -o /dev/null -w "%{http_code}" http://localhost:8080/dashboard | grep -q "200"; then
        log_info "Dashboard is accessible"
    else
        log_warn "Dashboard may not be accessible (this is normal if not configured)"
    fi
    
    log_info "Installation verified successfully!"
}

# Print next steps
print_next_steps() {
    cat << EOF

${GREEN}=== LLM Proxy Gateway Installation Complete ===${NC}

${YELLOW}Next steps:${NC}

1. ${GREEN}Configure API keys${NC}
   Edit: $ENV_FILE
   Add your API keys for the providers you want to use

2. ${GREEN}Configure authentication${NC}
   Edit: $CONFIG_DIR/config.toml
   Uncomment and set auth_tokens for client authentication

3. ${GREEN}Configure nginx${NC}
   Edit: /etc/nginx/sites-available/$APP_NAME
   Update domain name and SSL certificate paths

4. ${GREEN}Test the API${NC}
   curl -X POST http://localhost:8080/v1/chat/completions \\
     -H "Content-Type: application/json" \\
     -H "Authorization: Bearer your-token" \\
     -d '{
       "model": "gpt-4o",
       "messages": [{"role": "user", "content": "Hello!"}]
     }'

5. ${GREEN}Access the dashboard${NC}
   Open: http://your-server-ip:8080/dashboard
   Or: https://your-domain.com/dashboard (if nginx configured)

${YELLOW}Useful commands:${NC}
   systemctl status $APP_NAME    # Check service status
   journalctl -u $APP_NAME -f    # View logs
   systemctl restart $APP_NAME   # Restart service

${YELLOW}Configuration files:${NC}
   Service: $SERVICE_FILE
   Config: $CONFIG_DIR/config.toml
   Environment: $ENV_FILE
   Database: $DATA_DIR/llm_proxy.db
   Logs: $LOG_DIR/

${GREEN}For more information, see:${NC}
   https://git.dustin.coffee/hobokenchicken/llm-proxy
   $INSTALL_DIR/README.md
   $INSTALL_DIR/deployment.md

EOF
}

# Main deployment function
deploy() {
    log_info "Starting LLM Proxy Gateway deployment..."
    
    check_root
    install_dependencies
    install_rust
    setup_directories
    build_application
    create_configuration
    create_systemd_service
    initialize_database
    start_service
    verify_installation
    print_next_steps
    
    # Optional steps (uncomment if needed)
    # setup_nginx_proxy
    # setup_firewall
    
    log_info "Deployment completed successfully!"
}

# Update function
update() {
    log_info "Updating LLM Proxy Gateway..."
    
    check_root
    
    # Pull latest changes (while service keeps running)
    cd "$INSTALL_DIR"
    log_info "Pulling latest changes..."
    git pull
    
    # Build new binary (service stays up on the old binary)
    log_info "Building release binary (service still running)..."
    if ! cargo build --release; then
        log_error "Build failed — service was NOT interrupted. Fix the error and try again."
        exit 1
    fi
    
    # Verify binary exists
    if [[ ! -f "target/release/$APP_NAME" ]]; then
        log_error "Binary not found after build — aborting."
        exit 1
    fi
    
    # Restart service to pick up new binary
    log_info "Build succeeded. Restarting service..."
    systemctl restart "$APP_NAME"
    
    sleep 2
    if systemctl is-active --quiet "$APP_NAME"; then
        log_info "Update completed successfully!"
        systemctl status "$APP_NAME" --no-pager
    else
        log_error "Service failed to start after update. Check logs:"
        journalctl -u "$APP_NAME" -n 20 --no-pager
        exit 1
    fi
}

# Uninstall function
uninstall() {
    log_info "Uninstalling LLM Proxy Gateway..."
    
    check_root
    
    # Stop and disable service
    systemctl stop "$APP_NAME" 2>/dev/null || true
    systemctl disable "$APP_NAME" 2>/dev/null || true
    rm -f "$SERVICE_FILE"
    systemctl daemon-reload
    
    # Remove application files
    rm -rf "$INSTALL_DIR"
    rm -rf "$CONFIG_DIR"
    
    # Keep data and logs (comment out to remove)
    log_warn "Data directory $DATA_DIR and logs $LOG_DIR have been preserved"
    log_warn "Remove manually if desired:"
    log_warn "  rm -rf $DATA_DIR $LOG_DIR"
    
    # Remove user (optional)
    read -p "Remove user $APP_USER? [y/N]: " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        userdel "$APP_USER" 2>/dev/null || true
        groupdel "$APP_GROUP" 2>/dev/null || true
    fi
    
    log_info "Uninstallation completed!"
}

# Show usage
usage() {
    cat << EOF
LLM Proxy Gateway Deployment Script

Usage: $0 [command]

Commands:
  deploy     - Install and configure LLM Proxy Gateway
  update     - Pull latest changes, rebuild, and restart
  status     - Show service status and health check
  logs       - Tail the service logs (Ctrl+C to stop)
  uninstall  - Remove LLM Proxy Gateway
  help       - Show this help message

Examples:
  $0 deploy     # Full installation
  $0 update     # Update to latest version
  $0 status     # Check if service is healthy
  $0 logs       # Follow live logs

EOF
}

# Status function
status() {
    echo ""
    log_info "Service status:"
    systemctl status "$APP_NAME" --no-pager 2>/dev/null || log_warn "Service not found"
    echo ""

    # Health check
    if curl -sf http://localhost:8080/health &>/dev/null; then
        log_info "Health check: OK"
    else
        log_warn "Health check: FAILED (service may not be running or port 8080 not responding)"
    fi

    # Show current git commit
    if [[ -d "$INSTALL_DIR/.git" ]]; then
        echo ""
        log_info "Installed version:"
        git -C "$INSTALL_DIR" log -1 --format="  %h  %s  (%cr)" 2>/dev/null
    fi
}

# Logs function
logs() {
    log_info "Tailing $APP_NAME logs (Ctrl+C to stop)..."
    journalctl -u "$APP_NAME" -f
}

# Parse command line arguments
case "${1:-}" in
    deploy)
        deploy
        ;;
    update)
        update
        ;;
    status)
        status
        ;;
    logs)
        logs
        ;;
    uninstall)
        uninstall
        ;;
    help|--help|-h)
        usage
        ;;
    *)
        usage
        exit 1
        ;;
esac