#!/usr/bin/env bash
set -euo pipefail

# --- Configuration
USER_HOME="$HOME"
SP_REPO="${USER_HOME}/moonraker-sp"
SP_VENV="${USER_HOME}/moonraker-sp-env"
SP_DATA="${USER_HOME}/sp_printer_data"
SP_CONFIG="${SP_DATA}/config/moonraker.conf"
SP_ENV_FILE="${SP_DATA}/systemd/moonraker-sp.env"
SP_SERVICE="moonraker-sp"
SP_SERVICE_FILE="/etc/systemd/system/${SP_SERVICE}.service"
SP_PORT=7126

# Auto-detect printer layout
if [[ -f "${USER_HOME}/printer_data/config/moonraker.conf" ]]; then
    # Modern printer_data layout (Sovol SV08, OpenNept4une, KIAUH)
    EXISTING_CONFIG="${USER_HOME}/printer_data/config/moonraker.conf"
    EXISTING_GCODES="${USER_HOME}/printer_data/gcodes"
elif [[ -f "${USER_HOME}/klipper_config/moonraker.conf" ]]; then
    # Legacy layout (Elegoo Neptune 4, Qidi, Artillery X4, Sovol SV07)
    EXISTING_CONFIG="${USER_HOME}/klipper_config/moonraker.conf"
    EXISTING_GCODES="${USER_HOME}/gcode_files"
else
    EXISTING_CONFIG=""
    EXISTING_GCODES=""
fi

MOONRAKER_REPO="https://github.com/Arksine/moonraker.git"

MIN_DISK_MB=200
VERIFY_TIMEOUT=15

# --- Helpers
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

info()  { echo -e "${GREEN}[INFO]${NC} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*"; }
step()  { echo -e "${CYAN}[STEP]${NC} ${BOLD}$*${NC}"; }

die() { error "$@"; exit 1; }

confirm() {
    local prompt="${1:-Continue?}"
    if [[ "${YES:-}" == "1" ]]; then
        return 0
    fi
    read -rp "$(echo -e "${YELLOW}${prompt} [y/N]:${NC} ")" answer
    [[ "${answer,,}" == "y" || "${answer,,}" == "yes" ]]
}

is_installed() {
    [[ -f "${SP_SERVICE_FILE}" ]] && [[ -d "${SP_REPO}/.git" ]] && [[ -d "${SP_VENV}" ]]
}

