#!/bin/bash
#
# AtlasOS Installation Script
# Installs AtlasOS storage controller on a Linux system
#
# Usage: sudo ./install.sh [options]
#
# Note: Run this script from the atlas repository root directory
#

set -e

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

# Default values
INSTALL_DIR="/opt/atlas"
DATA_DIR="/var/lib/atlas"
CONFIG_DIR="/etc/atlas"
SERVICE_USER="atlas"
LOG_DIR="/var/log/atlas"
BACKUP_DIR="/var/lib/atlas/backups"
HTTP_ADDR=":8080"
DB_PATH="/var/lib/atlas/atlas.db"
BUILD_BINARIES=true
SKIP_DEPS=false
REPO_DIR=""

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --install-dir)
            INSTALL_DIR="$2"
            shift 2
            ;;
        --data-dir)
            DATA_DIR="$2"
            shift 2
            ;;
        --skip-deps)
            SKIP_DEPS=true
            shift
            ;;
        --skip-build)
            BUILD_BINARIES=false
            shift
            ;;
        --http-addr)
            HTTP_ADDR="$2"
            shift 2
            ;;
        --repo-dir)
            REPO_DIR="$2"
            shift 2
            ;;
        -h|--help)
            echo "AtlasOS Installation Script"
            echo ""
            echo "Usage: sudo ./install.sh [options]"
            echo ""
            echo "Options:"
            echo "  --install-dir DIR    Installation directory (default: /opt/atlas)"
            echo "  --data-dir DIR       Data directory (default: /var/lib/atlas)"
            echo "  --skip-deps          Skip dependency installation"
            echo "  --skip-build         Skip building binaries (use existing)"
            echo "  --http-addr ADDR     HTTP address (default: :8080)"
            echo "  --repo-dir DIR       Repository directory (if not in current dir)"
            echo "  -h, --help           Show this help message"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Check if running as root
if [[ $EUID -ne 0 ]]; then
    echo -e "${RED}Error: This script must be run as root (use sudo)${NC}"
    exit 1
fi

# Get script directory early (for path resolution)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Detect distribution
detect_distro() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        DISTRO=$ID
        VERSION=$VERSION_ID
    else
        echo -e "${RED}Error: Cannot detect Linux distribution${NC}"
        exit 1
    fi
}

# Install dependencies
install_dependencies() {
    echo -e "${GREEN}Installing dependencies...${NC}"
    
    case $DISTRO in
        ubuntu|debian)
            apt-get update
            
            # Try targetcli-fb first (common on newer Ubuntu/Debian), fallback to targetcli
            TARGETCLI_PKG="targetcli-fb"
            if ! apt-cache show targetcli-fb &>/dev/null; then
                TARGETCLI_PKG="targetcli"
            fi
            
            apt-get install -y \
                zfsutils-linux \
                samba \
                nfs-kernel-server \
                $TARGETCLI_PKG \
                sqlite3 \
                golang-go \
                git \
                build-essential \
                curl
            ;;
        fedora|rhel|centos)
            if command -v dnf &> /dev/null; then
                dnf install -y \
                    zfs \
                    samba \
                    nfs-utils \
                    targetcli \
                    sqlite \
                    golang \
                    git \
                    gcc \
                    make \
                    curl
            else
                yum install -y \
                    zfs \
                    samba \
                    nfs-utils \
                    targetcli \
                    sqlite \
                    golang \
                    git \
                    gcc \
                    make \
                    curl
            fi
            ;;
        *)
            echo -e "${YELLOW}Warning: Unknown distribution. Please install dependencies manually:${NC}"
            echo "  - ZFS utilities"
            echo "  - Samba (SMB/CIFS)"
            echo "  - NFS server"
            echo "  - targetcli or targetcli-fb (iSCSI)"
            echo "  - SQLite"
            echo "  - Go compiler"
            echo "  - Build tools"
            ;;
    esac
    
    echo -e "${GREEN}Dependencies installed${NC}"
}

# Create system user
create_user() {
    echo -e "${GREEN}Creating system user...${NC}"
    
    if ! id "$SERVICE_USER" &>/dev/null; then
        useradd -r -s /bin/false -d "$DATA_DIR" "$SERVICE_USER"
        echo -e "${GREEN}User $SERVICE_USER created${NC}"
    else
        echo -e "${YELLOW}User $SERVICE_USER already exists${NC}"
    fi
}

