#!/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

    # 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