#!/bin/bash
# unicornscan-alicorn - Alicorn Web UI management script
# Copyright (c) 2025 Robert E. Lee <robert@unicornscan.org>
#
# Usage:
#   unicornscan-alicorn start     Start the web UI (PostgreSQL + PostgREST + Alicorn)
#   unicornscan-alicorn stop      Stop all containers
#   unicornscan-alicorn status    Show container status
#   unicornscan-alicorn logs      View container logs
#   unicornscan-alicorn shell     Open PostgreSQL shell

set -e

INSTALL_DIR="/usr/share/unicornscan/alicorn"
VAR_DIR="/var/lib/unicornscan/alicorn"
COMPOSE_FILE="$INSTALL_DIR/docker-compose.yml"
ENV_FILE="$VAR_DIR/.env"
PASSWORD_FILE="$VAR_DIR/.db_password"
MODULES_CONF="/etc/unicornscan/modules.conf"

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

print_status() {
    echo -e "${BLUE}[unicornscan-alicorn]${NC} $1"
}

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

print_warning() {
    echo -e "${YELLOW}[!]${NC} $1"
}

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

# Generate a random 32-character password
generate_password() {
    tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 32
}

# Check if postgres volume exists (from previous install)
check_orphan_volume() {
    local volume_name="alicorn_postgres_data"
    if docker volume inspect "$volume_name" &>/dev/null; then
        return 0  # Volume exists
    fi
    return 1  # No volume
}

# Check if required ports are available
check_ports() {
    local ports_in_use=""
    local port_info=""

    for port in 31337 31338 5432; do
        if ss -tuln 2>/dev/null | grep -q ":${port} " || \
           netstat -tuln 2>/dev/null | grep -q ":${port} "; then
            ports_in_use="$ports_in_use $port"
            # Try to identify what's using the port
            local proc=$(ss -tulnp 2>/dev/null | grep ":${port} " | sed 's/.*users:(("\([^"]*\)".*/\1/' | head -1)
            [ -z "$proc" ] && proc=$(lsof -i :${port} 2>/dev/null | tail -1 | awk '{print $1}')
            [ -n "$proc" ] && port_info="$port_info\n    Port $port: $proc"
        fi
    done

    if [ -n "$ports_in_use" ]; then
        echo ""
        print_error "Required ports are already in use:$ports_in_use"
        [ -n "$port_info" ] && echo -e "$port_info"
        echo ""
        echo "  Port 31337 = Alicorn Web UI"
        echo "  Port 31338 = Alicorn API (PostgREST)"
        echo "  Port 5432  = PostgreSQL database"
        echo ""
        echo "Please stop the conflicting service(s) or change their ports."
        echo "If this is from a previous Alicorn run, try: sudo unicornscan-alicorn stop"
        return 1
    fi
    return 0
}

# Sync password to modules.conf
sync_modules_conf() {
    local password="$1"
    if [ -f "$MODULES_CONF" ]; then
        if grep -q 'module "pgsqldb"' "$MODULES_CONF"; then
            # Check if password already matches
            if grep -q "password=${password}" "$MODULES_CONF"; then
                return 0  # Already synced
            fi
            # Update password in existing pgsqldb config
            if [ -w "$MODULES_CONF" ]; then
                sed -i "s/password=[^;[:space:]\"]*/password=${password}/" "$MODULES_CONF"
                # Verify the update worked
                if grep -q "password=${password}" "$MODULES_CONF"; then
                    print_success "Updated $MODULES_CONF with database password"
                else
                    print_warning "Failed to update password in $MODULES_CONF"
                    print_warning "Manually update: password=${password}"
                fi
            else
                print_warning "Cannot write to $MODULES_CONF (run with sudo)"
                print_warning "Manually update: password=${password}"
            fi
        else
            print_warning "$MODULES_CONF exists but has no pgsqldb module configured"
        fi
    else
        print_warning "$MODULES_CONF not found - unicornscan may not export to database"
    fi
}