# --- Prerequisites
check_prerequisites() {
    step "Checking prerequisites..."
    local ok=1

    # Not root
    if [[ "$(id -u)" -eq 0 ]]; then
        error "Please don't run this script as root - run it without sudo or as the default user instead."
        ok=0
    fi

    # Python 3.7+
    if command -v python3 &>/dev/null; then
        local pyver
        pyver=$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
        if python3 -c 'import sys; exit(0 if sys.version_info >= (3,7) else 1)'; then
            info "Python ${pyver} - OK"
        else
            error "Python ${pyver} is too old (need 3.7 or newer)."
            ok=0
        fi
    else
        error "Python 3 is not installed."
        ok=0
    fi

    # venv / virtualenv
    if python3 -m virtualenv --version &>/dev/null 2>&1; then
        info "virtualenv available - OK"
    elif python3 -m venv --help &>/dev/null 2>&1; then
        info "python3 -m venv available - OK"
    else
        error "Python virtual environment support is missing."
        ok=0
    fi

    # Disk space
    local free_mb
    free_mb=$(df -m "${USER_HOME}" | awk 'NR==2{print $4}')
    if [[ "${free_mb}" -ge "${MIN_DISK_MB}" ]]; then
        info "Disk free: ${free_mb} MB - OK"
    else
        error "Not enough disk space (${free_mb} MB free, need ${MIN_DISK_MB} MB). Try removing old gcode files."
        ok=0
    fi

    # Klipper socket (extract path from existing config, fall back to /tmp/klippy_uds)
    local klipper_socket="/tmp/klippy_uds"
    if [[ -n "${EXISTING_CONFIG}" ]] && [[ -f "${EXISTING_CONFIG}" ]]; then
        local parsed_socket
        parsed_socket=$(awk '/^klippy_uds_address[[:space:]]*[:=]/ {
            sub(/^klippy_uds_address[[:space:]]*[:=][[:space:]]*/, "")
            print; exit
        }' "${EXISTING_CONFIG}")
        [[ -n "${parsed_socket}" ]] && klipper_socket="${parsed_socket}"
    fi
    if [[ -S "${klipper_socket}" ]]; then
        info "Klipper socket exists - OK"
    else
        warn "Klipper doesn't seem to be running. Install will continue, but SimplyPrint won't work until Klipper is running."
    fi

    # Port not in use (unless already ours)
    if ss -tlnp 2>/dev/null | grep -q ":${SP_PORT} "; then
        if systemctl is-active --quiet "${SP_SERVICE}" 2>/dev/null; then
            info "Port ${SP_PORT} in use by ${SP_SERVICE} - OK (already installed)"
        else
            error "Port ${SP_PORT} is already in use by something else."
            ok=0
        fi
    else
        info "Port ${SP_PORT} available - OK"
    fi

    # Git
    if command -v git &>/dev/null; then
        info "git available - OK"
    else
        error "Git is not installed."
        ok=0
    fi

    # Network (quick check)
    if git ls-remote --exit-code "${MOONRAKER_REPO}" HEAD &>/dev/null; then
        info "Network connectivity - OK"
    else
        error "Can't connect to the internet. Please check your printer's network connection."
        ok=0
    fi

    # Sudo access
    if sudo -n true 2>/dev/null || sudo true; then
        info "Sudo access - OK"
    else
        error "Could not get administrator (sudo) access."
        ok=0
    fi

    # Existing config
    if [[ -n "${EXISTING_CONFIG}" ]] && [[ -f "${EXISTING_CONFIG}" ]]; then
        info "Existing moonraker.conf found at ${EXISTING_CONFIG} - OK"
    else
        error "Could not find moonraker.conf in ~/printer_data/config/ or ~/klipper_config/"
        ok=0
    fi

    if [[ "${ok}" -eq 0 ]]; then
        echo
        error "Some checks above failed. Please fix them and try again."
        error "If you're not sure what to do, contact SimplyPrint support."
        exit 1
    fi
    info "All prerequisites passed."
    echo
}

# --- Clone Repo
clone_repo() {
    step "Setting up moonraker repository..."

    if [[ -d "${SP_REPO}/.git" ]]; then
        info "Existing git repo found, updating..."
        cd "${SP_REPO}"
        git fetch origin
        git reset --hard origin/master
        cd "${USER_HOME}"
    elif [[ -d "${SP_REPO}" ]]; then
        warn "${SP_REPO} exists but is not a git repo - backing up and re-cloning"
        mv "${SP_REPO}" "${SP_REPO}.bak.$(date +%s)"
        git clone "${MOONRAKER_REPO}" "${SP_REPO}"
    else
        info "Cloning moonraker..."
        git clone "${MOONRAKER_REPO}" "${SP_REPO}"
    fi

    info "Moonraker repo ready at ${SP_REPO}"
    echo
}

# --- Python Venv
setup_venv() {
    step "Setting up Python virtual environment..."

    if [[ ! -d "${SP_VENV}" ]]; then
        if python3 -m virtualenv --version &>/dev/null 2>&1; then
            python3 -m virtualenv -p python3 "${SP_VENV}"
        else
            python3 -m venv "${SP_VENV}"
        fi
        info "Virtual environment created"
    else
        info "Virtual environment already exists"
    fi

    info "Upgrading pip..."
    "${SP_VENV}/bin/pip" install --upgrade pip --quiet

    info "Installing moonraker dependencies..."
    local req_file="${SP_REPO}/scripts/moonraker-requirements.txt"
    if [[ ! -f "${req_file}" ]]; then
        die "Requirements file not found at ${req_file} - try running Uninstall, then Install again."
    fi

    SKIP_CYTHON=1 "${SP_VENV}/bin/pip" install -r "${req_file}" --quiet

    info "Python environment ready"
    echo
}