# Create directories
create_directories() {
    echo -e "${GREEN}Creating directories...${NC}"
    
    mkdir -p "$INSTALL_DIR/bin"
    mkdir -p "$DATA_DIR"
    mkdir -p "$CONFIG_DIR"
    mkdir -p "$LOG_DIR"
    mkdir -p "$BACKUP_DIR"
    mkdir -p "$CONFIG_DIR/tls"
    
    # Set ownership
    chown -R "$SERVICE_USER:$SERVICE_USER" "$DATA_DIR"
    chown -R "$SERVICE_USER:$SERVICE_USER" "$LOG_DIR"
    chown -R "$SERVICE_USER:$SERVICE_USER" "$BACKUP_DIR"
    chown -R "$SERVICE_USER:$SERVICE_USER" "$CONFIG_DIR"
    
    # Set permissions
    chmod 755 "$INSTALL_DIR"
    chmod 755 "$INSTALL_DIR/bin"
    chmod 700 "$DATA_DIR"
    chmod 700 "$CONFIG_DIR"
    chmod 750 "$LOG_DIR"
    chmod 750 "$BACKUP_DIR"
    
    echo -e "${GREEN}Directories created${NC}"
}

# Build binaries
build_binaries() {
    if [[ "$BUILD_BINARIES" == "false" ]]; then
        echo -e "${YELLOW}Skipping binary build${NC}"
        return
    fi
    
    echo -e "${GREEN}Building binaries...${NC}"
    
    # Check if we're in the right directory
    # Use REPO_DIR if specified, otherwise try current directory, then script directory
    if [[ -n "$REPO_DIR" ]] && [[ -d "$REPO_DIR/cmd" ]] && [[ -f "$REPO_DIR/cmd/atlas-api/main.go" ]]; then
        BUILD_DIR="$(cd "$REPO_DIR" && pwd)"
    elif [[ -d "./cmd" ]] && [[ -f "./cmd/atlas-api/main.go" ]]; then
        BUILD_DIR="$(pwd)"
    elif [[ -d "$SCRIPT_DIR/cmd" ]] && [[ -f "$SCRIPT_DIR/cmd/atlas-api/main.go" ]]; then
        BUILD_DIR="$SCRIPT_DIR"
    elif [[ -d "$SCRIPT_DIR/../cmd" ]] && [[ -f "$SCRIPT_DIR/../cmd/atlas-api/main.go" ]]; then
        BUILD_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
    else
        echo -e "${RED}Error: Cannot find source code directory${NC}"
        echo "  Current directory: $(pwd)"
        echo "  Script directory: $SCRIPT_DIR"
        echo "  Looking for: cmd/atlas-api/main.go"
        echo ""
        echo "  Please ensure you're running the installer from the atlas repository root"
        echo "  Or specify the repository path with: --repo-dir /path/to/atlas"
        exit 1
    fi
    
    cd "$BUILD_DIR"
    echo "Building from: $(pwd)"
    
    # Verify source files exist
    if [[ ! -f "$BUILD_DIR/cmd/atlas-api/main.go" ]]; then
        echo -e "${RED}Error: Cannot find cmd/atlas-api/main.go${NC}"
        echo "  Expected at: $BUILD_DIR/cmd/atlas-api/main.go"
        exit 1
    fi
    
    # Build binaries
    echo "Building atlas-api..."
    if ! go build -o "$INSTALL_DIR/bin/atlas-api" ./cmd/atlas-api; then
        echo -e "${RED}Error: Failed to build atlas-api${NC}"
        exit 1
    fi
    
    echo "Building atlas-tui..."
    if ! go build -o "$INSTALL_DIR/bin/atlas-tui" ./cmd/atlas-tui; then
        echo -e "${RED}Error: Failed to build atlas-tui${NC}"
        exit 1
    fi
    
    # Set permissions
    chown root:root "$INSTALL_DIR/bin/atlas-api"
    chown root:root "$INSTALL_DIR/bin/atlas-tui"
    chmod 755 "$INSTALL_DIR/bin/atlas-api"
    chmod 755 "$INSTALL_DIR/bin/atlas-tui"
    
    echo -e "${GREEN}Binaries built successfully${NC}"
}