# Initialize environment with random password if not exists
init_env() {
    # If env already exists, just sync modules.conf and return
    if [ -f "$ENV_FILE" ] && [ -f "$PASSWORD_FILE" ]; then
        sync_modules_conf "$(cat "$PASSWORD_FILE")"
        return 0
    fi

    # Check for orphan volume with old password
    if check_orphan_volume; then
        print_error "Database volume exists but no password file found!"
        print_error "This happens after reinstalling without removing Docker volumes."
        echo ""
        echo "To fix this, reset the database:"
        echo "  sudo unicornscan-alicorn reset"
        echo "  sudo unicornscan-alicorn start"
        echo ""
        echo "WARNING: This will delete any existing scan data."
        exit 1
    fi

    print_status "First run - generating secure database password..."

    PASSWORD=$(generate_password)

    # Create .env file
    cat > "$ENV_FILE" << EOF
# Alicorn Docker Configuration
# Generated on $(date -Iseconds)
POSTGRES_USER=alicorn
POSTGRES_PASSWORD=${PASSWORD}
POSTGRES_DB=unicornscan
EOF
    chmod 600 "$ENV_FILE"

    # Save password to separate file for easy access
    echo "$PASSWORD" > "$PASSWORD_FILE"
    chmod 600 "$PASSWORD_FILE"

    # Update modules.conf
    sync_modules_conf "$PASSWORD"

    print_success "Generated secure password (saved to $PASSWORD_FILE)"
}

# Get the current password (from .env or default)
get_password() {
    if [ -f "$PASSWORD_FILE" ]; then
        cat "$PASSWORD_FILE"
    elif [ -f "$ENV_FILE" ]; then
        grep '^POSTGRES_PASSWORD=' "$ENV_FILE" | cut -d'=' -f2
    else
        echo "changeme"
    fi
}

# Detect docker compose command (v2 plugin vs v1 standalone)
check_compose_available() {
    if docker compose version &> /dev/null; then
        return 0
    elif command -v docker-compose &> /dev/null; then
        return 0
    else
        print_error "Docker Compose is not installed."
        echo "Install with: sudo apt install docker-compose-v2"
        exit 1
    fi
}

get_compose_cmd() {
    if docker compose version &> /dev/null; then
        echo "docker compose"
    else
        echo "docker-compose"
    fi
}

check_docker() {
    if ! command -v docker &> /dev/null; then
        print_error "Docker is not installed. Please install Docker first:"
        echo "  sudo apt install docker.io docker-compose-v2"
        exit 1
    fi

    if ! docker info &> /dev/null; then
        print_error "Docker daemon is not running or you don't have permission."
        echo "Try: sudo systemctl start docker"
        echo "Or add yourself to docker group: sudo usermod -aG docker \$USER"
        exit 1
    fi
}

check_write_permission() {
    # Ensure VAR_DIR exists and is writable
    if [ ! -d "$VAR_DIR" ]; then
        mkdir -p "$VAR_DIR" 2>/dev/null || true
    fi
    if [ ! -w "$VAR_DIR" ]; then
        print_error "Cannot write to $VAR_DIR"
        echo "Run with sudo: sudo unicornscan-alicorn $1"
        exit 1
    fi
}

check_compose_file() {
    if [ ! -f "$COMPOSE_FILE" ]; then
        print_error "Docker Compose file not found at $COMPOSE_FILE"
        exit 1
    fi
}

cmd_start() {
    check_docker
    check_compose_available
    check_compose_file
    check_write_permission start

    # Check if required ports are available before proceeding
    check_ports || exit 1

    COMPOSE_CMD=$(get_compose_cmd)

    # Generate password on first run
    init_env

    print_status "Starting Alicorn Web UI stack..."
    cd "$INSTALL_DIR"

    # Clean up orphan containers (can happen after purge+reinstall or failed starts)
    # These are containers with alicorn- names not managed by current compose project
    ORPHAN_CONTAINERS=$(docker ps -a --filter "name=alicorn-" --format '{{.Names}}' 2>/dev/null | tr '\n' ' ')
    if [ -n "$ORPHAN_CONTAINERS" ]; then
        echo ""
        print_warning "Found existing Alicorn containers: $ORPHAN_CONTAINERS"
        echo "These may be from a previous install and need to be removed."
        echo ""
        read -p "Remove these containers to continue? [y/N] " -n 1 -r
        echo ""
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            print_status "Removing orphan containers..."
            docker rm -f alicorn-postgres alicorn-postgrest alicorn-web alicorn-geoip 2>/dev/null || true
        else
            print_error "Cannot start: existing containers conflict. Run 'sudo unicornscan-alicorn stop' first."
            exit 1
        fi
    fi

    # Build and start containers (first run will build images, takes ~2 minutes)
    print_status "Building containers (first run may take 1-2 minutes)..."
    $COMPOSE_CMD --env-file "$ENV_FILE" up -d --build

    # Wait for services to be healthy
    print_status "Waiting for services to start..."
    sleep 5

    PASSWORD=$(get_password)

    echo ""
    print_success "Alicorn Web UI started!"
    echo ""
    echo "  Web UI:    http://localhost:31337"
    echo "  API:       http://localhost:31338"
    echo "  Database:  postgresql://localhost:5432/unicornscan"
    echo ""
    echo "Run a scan with database export:"
    echo "  unicornscan -epgsqldb <target>"
    echo ""
    echo "Or with explicit connection (password: $PASSWORD):"
    echo "  unicornscan -epgsqldb,host=localhost,user=alicorn,pass=$PASSWORD,db=unicornscan <target>"
    echo ""
}