# --- Data Directory
init_data_dir() {
    step "Initializing data directory..."

    mkdir -p "${SP_DATA}/config"
    mkdir -p "${SP_DATA}/logs"
    mkdir -p "${SP_DATA}/systemd"

    # Symlink gcodes
    if [[ -z "${EXISTING_GCODES}" ]] || [[ ! -d "${EXISTING_GCODES}" ]]; then
        warn "Gcode directory not found - skipping symlink."
    elif [[ -L "${SP_DATA}/gcodes" ]]; then
        info "gcodes symlink already exists"
    elif [[ -d "${SP_DATA}/gcodes" ]]; then
        warn "${SP_DATA}/gcodes is a directory, replacing with symlink"
        rmdir "${SP_DATA}/gcodes" 2>/dev/null || true
        ln -sf "${EXISTING_GCODES}" "${SP_DATA}/gcodes"
    else
        ln -sf "${EXISTING_GCODES}" "${SP_DATA}/gcodes"
    fi

    info "Data directory ready at ${SP_DATA}"
    echo
}

# --- Generate Config
extract_section() {
    # Extract a full INI section (including multi-line indented values)
    local section="$1"
    local file="$2"
    awk -v sec="[${section}]" '
        BEGIN { found=0 }
        $0 == sec || $0 ~ "^\\[" sec "\\]" {
            found=1
            print
            next
        }
        found && /^\[/ { exit }
        found { print }
    ' "${file}"
}

generate_config() {
    step "Generating moonraker.conf..."

    if [[ -f "${SP_CONFIG}" ]]; then
        info "Config already exists - preserving (tokens/settings intact)"
        echo
        return
    fi

    # Extract klippy_uds_address from existing config
    local uds_addr
    uds_addr=$(awk '/^klippy_uds_address[[:space:]]*[:=]/ {
        sub(/^klippy_uds_address[[:space:]]*[:=][[:space:]]*/, "")
        print
        exit
    }' "${EXISTING_CONFIG}")
    uds_addr="${uds_addr:-/tmp/klippy_uds}"

    # Extract authorization section
    local auth_section
    auth_section=$(extract_section "authorization" "${EXISTING_CONFIG}")

    # Build config
    cat > "${SP_CONFIG}" <<CONF
# Moonraker configuration - SimplyPrint dedicated instance
# Auto-generated by moonraker-sp.sh - $(date -Iseconds)

[server]
host: 0.0.0.0
port: ${SP_PORT}
klippy_uds_address: ${uds_addr}

${auth_section}

[machine]
provider: systemd_dbus
validate_service: False

[file_manager]

[simplyprint]
CONF

    info "Config written to ${SP_CONFIG}"
    echo
}

# --- Environment File
generate_env_file() {
    step "Generating environment file..."

    cat > "${SP_ENV_FILE}" <<EOF
MOONRAKER_DATA_PATH="${SP_DATA}"
MOONRAKER_ARGS="-m moonraker"
PYTHONPATH="${SP_REPO}"
EOF

    info "Environment file written to ${SP_ENV_FILE}"
    echo
}

# --- Systemd Service
install_service() {
    step "Installing systemd service..."

    local service_content
    service_content="[Unit]
Description=Moonraker API Server (SimplyPrint)
Requires=network-online.target
After=network-online.target

[Install]
WantedBy=multi-user.target

[Service]
Type=simple
User=$(whoami)
SupplementaryGroups=moonraker-admin
RemainAfterExit=yes
WorkingDirectory=${SP_REPO}
EnvironmentFile=${SP_ENV_FILE}
ExecStart=${SP_VENV}/bin/python \$MOONRAKER_ARGS
Restart=always
RestartSec=10
RestartForceExitStatus=42"

    echo "${service_content}" | sudo tee "${SP_SERVICE_FILE}" >/dev/null
    sudo systemctl daemon-reload
    sudo systemctl enable "${SP_SERVICE}" 2>/dev/null

    info "Service ${SP_SERVICE} installed and enabled"
    echo
}

# --- Start & Verify
start_and_verify() {
    step "Starting ${SP_SERVICE}..."

    sudo systemctl start "${SP_SERVICE}"

    info "Waiting for service to come up (max ${VERIFY_TIMEOUT}s)..."
    local elapsed=0
    while [[ "${elapsed}" -lt "${VERIFY_TIMEOUT}" ]]; do
        if curl -sf "http://localhost:${SP_PORT}/server/info" &>/dev/null; then
            info "Service is up and responding on port ${SP_PORT}!"
            echo
            return 0
        fi
        sleep 1
        elapsed=$((elapsed + 1))
    done

    warn "The service hasn't responded yet. It may still be starting up."
    echo
    warn "--- Recent log output (for troubleshooting) ---"
    sudo journalctl -u "${SP_SERVICE}" --no-pager -n 30 2>/dev/null || true
    echo
    local logfile="${SP_DATA}/logs/moonraker.log"
    if [[ -f "${logfile}" ]]; then
        warn "--- Log file: ${logfile} ---"
        tail -n 30 "${logfile}" 2>/dev/null || true
    fi
    echo
    echo -e "You can check the status any time by running: ${BOLD}systemctl status ${SP_SERVICE}${NC}"
    echo -e "If the problem persists, contact SimplyPrint support and share the log output above."
    return 1
}

# --- Setup Code
fetch_setup_code() {
    # Fetches the SimplyPrint setup code from moonraker's announcements API.
    local response
    response=$(curl -sf "http://localhost:${SP_PORT}/server/announcements/list" 2>/dev/null) || return 1

    # Parse JSON with python (available on the system - no jq dependency)
    echo "${response}" | python3 -c "
import sys, json
try:
    data = json.load(sys.stdin)
    for entry in data.get('result', {}).get('entries', []):
        if entry.get('feed', '') == 'simplyprint' and 'setup code' in entry.get('description', '').lower():
            desc = entry['description']
            # Format: '...setup code:\n\n{CODE}\n\n'
            parts = desc.split('\n\n')
            if len(parts) >= 2:
                print(parts[1].strip())
                sys.exit(0)
except Exception:
    pass
" 2>/dev/null
}

show_setup_code() {
    # Displays the setup code, polling briefly if needed (WebSocket may take a moment).
    local code=""
    local max_wait="${1:-10}"
    local elapsed=0

    while [[ "${elapsed}" -lt "${max_wait}" ]]; do
        code=$(fetch_setup_code)
        if [[ -n "${code}" ]]; then
            echo
            echo -e "${GREEN}============================================${NC}"
            echo -e "${GREEN}  SimplyPrint Setup Code${NC}"
            echo -e "${GREEN}============================================${NC}"
            echo
            echo -e "  Your setup code is:  ${BOLD}${code}${NC}"
            echo
            echo -e "  Go to ${CYAN}https://simplyprint.io${NC} and enter this"
            echo -e "  code to connect your printer."
            echo
            return 0
        fi
        sleep 1
        elapsed=$((elapsed + 1))
    done

    echo
    warn "No setup code available."
    echo -e "  This usually means one of:"
    echo -e "  - The printer is already linked to SimplyPrint"
    echo -e "  - The service is still starting (try again in a moment)"
    echo -e "  - The service isn't running (check with: ${BOLD}sudo service ${SP_SERVICE} status${NC})"
    echo
    local local_ip
    local_ip=$(hostname -I 2>/dev/null | awk '{print $1}')
    local_ip="${local_ip:-<your-printer-ip>}"
    echo -e "  You can also look for the setup code in the Moonraker web interface:"
    echo -e "  ${CYAN}http://${local_ip}:${SP_PORT}${NC}"
    echo
    return 1
}

do_setup_code() {
    echo
    step "Retrieving SimplyPrint setup code..."

    if ! curl -sf "http://localhost:${SP_PORT}/server/info" &>/dev/null; then
        error "The SimplyPrint Moonraker service is not responding."
        echo -e "  Make sure it's running: ${BOLD}sudo systemctl start ${SP_SERVICE}${NC}"
        echo
        return 1
    fi

    show_setup_code 5
}

# --- Print Summary
print_summary() {
    echo
    echo -e "${GREEN}============================================${NC}"
    echo -e "${GREEN}  Installation Complete!${NC}"
    echo -e "${GREEN}============================================${NC}"
    echo
    echo -e "  Service:   ${BOLD}${SP_SERVICE}${NC}"
    echo -e "  Port:      ${BOLD}${SP_PORT}${NC}"
    echo -e "  Config:    ${SP_CONFIG}"
    echo -e "  Data dir:  ${SP_DATA}"
    echo -e "  Logs:      ${SP_DATA}/logs/moonraker.log"
    echo -e "  Venv:      ${SP_VENV}"
    echo -e "  Repo:      ${SP_REPO}"
    echo
}

# --- Main Flows
do_install() {
    echo
    echo -e "${BOLD}Starting SimplyPrint Moonraker installation...${NC}"
    echo

    if is_installed; then
        warn "SimplyPrint Moonraker is already installed."
        if confirm "Update SimplyPrint Moonraker instead?"; then
            do_update
            return
        fi
        info "Aborted."
        return
    fi

    check_prerequisites
    clone_repo
    setup_venv
    init_data_dir
    generate_config
    generate_env_file
    install_service
    start_and_verify
    print_summary

    # Try to show the setup code (may take a few seconds to arrive)
    info "Waiting for SimplyPrint setup code..."
    show_setup_code 15
}

do_update() {
    echo
    step "Updating SimplyPrint Moonraker..."
    echo

    if [[ ! -d "${SP_REPO}/.git" ]]; then
        die "SimplyPrint Moonraker is not installed yet. Please run Install first."
    fi

    info "Pulling latest code..."
    cd "${SP_REPO}"
    git pull
    cd "${USER_HOME}"

    info "Updating pip dependencies..."
    SKIP_CYTHON=1 "${SP_VENV}/bin/pip" install -r "${SP_REPO}/scripts/moonraker-requirements.txt" --quiet

    info "Restarting service..."
    sudo systemctl restart "${SP_SERVICE}"

    # Verify
    info "Waiting for service to come up (max ${VERIFY_TIMEOUT}s)..."
    local elapsed=0
    while [[ "${elapsed}" -lt "${VERIFY_TIMEOUT}" ]]; do
        if curl -sf "http://localhost:${SP_PORT}/server/info" &>/dev/null; then
            info "Update complete - service is running on port ${SP_PORT}"
            return 0
        fi
        sleep 1
        elapsed=$((elapsed + 1))
    done

    warn "Service didn't respond yet. You can check its status with: systemctl status ${SP_SERVICE}"
    return 1
}

do_uninstall() {
    echo
    step "Uninstalling SimplyPrint Moonraker..."
    echo

    if ! confirm "This will remove the SimplyPrint moonraker instance, its dependencies, and data. Continue?"; then
        info "Aborted."
        return
    fi

    if systemctl is-active --quiet "${SP_SERVICE}" 2>/dev/null; then
        info "Stopping ${SP_SERVICE}..."
        sudo systemctl stop "${SP_SERVICE}"
    fi

    if systemctl is-enabled --quiet "${SP_SERVICE}" 2>/dev/null; then
        info "Disabling ${SP_SERVICE}..."
        sudo systemctl disable "${SP_SERVICE}" 2>/dev/null
    fi

    if [[ -f "${SP_SERVICE_FILE}" ]]; then
        info "Removing service file..."
        sudo rm -f "${SP_SERVICE_FILE}"
        sudo systemctl daemon-reload
    fi

    for dir in "${SP_REPO}" "${SP_VENV}" "${SP_DATA}"; do
        if [[ -e "${dir}" ]]; then
            info "Removing ${dir}..."
            rm -rf "${dir}"
        fi
    done

    echo
    info "Uninstall complete."
    echo
}

do_status() {
    echo
    echo -e "${BOLD}SimplyPrint Moonraker Status${NC}"
    echo "--------------------------------------"

    # Service status
    if systemctl is-active --quiet "${SP_SERVICE}" 2>/dev/null; then
        echo -e "  Service:  ${GREEN}active (running)${NC}"
    elif systemctl is-enabled --quiet "${SP_SERVICE}" 2>/dev/null; then
        echo -e "  Service:  ${YELLOW}inactive (enabled but stopped)${NC}"
    elif [[ -f "${SP_SERVICE_FILE}" ]]; then
        echo -e "  Service:  ${RED}inactive (disabled)${NC}"
    else
        echo -e "  Service:  ${RED}not installed${NC}"
        echo
        return
    fi

    # Port
    if ss -tlnp 2>/dev/null | grep -q ":${SP_PORT} "; then
        echo -e "  Port:     ${GREEN}${SP_PORT} (listening)${NC}"
    else
        echo -e "  Port:     ${YELLOW}${SP_PORT} (not listening)${NC}"
    fi

    # API check
    if curl -sf "http://localhost:${SP_PORT}/server/info" &>/dev/null; then
        echo -e "  API:      ${GREEN}responding${NC}"
    else
        echo -e "  API:      ${RED}not responding${NC}"
    fi

    # Paths
    echo -e "  Repo:     ${SP_REPO} $(test -d "${SP_REPO}/.git" && echo -e "${GREEN}✓${NC}" || echo -e "${RED}✗${NC}")"
    echo -e "  Venv:     ${SP_VENV} $(test -d "${SP_VENV}" && echo -e "${GREEN}✓${NC}" || echo -e "${RED}✗${NC}")"
    echo -e "  Data:     ${SP_DATA} $(test -d "${SP_DATA}" && echo -e "${GREEN}✓${NC}" || echo -e "${RED}✗${NC}")"

    # Disk usage
    if [[ -d "${SP_REPO}" ]] || [[ -d "${SP_VENV}" ]] || [[ -d "${SP_DATA}" ]]; then
        local usage
        usage=$(du -shc "${SP_REPO}" "${SP_VENV}" "${SP_DATA}" 2>/dev/null | tail -1 | awk '{print $1}')
        echo -e "  Disk:     ${usage}"
    fi

    echo
}

# --- Restart
do_restart() {
    if ! is_installed; then
        warn "SimplyPrint Moonraker is not installed."
        echo
        return 1
    fi

    info "Restarting SimplyPrint Moonraker..."
    sudo systemctl restart "${SP_SERVICE}"

    # Wait briefly and verify it came back up
    sleep 2
    if sudo systemctl status "${SP_SERVICE}" &>/dev/null; then
        info "SimplyPrint Moonraker restarted successfully."
    else
        error "SimplyPrint Moonraker failed to start after restart."
        echo -e "  Check the logs: ${BOLD}sudo systemctl status ${SP_SERVICE}${NC}"
    fi
    echo
}

# --- Menu
show_menu() {
    while true; do
        echo
        echo "============================================"
        echo "  Moonraker SimplyPrint Instance Manager"
        echo "============================================"
        echo
        echo "  1) Install"
        echo "  2) Update"
        echo "  3) Uninstall"
        echo "  4) Restart"
        echo "  5) Status"
        echo "  6) Setup Code"
        echo "  7) Exit"
        echo
        read -rp "  Choose an option [1-7]: " choice
        case "${choice}" in
            1) do_install ;;
            2) do_update ;;
            3) do_uninstall ;;
            4) do_restart ;;
            5) do_status ;;
            6) do_setup_code ;;
            7) echo "Bye!"; exit 0 ;;
            *) warn "Invalid choice" ;;
        esac
    done
}

# --- CLI Arguments
main() {
    local action=""
    YES="${YES:-0}"

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --install)    action="install" ;;
            --update)     action="update" ;;
            --uninstall)  action="uninstall" ;;
            --restart)    action="restart" ;;
            --status)     action="status" ;;
            --setup-code) action="setup-code" ;;
            --yes|-y)    YES=1 ;;
            --help|-h)
                echo "Usage: $0 [--install|--update|--uninstall|--restart|--status|--setup-code] [--yes]"
                exit 0
                ;;
            *)
                die "Unknown argument: $1 (try --help)"
                ;;
        esac
        shift
    done

    if [[ -n "${action}" ]]; then
        case "${action}" in
            install)    do_install ;;
            update)     do_update ;;
            uninstall)  do_uninstall ;;
            restart)    do_restart ;;
            status)     do_status ;;
            setup-code) do_setup_code ;;
        esac
    else
        show_menu
    fi
}

main "$@"