# Create systemd service
create_systemd_service() {
    echo -e "${GREEN}Creating systemd service...${NC}"
    
    cat > /etc/systemd/system/atlas-api.service <<EOF
[Unit]
Description=AtlasOS Storage Controller API
After=network.target zfs.target

[Service]
Type=simple
User=$SERVICE_USER
Group=$SERVICE_USER
WorkingDirectory=$INSTALL_DIR
ExecStart=$INSTALL_DIR/bin/atlas-api
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal
SyslogIdentifier=atlas-api

# Environment variables
Environment="ATLAS_HTTP_ADDR=$HTTP_ADDR"
Environment="ATLAS_DB_PATH=$DB_PATH"
Environment="ATLAS_BACKUP_DIR=$BACKUP_DIR"
Environment="ATLAS_LOG_LEVEL=INFO"
Environment="ATLAS_LOG_FORMAT=json"

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

[Install]
WantedBy=multi-user.target
EOF
    
    systemctl daemon-reload
    echo -e "${GREEN}Systemd service created${NC}"
}

# Create configuration file
create_config() {
    echo -e "${GREEN}Creating configuration...${NC}"
    
    cat > "$CONFIG_DIR/atlas.conf" <<EOF
# AtlasOS Configuration
# This file is sourced by the systemd service

# HTTP Server
ATLAS_HTTP_ADDR=$HTTP_ADDR

# Database
ATLAS_DB_PATH=$DB_PATH

# Backup Directory
ATLAS_BACKUP_DIR=$BACKUP_DIR

# Logging
ATLAS_LOG_LEVEL=INFO
ATLAS_LOG_FORMAT=json

# TLS (optional - uncomment to enable)
# ATLAS_TLS_ENABLED=true
# ATLAS_TLS_CERT=$CONFIG_DIR/tls/cert.pem
# ATLAS_TLS_KEY=$CONFIG_DIR/tls/key.pem

# JWT Secret (generate with: openssl rand -hex 32)
# ATLAS_JWT_SECRET=your-secret-here
EOF
    
    chown "$SERVICE_USER:$SERVICE_USER" "$CONFIG_DIR/atlas.conf"
    chmod 600 "$CONFIG_DIR/atlas.conf"
    
    echo -e "${GREEN}Configuration created${NC}"
}

# Generate JWT secret
generate_jwt_secret() {
    echo -e "${GREEN}Generating JWT secret...${NC}"
    
    if command -v openssl &> /dev/null; then
        JWT_SECRET=$(openssl rand -hex 32)
        echo "ATLAS_JWT_SECRET=$JWT_SECRET" >> "$CONFIG_DIR/atlas.conf"
        echo -e "${GREEN}JWT secret generated${NC}"
    else
        echo -e "${YELLOW}Warning: openssl not found. Please set ATLAS_JWT_SECRET manually${NC}"
    fi
}

# Setup ZFS (if needed)
setup_zfs() {
    echo -e "${GREEN}Checking ZFS...${NC}"
    
    if ! command -v zpool &> /dev/null; then
        echo -e "${YELLOW}Warning: ZFS not found. Please install ZFS utilities${NC}"
        return
    fi
    
    # Check if ZFS module is loaded
    if ! lsmod | grep -q zfs; then
        echo -e "${YELLOW}Warning: ZFS kernel module not loaded${NC}"
        echo "  Run: modprobe zfs"
    fi
    
    echo -e "${GREEN}ZFS check complete${NC}"
}

# Setup Samba
setup_samba() {
    echo -e "${GREEN}Setting up Samba...${NC}"
    
    if ! command -v smbd &> /dev/null; then
        echo -e "${YELLOW}Warning: Samba not found${NC}"
        return
    fi
    
    # Enable and start Samba (if not already)
    systemctl enable smbd 2>/dev/null || true
    systemctl enable nmbd 2>/dev/null || true
    
    echo -e "${GREEN}Samba setup complete${NC}"
}

# Setup NFS
setup_nfs() {
    echo -e "${GREEN}Setting up NFS...${NC}"
    
    if ! command -v exportfs &> /dev/null; then
        echo -e "${YELLOW}Warning: NFS not found${NC}"
        return
    fi
    
    # Enable and start NFS (if not already)
    systemctl enable nfs-server 2>/dev/null || true
    systemctl enable rpcbind 2>/dev/null || true
    
    echo -e "${GREEN}NFS setup complete${NC}"
}