cmd_stop() {
    check_docker
    check_compose_file

    COMPOSE_CMD=$(get_compose_cmd)

    print_status "Stopping Alicorn Web UI stack..."
    cd "$INSTALL_DIR"
    $COMPOSE_CMD --env-file "$ENV_FILE" down

    print_success "All containers stopped."
}

cmd_status() {
    check_docker
    check_compose_file

    COMPOSE_CMD=$(get_compose_cmd)

    print_status "Container status:"
    cd "$INSTALL_DIR"
    $COMPOSE_CMD --env-file "$ENV_FILE" ps
}

cmd_logs() {
    check_docker
    check_compose_file

    COMPOSE_CMD=$(get_compose_cmd)

    cd "$INSTALL_DIR"
    $COMPOSE_CMD --env-file "$ENV_FILE" logs -f "${@:2}"
}

cmd_shell() {
    check_docker
    check_compose_file

    COMPOSE_CMD=$(get_compose_cmd)

    print_status "Connecting to PostgreSQL..."
    cd "$INSTALL_DIR"
    PASSWORD=$(get_password)
    $COMPOSE_CMD --env-file "$ENV_FILE" exec -e PGPASSWORD="$PASSWORD" postgres psql -U alicorn -d unicornscan
}

cmd_password() {
    if [ -f "$PASSWORD_FILE" ]; then
        echo "Database password: $(cat $PASSWORD_FILE)"
        echo ""
        echo "Password file: $PASSWORD_FILE"
    elif [ -f "$ENV_FILE" ]; then
        PASSWORD=$(grep '^POSTGRES_PASSWORD=' "$ENV_FILE" | cut -d'=' -f2)
        echo "Database password: $PASSWORD"
        echo ""
        echo "Config file: $ENV_FILE"
    else
        print_warning "No password file found. Run 'unicornscan-alicorn start' first."
    fi
    echo ""
    echo "Application: $INSTALL_DIR"
    echo "Runtime data: $VAR_DIR"
}

cmd_reset() {
    check_docker
    check_compose_file
    check_write_permission reset

    COMPOSE_CMD=$(get_compose_cmd)

    print_warning "This will delete ALL scan data! Are you sure? (y/N)"
    read -r confirm
    if [ "$confirm" != "y" ] && [ "$confirm" != "Y" ]; then
        echo "Aborted."
        exit 0
    fi

    cd "$INSTALL_DIR"

    # Handle missing .env file gracefully (happens after apt purge + reinstall)
    if [ -f "$ENV_FILE" ]; then
        $COMPOSE_CMD --env-file "$ENV_FILE" down -v 2>/dev/null || true
    else
        # No env file exists - still try to bring down containers/volumes
        $COMPOSE_CMD down -v 2>/dev/null || true
    fi

    # Fallback: directly remove the volume if compose didn't catch it
    # (handles orphan volumes from old installs or different project names)
    docker volume rm alicorn_postgres_data 2>/dev/null || true

    # Remove credentials from VAR_DIR so fresh password is generated on next start
    rm -f "$ENV_FILE" "$PASSWORD_FILE"

    print_success "Database reset. Run 'sudo unicornscan-alicorn start' to create fresh database."
}

