#!/usr/bin/env bash
#
# TrueNAS Scale Collection Script v1.1.0
# Collects TrueNAS Scale infrastructure configuration via API
#

set -euo pipefail

# Default Configuration
TRUENAS_HOST="${TRUENAS_HOST:-192.168.2.150}"
TRUENAS_API_KEY="${TRUENAS_API_KEY:-}"
TIMESTAMP="$(date +%Y%m%d-%H%M%S)"
OUTPUT_DIR=""
COLLECTION_LEVEL="standard"

# Colors
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
BLUE='\033[0;34m'; CYAN='\033[0;36m'; NC='\033[0m'

# Counters
COLLECTED=0; SKIPPED=0; ERRORS=0

# Usage/Help function
usage() {
    cat << 'EOF'
Usage: collect-truenas-config.sh [OPTIONS] [OUTPUT_DIR]

TrueNAS Scale configuration collection script.

OPTIONS:
    -l, --level LEVEL       Collection level: basic, standard, full, paranoid
                            (default: standard)
    -o, --output DIR        Output directory path
                            (default: ./truenas-export-YYYYMMDD-HHMMSS)
    -h, --host HOST         TrueNAS host IP or hostname
                            (default: $TRUENAS_HOST or 192.168.2.150)
    --help                  Show this help message

ENVIRONMENT VARIABLES:
    TRUENAS_HOST           TrueNAS host (default: 192.168.2.150)
    TRUENAS_API_KEY        TrueNAS API key (required)
    COLLECTION_LEVEL       Default collection level (default: standard)

EXAMPLES:
    # Standard collection with named arguments
    collect-truenas-config.sh --level standard --output ./truenas-exports

    # Full collection to disaster recovery directory
    collect-truenas-config.sh --level full --output ../disaster-recovery/

    # Backward compatible (positional argument for output directory)
    collect-truenas-config.sh ./my-export-dir

    # Custom host
    collect-truenas-config.sh --host 192.168.2.151 --level basic --output ./exports

COLLECTION LEVELS:
    basic      - System info, storage, shares, network, services
    standard   - Basic + tasks and users (default)
    full       - Standard + SMART data and metrics
    paranoid   - Full + all available diagnostics

EOF
    exit 0
}

# Argument Parsing
parse_arguments() {
    # Handle empty arguments
    if [[ $# -eq 0 ]]; then
        OUTPUT_DIR="./truenas-export-${TIMESTAMP}"
        return
    fi

    # Check for help first
    for arg in "$@"; do
        if [[ "$arg" == "--help" ]]; then
            usage
        fi
    done

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -l|--level)
                if [[ -z "${2:-}" ]]; then
                    echo "Error: --level requires an argument" >&2
                    usage
                fi
                COLLECTION_LEVEL="$2"
                shift 2
                ;;
            -o|--output)
                if [[ -z "${2:-}" ]]; then
                    echo "Error: --output requires an argument" >&2
                    usage
                fi
                OUTPUT_DIR="$2"
                shift 2
                ;;
            -h|--host)
                if [[ -z "${2:-}" ]]; then
                    echo "Error: --host requires an argument" >&2
                    usage
                fi
                TRUENAS_HOST="$2"
                shift 2
                ;;
            --help)
                usage
                ;;
            -*)
                echo "Error: Unknown option: $1" >&2
                usage
                ;;
            *)
                # Backward compatibility: positional argument for output directory
                if [[ -z "$OUTPUT_DIR" ]]; then
                    OUTPUT_DIR="$1"
                    shift
                else
                    echo "Error: Unexpected argument: $1" >&2
                    usage
                fi
                ;;
        esac
    done

    # Set defaults if not provided
    if [[ -z "$OUTPUT_DIR" ]]; then
        OUTPUT_DIR="./truenas-export-${TIMESTAMP}"
    fi

    # Validate collection level
    case "$COLLECTION_LEVEL" in
        basic|standard|full|paranoid)
            # Valid level
            ;;
        *)
            echo "Error: Invalid collection level: $COLLECTION_LEVEL" >&2
            echo "Valid levels: basic, standard, full, paranoid" >&2
            exit 1
            ;;
    esac
}

# Parse arguments before setting up API configuration
parse_arguments "$@"

# Now validate API key and set up API base
TRUENAS_API_KEY="${TRUENAS_API_KEY:?API key required. Set TRUENAS_API_KEY environment variable}"
TRUENAS_API_BASE="https://${TRUENAS_HOST}/api/v2.0"

# Logging
log() {
    local level=$1; shift
    case "$level" in
        INFO) echo -e "${BLUE}[INFO]${NC} $*" ;;
        OK)   echo -e "${GREEN}[✓]${NC} $*"; ((COLLECTED++)) ;;
        WARN) echo -e "${YELLOW}[WARN]${NC} $*"; ((SKIPPED++)) ;;
        ERROR) echo -e "${RED}[ERROR]${NC} $*" >&2; ((ERRORS++)) ;;
    esac
}