# Setup iSCSI
setup_iscsi() {
    echo -e "${GREEN}Setting up iSCSI...${NC}"
    
    # Check for targetcli or targetcli-fb
    TARGETCLI_CMD=""
    if command -v targetcli &> /dev/null; then
        TARGETCLI_CMD="targetcli"
    elif command -v targetcli-fb &> /dev/null; then
        TARGETCLI_CMD="targetcli-fb"
        # Create symlink if targetcli doesn't exist
        if ! command -v targetcli &> /dev/null; then
            ln -s $(which targetcli-fb) /usr/local/bin/targetcli 2>/dev/null || true
        fi
    fi
    
    if [[ -z "$TARGETCLI_CMD" ]]; then
        echo -e "${YELLOW}Warning: targetcli or targetcli-fb not found${NC}"
        echo "  Install with: apt-get install targetcli-fb (Ubuntu/Debian)"
        return
    fi
    
    # Enable and start iSCSI target (if not already)
    systemctl enable target 2>/dev/null || true
    
    echo -e "${GREEN}iSCSI setup complete (using $TARGETCLI_CMD)${NC}"
}

# Create initial admin user
create_admin_user() {
    echo -e "${GREEN}Creating initial admin user...${NC}"
    
    echo ""
    echo -e "${YELLOW}Please set up the initial admin user:${NC}"
    echo "  Username: admin"
    echo "  Password: (you will be prompted)"
    echo ""
    echo "After starting the service, you can create the admin user via:"
    echo "  curl -X POST http://localhost:8080/api/v1/users \\"
    echo "    -H 'Content-Type: application/json' \\"
    echo "    -d '{\"username\":\"admin\",\"password\":\"your-password\",\"role\":\"administrator\"}'"
    echo ""
    echo "Or use the TUI:"
    echo "  $INSTALL_DIR/bin/atlas-tui"
    echo ""
}

# Start service
start_service() {
    echo -e "${GREEN}Starting AtlasOS service...${NC}"
    
    systemctl enable atlas-api
    systemctl start atlas-api
    
    # Wait a moment for service to start
    sleep 2
    
    if systemctl is-active --quiet atlas-api; then
        echo -e "${GREEN}AtlasOS service started successfully${NC}"
    else
        echo -e "${RED}Error: Service failed to start${NC}"
        echo "Check logs with: journalctl -u atlas-api -n 50"
        exit 1
    fi
}

# Print summary
print_summary() {
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}AtlasOS Installation Complete!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Installation Directory: $INSTALL_DIR"
    echo "Data Directory: $DATA_DIR"
    echo "Config Directory: $CONFIG_DIR"
    echo "Log Directory: $LOG_DIR"
    echo ""
    echo "Service Status:"
    systemctl status atlas-api --no-pager -l || true
    echo ""
    echo "Useful Commands:"
    echo "  Service: systemctl {start|stop|restart|status} atlas-api"
    echo "  Logs:   journalctl -u atlas-api -f"
    echo "  TUI:    $INSTALL_DIR/bin/atlas-tui"
    echo ""
    echo "Web Interface:"
    echo "  http://localhost:8080"
    echo ""
    echo "API Documentation:"
    echo "  http://localhost:8080/api/docs"
    echo ""
    echo -e "${YELLOW}Next Steps:${NC}"
    echo "1. Create initial admin user (see instructions above)"
    echo "2. Configure TLS certificates (optional)"
    echo "3. Review configuration in $CONFIG_DIR/atlas.conf"
    echo ""
}

# Main installation
main() {
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}AtlasOS Installation Script${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    
    detect_distro
    echo "Detected distribution: $DISTRO $VERSION"
    echo ""
    
    if [[ "$SKIP_DEPS" == "false" ]]; then
        install_dependencies
    else
        echo -e "${YELLOW}Skipping dependency installation${NC}"
    fi
    
    create_user
    create_directories
    build_binaries
    create_config
    generate_jwt_secret
    create_systemd_service
    
    setup_zfs
    setup_samba
    setup_nfs
    setup_iscsi
    
    create_admin_user
    
    # Ask if user wants to start service
    echo ""
    read -p "Start AtlasOS service now? (y/n) " -n 1 -r
    echo ""
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        start_service
    else
        echo -e "${YELLOW}Service not started. Start manually with: systemctl start atlas-api${NC}"
    fi
    
    print_summary
}

# Run main
main