cmd_purge_volume() {
    check_docker

    # Check if volume exists
    if ! docker volume inspect alicorn_postgres_data &>/dev/null; then
        print_warning "No Alicorn database volume found."
        echo "Volume 'alicorn_postgres_data' does not exist."
        exit 0
    fi

    # Show volume info
    print_status "Found Docker volume: alicorn_postgres_data"
    docker volume inspect alicorn_postgres_data --format '  Created: {{.CreatedAt}}' 2>/dev/null || true

    echo ""
    print_warning "This will PERMANENTLY DELETE the database volume and all scan data!"
    print_warning "This action cannot be undone."
    echo ""
    echo "Type 'yes' to confirm deletion:"
    read -r confirm
    if [ "$confirm" != "yes" ]; then
        echo "Aborted. Volume not removed."
        exit 0
    fi

    # Stop any running containers that might be using the volume
    if docker ps -q --filter "name=alicorn-" 2>/dev/null | grep -q .; then
        print_status "Stopping Alicorn containers..."
        docker stop alicorn-postgres alicorn-postgrest alicorn-web alicorn-geoip 2>/dev/null || true
        docker rm alicorn-postgres alicorn-postgrest alicorn-web alicorn-geoip 2>/dev/null || true
    fi

    # Remove the volume
    if docker volume rm alicorn_postgres_data 2>/dev/null; then
        print_success "Volume 'alicorn_postgres_data' removed."
    else
        print_error "Failed to remove volume. It may be in use by a container."
        echo "Try: docker ps -a | grep alicorn"
        exit 1
    fi

    # Also clean up credentials if they exist
    rm -f "$ENV_FILE" "$PASSWORD_FILE" 2>/dev/null || true

    echo ""
    print_success "Cleanup complete. Run 'sudo unicornscan-alicorn start' for fresh install."
}

cmd_setup() {
    if ! command -v node &> /dev/null; then
        print_error "Node.js is required for the setup wizard."
        echo "Install Node.js 18+ from: https://nodejs.org/"
        exit 1
    fi

    if ! command -v npx &> /dev/null; then
        print_error "npm/npx is required for the setup wizard."
        exit 1
    fi

    print_status "Starting Alicorn Setup Wizard..."
    cd "$INSTALL_DIR"

    # Install dependencies if needed
    if [ ! -d "node_modules" ]; then
        print_status "Installing dependencies (first run)..."
        npm install --silent
    fi

    # Run the setup wizard
    npx tsx cli/setup.ts
}

cmd_dev() {
    if ! command -v node &> /dev/null; then
        print_error "Node.js is required for development mode."
        exit 1
    fi

    cd "$INSTALL_DIR"

    # Install dependencies if needed
    if [ ! -d "node_modules" ]; then
        print_status "Installing dependencies..."
        npm install
    fi

    print_status "Starting development server..."
    npm run dev
}

cmd_build() {
    if ! command -v node &> /dev/null; then
        print_error "Node.js is required to build."
        exit 1
    fi

    cd "$INSTALL_DIR"

    if [ ! -d "node_modules" ]; then
        print_status "Installing dependencies..."
        npm install
    fi

    print_status "Building production bundle..."
    npm run build
    print_success "Build complete! Output in: $INSTALL_DIR/dist/"
}

cmd_help() {
    echo "Alicorn Web UI - unicornscan results viewer"
    echo ""
    echo "Usage: unicornscan-alicorn <command>"
    echo ""
    echo "Docker Commands (recommended):"
    echo "  start        Start the web UI stack (PostgreSQL + PostgREST + Alicorn)"
    echo "  stop         Stop all containers"
    echo "  status       Show container status"
    echo "  logs         View container logs (use: logs [service])"
    echo "  shell        Open PostgreSQL shell"
    echo "  password     Show database password"
    echo "  reset        Delete all data and reset database (WARNING: destructive!)"
    echo "  purge-volume Remove orphan Docker volume (use after apt purge)"
    echo ""
    echo "Development Commands (requires Node.js 18+):"
    echo "  setup     Run interactive setup wizard"
    echo "  dev       Start development server with hot reload"
    echo "  build     Build production bundle"
    echo ""
    echo "  help      Show this help message"
    echo ""
    echo "Quick Start:"
    echo "  1. sudo unicornscan-alicorn start"
    echo "  2. unicornscan -epgsqldb <target>"
    echo "  3. Open http://localhost:31337"
    echo ""
}

# Main command dispatcher
case "${1:-help}" in
    start)
        cmd_start
        ;;
    stop)
        cmd_stop
        ;;
    status)
        cmd_status
        ;;
    logs)
        cmd_logs "$@"
        ;;
    shell)
        cmd_shell
        ;;
    password|pass|pw)
        cmd_password
        ;;
    reset)
        cmd_reset
        ;;
    purge-volume)
        cmd_purge_volume
        ;;
    setup)
        cmd_setup
        ;;
    dev)
        cmd_dev
        ;;
    build)
        cmd_build
        ;;
    help|--help|-h)
        cmd_help
        ;;
    *)
        print_error "Unknown command: $1"
        cmd_help
        exit 1
        ;;
esac