# API call wrapper with timeout protection
api_call() {
    local endpoint=$1 output=$2 desc=$3
    mkdir -p "$(dirname "$output")"

    log INFO "Fetching: ${endpoint}"

    local code
    local start_time=$(date +%s)

    # FIXED: Added connection timeout (10s) and max time (30s) to prevent hangs
    code=$(timeout 30 curl -s -w "%{http_code}" -X GET \
        --connect-timeout 10 \
        --max-time 30 \
        "${TRUENAS_API_BASE}${endpoint}" \
        -H "Authorization: Bearer ${TRUENAS_API_KEY}" \
        -H "Content-Type: application/json" \
        --insecure -o "$output" 2>/dev/null || echo "000")

    local end_time=$(date +%s)
    local duration=$((end_time - start_time))

    if [[ "$code" == "200" ]]; then
        # FIXED: Validate JSON response (only if jq is available)
        if command -v jq &>/dev/null; then
            if jq empty "$output" 2>/dev/null; then
                log OK "$desc (${duration}s)"
                return 0
            else
                log ERROR "$desc - Invalid JSON response"
                rm -f "$output"
                return 1
            fi
        else
            # jq not available, skip validation
            log OK "$desc (${duration}s)"
            return 0
        fi
    elif [[ "$code" == "000" ]]; then
        log ERROR "$desc - Connection timeout or failure after ${duration}s (endpoint: ${endpoint})"
        rm -f "$output"
        return 1
    else
        log WARN "$desc (HTTP $code after ${duration}s, endpoint: ${endpoint})"
        rm -f "$output"
        return 1
    fi
}

# Main collection
main() {
    echo -e "${CYAN}======================================${NC}"
    echo -e "${CYAN}TrueNAS Scale Collection v1.1.0${NC}"
    echo -e "${CYAN}======================================${NC}"
    echo
    log INFO "Host: $TRUENAS_HOST"
    log INFO "Level: $COLLECTION_LEVEL"
    log INFO "Output: $OUTPUT_DIR"
    echo

    # Create directory structure
    mkdir -p "$OUTPUT_DIR"/{configs/{system,storage,sharing,network,services,tasks,users},exports/{storage,system,logs},metrics}

    # System Information
    echo -e "${CYAN}=== System Information ===${NC}"
    api_call "/system/version" "$OUTPUT_DIR/exports/system/version.json" "TrueNAS version" || true
    api_call "/system/info" "$OUTPUT_DIR/exports/system/info.json" "System information" || true
    api_call "/system/general" "$OUTPUT_DIR/configs/system/general.json" "General config" || true
    api_call "/system/advanced" "$OUTPUT_DIR/configs/system/advanced.json" "Advanced config" || true
    echo

    # Storage
    echo -e "${CYAN}=== Storage ===${NC}"
    api_call "/pool" "$OUTPUT_DIR/exports/storage/pools.json" "Storage pools" || true
    api_call "/pool/dataset" "$OUTPUT_DIR/exports/storage/datasets.json" "Datasets" || true
    api_call "/disk" "$OUTPUT_DIR/exports/storage/disks.json" "Disk inventory" || true
    api_call "/zfs/snapshot" "$OUTPUT_DIR/exports/storage/snapshots.json" "ZFS snapshots" || true
    echo

    # Shares
    echo -e "${CYAN}=== Shares ===${NC}"
    api_call "/sharing/nfs" "$OUTPUT_DIR/configs/sharing/nfs.json" "NFS shares" || true
    api_call "/sharing/smb" "$OUTPUT_DIR/configs/sharing/smb.json" "SMB shares" || true
    api_call "/iscsi/target" "$OUTPUT_DIR/configs/sharing/iscsi-targets.json" "iSCSI targets" || true
    echo

    # Network
    echo -e "${CYAN}=== Network ===${NC}"
    api_call "/interface" "$OUTPUT_DIR/configs/network/interfaces.json" "Network interfaces" || true
    api_call "/network/configuration" "$OUTPUT_DIR/configs/network/config.json" "Network config" || true
    api_call "/staticroute" "$OUTPUT_DIR/configs/network/routes.json" "Static routes" || true
    echo

    # Services
    echo -e "${CYAN}=== Services ===${NC}"
    api_call "/service" "$OUTPUT_DIR/configs/services/status.json" "Services status" || true
    api_call "/ssh" "$OUTPUT_DIR/configs/services/ssh.json" "SSH config" || true
    echo

    # Tasks
    if [[ "$COLLECTION_LEVEL" != "basic" ]]; then
        echo -e "${CYAN}=== Tasks ===${NC}"
        api_call "/cronjob" "$OUTPUT_DIR/configs/tasks/cron.json" "Cron jobs" || true
        api_call "/pool/snapshottask" "$OUTPUT_DIR/configs/tasks/snapshots.json" "Snapshot tasks" || true
        api_call "/replication" "$OUTPUT_DIR/configs/tasks/replication.json" "Replication" || true
        api_call "/cloudsync" "$OUTPUT_DIR/configs/tasks/cloudsync.json" "Cloud sync" || true
        echo
    fi

    # Users
    if [[ "$COLLECTION_LEVEL" != "basic" ]]; then
        echo -e "${CYAN}=== Users ===${NC}"
        api_call "/user" "$OUTPUT_DIR/configs/users/users.json" "User accounts" || true
        api_call "/group" "$OUTPUT_DIR/configs/users/groups.json" "Groups" || true
        echo
    fi

    # SMART data (full/paranoid only)
    if [[ "$COLLECTION_LEVEL" == "full" ]] || [[ "$COLLECTION_LEVEL" == "paranoid" ]]; then
        echo -e "${CYAN}=== SMART Data ===${NC}"
        api_call "/smart/test/results" "$OUTPUT_DIR/metrics/smart-results.json" "SMART test results" || true
        echo
    fi

    # Docker/Apps Collection (all levels)
    echo -e "${CYAN}=== Docker & Apps ===${NC}"
    mkdir -p "$OUTPUT_DIR/exports/apps"
    mkdir -p "$OUTPUT_DIR/configs/apps"

    # Detect SSH user (try common TrueNAS usernames)
    SSH_USER=""
    for user in root admin truenas_admin; do
        if command -v ssh &>/dev/null && ssh -o ConnectTimeout=5 -o BatchMode=yes "${user}@${TRUENAS_HOST}" "echo test" &>/dev/null 2>&1; then
            SSH_USER="$user"
            break
        fi
    done

    # Check if we can access Docker via SSH
    if [[ -n "$SSH_USER" ]]; then
        log INFO "SSH access available as ${SSH_USER}@${TRUENAS_HOST}"

        # Docker containers (convert newline-delimited JSON to proper JSON array)
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "command -v docker &>/dev/null && sudo docker ps -a --format '{{json .}}' 2>/dev/null | jq -s '.' 2>/dev/null || sudo docker ps -a --format '{{json .}}' 2>/dev/null | sed '1s/^/[/; \$!s/\$/,/; \$s/\$/]/'" > "$OUTPUT_DIR/exports/apps/docker-containers.json" 2>/dev/null && [[ -s "$OUTPUT_DIR/exports/apps/docker-containers.json" ]]; then
            log OK "Docker containers list"
        else
            log WARN "Docker containers (command failed or no containers)"
        fi

        # Docker human-readable list
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "command -v docker &>/dev/null && sudo docker ps -a 2>/dev/null" > "$OUTPUT_DIR/exports/apps/docker-containers.txt" 2>/dev/null && [[ -s "$OUTPUT_DIR/exports/apps/docker-containers.txt" ]]; then
            log OK "Docker containers (text format)"
        else
            log WARN "Docker containers text list (command failed)"
        fi

        # Docker images
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "command -v docker &>/dev/null && sudo docker images --format '{{json .}}' 2>/dev/null | jq -s '.' 2>/dev/null || sudo docker images --format '{{json .}}' 2>/dev/null | sed '1s/^/[/; \$!s/\$/,/; \$s/\$/]/'" > "$OUTPUT_DIR/exports/apps/docker-images.json" 2>/dev/null && [[ -s "$OUTPUT_DIR/exports/apps/docker-images.json" ]]; then
            log OK "Docker images"
        else
            log WARN "Docker images (command failed)"
        fi

        # Docker networks
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "command -v docker &>/dev/null && sudo docker network ls --format '{{json .}}' 2>/dev/null | jq -s '.' 2>/dev/null || sudo docker network ls --format '{{json .}}' 2>/dev/null | sed '1s/^/[/; \$!s/\$/,/; \$s/\$/]/'" > "$OUTPUT_DIR/exports/apps/docker-networks.json" 2>/dev/null && [[ -s "$OUTPUT_DIR/exports/apps/docker-networks.json" ]]; then
            log OK "Docker networks"
        else
            log WARN "Docker networks (command failed)"
        fi

        # Docker volumes
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "command -v docker &>/dev/null && sudo docker volume ls --format '{{json .}}' 2>/dev/null | jq -s '.' 2>/dev/null || sudo docker volume ls --format '{{json .}}' 2>/dev/null | sed '1s/^/[/; \$!s/\$/,/; \$s/\$/]/'" > "$OUTPUT_DIR/exports/apps/docker-volumes.json" 2>/dev/null && [[ -s "$OUTPUT_DIR/exports/apps/docker-volumes.json" ]]; then
            log OK "Docker volumes"
        else
            log WARN "Docker volumes (command failed)"
        fi

        # TrueNAS app metadata
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "sudo cat /mnt/.ix-apps/metadata.yaml 2>/dev/null" > "$OUTPUT_DIR/configs/apps/metadata.yaml" 2>/dev/null && [[ -s "$OUTPUT_DIR/configs/apps/metadata.yaml" ]]; then
            log OK "TrueNAS app metadata"
        else
            log WARN "TrueNAS app metadata (file not found)"
        fi

        # TrueNAS user config
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "sudo cat /mnt/.ix-apps/user_config.yaml 2>/dev/null" > "$OUTPUT_DIR/configs/apps/user_config.yaml" 2>/dev/null && [[ -s "$OUTPUT_DIR/configs/apps/user_config.yaml" ]]; then
            log OK "TrueNAS user config"
        else
            log WARN "TrueNAS user config (file not found)"
        fi

        # App configs directory listing
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "sudo ls -laR /mnt/.ix-apps/app_configs/ 2>/dev/null" > "$OUTPUT_DIR/exports/apps/app_configs_list.txt" 2>/dev/null && [[ -s "$OUTPUT_DIR/exports/apps/app_configs_list.txt" ]]; then
            log OK "App configs directory listing"
        else
            log WARN "App configs listing (directory not accessible)"
        fi

        # Docker storage sizes
        if ssh "${SSH_USER}@${TRUENAS_HOST}" "sudo du -sh /mnt/.ix-apps/docker/* 2>/dev/null" > "$OUTPUT_DIR/exports/apps/docker_sizes.txt" 2>/dev/null && [[ -s "$OUTPUT_DIR/exports/apps/docker_sizes.txt" ]]; then
            log OK "Docker storage sizes"
        else
            log WARN "Docker storage sizes (directory not accessible)"
        fi
    else
        log WARN "Docker/Apps collection (SSH not configured - see manual collection instructions)"
        cat > "$OUTPUT_DIR/exports/apps/MANUAL_COLLECTION.md" << 'MANUAL_EOF'
# Manual Docker/Apps Collection Instructions

SSH access is not configured. To collect Docker and app information, run these commands on your TrueNAS server:

```bash
# 1. Export Docker containers
sudo docker ps -a --format '{{json .}}' | jq -s '.' > /tmp/docker-containers.json

# 2. Export TrueNAS app metadata
sudo cp /mnt/.ix-apps/metadata.yaml /tmp/
sudo cp /mnt/.ix-apps/user_config.yaml /tmp/

# 3. List app configs
sudo ls -laR /mnt/.ix-apps/app_configs/ > /tmp/app_configs_list.txt

# 4. Copy files to collection directory
# Transfer these files to your disaster-recovery export
```

Then copy the files to: exports/apps/
MANUAL_EOF
        log INFO "Created manual collection instructions"
    fi
    echo

    # Generate summary
    cat > "$OUTPUT_DIR/SUMMARY.md" << EOF
# TrueNAS Scale Export Summary

**Date**: $(date '+%Y-%m-%d %H:%M:%S')
**Host**: $TRUENAS_HOST
**Level**: $COLLECTION_LEVEL

## Statistics
- Collected: $COLLECTED items
- Skipped: $SKIPPED items
- Errors: $ERRORS items

## Files Created
$(find "$OUTPUT_DIR" -type f -name "*.json" | wc -l) JSON files collected

See individual files in:
- configs/ - Configuration files
- exports/ - System state exports
- metrics/ - Performance data
EOF

    # Summary
    echo
    echo -e "${CYAN}======================================${NC}"
    echo -e "${CYAN}Collection Complete${NC}"
    echo -e "${CYAN}======================================${NC}"
    echo -e "${GREEN}✓ Collected:${NC} $COLLECTED items"
    echo -e "${YELLOW}⊘ Skipped:${NC} $SKIPPED items"
    echo -e "${RED}✗ Errors:${NC} $ERRORS items"
    echo
    echo -e "${BLUE}Output:${NC} $OUTPUT_DIR"
    echo -e "${BLUE}Summary:${NC} $OUTPUT_DIR/SUMMARY.md"
    echo

    # Compress
    if command -v tar &>/dev/null; then
        local archive="${OUTPUT_DIR}.tar.gz"
        log INFO "Compressing to $archive..."
        tar -czf "$archive" -C "$(dirname "$OUTPUT_DIR")" "$(basename "$OUTPUT_DIR")" 2>/dev/null
        log INFO "Archive: $(du -h "$archive" | cut -f1)"
    fi
}

main