#!/bin/bash -e

# Define colors for better output readability
RED="\033[31m"
GREEN="\033[32m"
YELLOW="\033[33m"
BLUE="\033[34m"
PURPLE="\033[35m"
WHITE="\033[97m"
BOLD="\033[1m"
NC="\033[0m" # No color

# Define log levels
ERROR=0
WARNING=1
INFO=2
DEBUG=3

# Determine LOG_LEVEL based on DEBUG variable
if [ "$DO_DEBUG" = "true" ]; then
    LOG_LEVEL=$DEBUG  # LOG_LEVEL=3
else
    LOG_LEVEL=$INFO  # Default to INFO level (2)
fi

# Logging functions with timestamps for log file output
log_error() {
    if [ "$LOG_LEVEL" -ge $ERROR ]; then
        echo -e "${RED}❌ $*${NC}"
    fi
}

log_warning() {
    if [ "$LOG_LEVEL" -ge $WARNING ]; then
        echo -e "${YELLOW}⚠️  $*${NC}"
    fi
}

log_info() {
    if [ "$LOG_LEVEL" -ge $INFO ]; then
        echo -e "${WHITE}🔵 $*${NC}"
    fi
}

log_success() {
    if [ "$LOG_LEVEL" -ge $INFO ]; then
        echo -e "${GREEN}✅ $*${NC}"
    fi
}

log_debug() {
    if [ "$LOG_LEVEL" -ge $DEBUG ]; then
        echo -e "${PURPLE}🐛 $*${NC}"
    fi
}

# Check and install prerequisites
check_prerequisites() {
    local packages=("jq" "curl" "rpi-eeprom")
    local missing_packages=()

    for pkg in "${packages[@]}"; do
        if ! dpkg -l | grep -qw "$pkg"; then
            missing_packages+=("$pkg")
        fi
    done

    if [ ${#missing_packages[@]} -eq 0 ]; then
        log_success "All prerequisites are already installed."
    else
        log_error "Missing packages: ${missing_packages[*]}"
        exit 1
    fi
}


            ############# Phase 3: Shelly Device Setup #############

get_shelly_ip() {
    shelly_ip="shelly-watter"
}

# Function to check the firmware version and update if needed
check_and_update_firmware() {
    # Define the minimum version required
    minimum_version="1.4.2"

    # Get the current firmware version using Shelly.GetDeviceInfo
    device_info=$(curl -s "http://$shelly_ip/rpc/Shelly.GetDeviceInfo")

    if [ -z "$device_info" ]; then
        log_error "Unable to fetch device info from Shelly. Please check the IP address or device status."
        return
    fi

    # Extract the current firmware version
    current_version=$(echo "$device_info" | jq -r '.ver')

    if [ "$current_version" != "null" ] && [ -n "$current_version" ]; then
        # Color-code the firmware version based on its value
        if dpkg --compare-versions "$current_version" ge "$minimum_version"; then
            firmware_color="$GREEN"
            log_info "Current firmware version: ${firmware_color}$current_version${NC}"
            log_success "Firmware is up to date (version $current_version). No update required."
            return
        elif dpkg --compare-versions "$current_version" ge "1.3.0"; then
            firmware_color="$YELLOW"
            log_warning "Current firmware version: ${firmware_color}$current_version${NC}"
            log_warning "Firmware is below the minimum required version ($minimum_version). Updating firmware..."
        else
            firmware_color="$RED"
            log_error "Current firmware version: ${firmware_color}$current_version${NC}"
            log_error "Firmware is too old. Please update manually."
            exit 1
        fi
    else
        log_error "Firmware version not found. Please check the device information."
        return
    fi

    # If the firmware version is less than the required version, check for updates
    log_info "Checking for firmware updates..."

    # Check for firmware update using Shelly.CheckForUpdate
    for attempt in {1..3}; do
        update_info=$(curl -s "http://$shelly_ip/rpc/Shelly.CheckForUpdate")
        stable_version=$(echo "$update_info" | jq -r '.stable.version')

        if [ "$stable_version" != "null" ] && [ -n "$stable_version" ]; then
            log_info "Firmware update available: $stable_version. Starting update..."

            # Trigger firmware update but discard any output
            curl -s "http://$shelly_ip/rpc/Shelly.Update" > /dev/null

            # Wait for the update to complete
            log_info "Waiting for firmware update to complete..."
            sleep 90  # Wait for 90 seconds for the update to be applied
            break
        else
            log_warning "No update found on attempt $attempt. Retrying..."
            sleep 2
        fi
    done

    if [ "$stable_version" == "null" ] || [ -z "$stable_version" ]; then
        log_info "No updates available or device is already running the latest firmware."
    fi
}

install_certs() {
    # Paths to device certificate, private key, and CA certificate
    local cert_path="controller-cert.pem"
    local key_path="controller-key.pem"
    local ca_path="rootCA.pem"

    # Check each file before uploading
    if [ ! -f "$cert_path" ]; then
        log_error "Device Certificate file not found at $cert_path. Please ensure the file exists."
        exit 1
    fi

    if [ ! -f "$key_path" ]; then
        log_error "Private Key file not found at $key_path. Please ensure the file exists."
        exit 1
    fi

    if [ ! -f "$ca_path" ]; then
        log_error "CA Certificate file not found at $ca_path. Please ensure the file exists."
        exit 1
    fi

    sudo rm -f /etc/watter/iot/{client.pem.crt,client.private.key,AmazonRootCA1.pem}
    sudo mkdir -p /etc/watter/iot
    sudo cp $cert_path /etc/watter/iot/client.pem.crt
    sudo cp $key_path /etc/watter/iot/client.private.key
    sudo cp $ca_path /etc/watter/iot/AmazonRootCA1.pem

    log_success "Copied RPi certs to /etc/watter/iot/"

    watter_config=$(printf '{
        "iot": {
	    "endpoint": {
	        "host": {
		    "production": "%s",
		}
            },
            "production": true,
            "device_id": "%s"
        }
    }' "$endpoint_url" "$thing_name")

    jq ". + $watter_config" /etc/watter/config.json > config.json
    sudo cp config.json /etc/watter/config.json
    rm -f config.json

    log_success "Updated /etc/watter/config.json"

    echo "pi-$thing_name" | sudo tee /etc/hostname > /dev/null 2>&1
    sudo hostname -F /etc/hostname
    sudo sed -i -e "s/^127\.0\.1\.1.*/127.0.1.1 pi-$thing_name/" /etc/hosts

    log_success "Updated /etc/hosts and /etc/hostname"
}

# Upload certificates to Shelly Device
upload_certs_to_shelly() {
    # Paths to device certificate, private key, and CA certificate
    local cert_path="shelly-cert.pem"
    local key_path="shelly-key.pem"
    local ca_path="rootCA.pem"

    # Check each file before uploading
    if [ ! -f "$cert_path" ]; then
        log_error "Device Certificate file not found at $cert_path. Please ensure the file exists."
        exit 1
    fi

    if [ ! -f "$key_path" ]; then
        log_error "Private Key file not found at $key_path. Please ensure the file exists."
        exit 1
    fi

    if [ ! -f "$ca_path" ]; then
        log_error "CA Certificate file not found at $ca_path. Please ensure the file exists."
        exit 1
    fi

    # Function to upload the certificate in one request
    upload_certificate_single() {
        local file_path=$1
        local endpoint=$2

        # Read the entire file content, including newlines
        local cert_data
        cert_data=$(<"$file_path")

        # Construct JSON payload using jq, ensuring proper JSON escaping
        local json_payload
        json_payload=$(jq -n \
            --arg id "1" \
            --arg method "$endpoint" \
            --arg data "$cert_data" \
            '{id: ($id | tonumber), method: $method, params: {data: $data, append: false}}')

        # For debugging: Print the payload
        log_debug "Uploading full certificate for $endpoint:"
        log_debug "$(echo "$json_payload" | jq)"

        # Send the POST request
        response=$(curl -s -X POST -H "Content-Type: application/json" \
        -d "$json_payload" \
        "http://$shelly_ip/rpc")

        # Check for success
        log_debug "Response: $response"
        if [[ "$response" != *"len"* ]]; then
            log_error "Failed to upload $file_path with response: $response"
            exit 1
        fi

        log_success "${BLUE}${BOLD}$(basename "$file_path")${NC}${GREEN} uploaded successfully!${NC}"
    }

    # Upload CA Certificate using Shelly.PutUserCA
    log_info "Uploading CA Certificate..."
    upload_certificate_single "$ca_path" "Shelly.PutUserCA"

    # Upload Device Certificate using Shelly.PutTLSClientCert
    log_info "Uploading Device Certificate..."
    upload_certificate_single "$cert_path" "Shelly.PutTLSClientCert"

    # Upload Private Key using Shelly.PutTLSClientKey
    log_info "Uploading Private Key..."
    upload_certificate_single "$key_path" "Shelly.PutTLSClientKey"

    log_success "Certificates uploaded successfully to Shelly device."
}

configure_shelly_name() {
    shelly_config=$(printf '{
        "id": 1,
        "method": "Sys.SetConfig",
        "params": {
            "config": {
                "device": {
                    "name": "Shelly-%s"
                }
            }
        }
    }' "$thing_name")

    # Send the configuration to the Shelly device
    log_info "Setting Shelly device name..."
    response=$(curl -s -X POST -H "Content-Type: application/json" -d "$shelly_config" "http://$shelly_ip/rpc")

    # Display the response under debug level
    log_debug "Response from Shelly device:"
    log_debug "$(echo "$response" | jq . || echo "$response")"

    # Check for errors in the response
    if echo "$response" | grep -q '"error"'; then
        log_error "Failed to set device name. Response: $response"
    else
        log_success "Device name changed successfully."
    fi
}

# Configure MQTT settings on Shelly Device
configure_shelly_mqtt() {
    # Construct the JSON payload for the MQTT configuration using printf
    mqtt_config=$(printf '{
        "id": 1,
        "method": "Mqtt.SetConfig",
        "params": {
            "config": {
                "enable": true,
                "server": "%s",
                "client_id": "Shelly-%s",
                "user": null,
                "ssl_ca": "ca.pem",
                "topic_prefix": "devices/%s/shelly",
                "rpc_ntf": false,
                "status_ntf": true,
                "use_client_cert": true,
                "enable_control": false
            }
        }
    }' "$endpoint_url" "$thing_name" "$thing_name")

    # Display the configuration payload under debug level
    log_debug "MQTT Configuration being sent to Shelly device:"
    log_debug "$(echo "$mqtt_config" | jq)"

    # Send the configuration to the Shelly device
    log_info "Sending MQTT configuration to Shelly device..."
    response=$(curl -s -X POST -H "Content-Type: application/json" -d "$mqtt_config" "http://$shelly_ip/rpc")

    # Display the response under debug level
    log_debug "Response from Shelly device:"
    log_debug "$(echo "$response" | jq . || echo "$response")"

    # Check for errors in the response
    if echo "$response" | grep -q '"error"'; then
        log_error "Failed to apply MQTT configuration. Response: $response"
    else
        log_success "MQTT configuration applied successfully."
    fi
}

# Reboot the Shelly device to apply changes
reboot_shelly_device() {
    log_info "Rebooting Shelly device to apply changes..."
    response=$(curl -s "http://$shelly_ip/rpc/Shelly.Reboot")

    # Print the response under debug level
    log_debug "Response from reboot command: $response"

    log_info "Waiting for the device to come back online..."

    # Sleep initially to give the device some time to reboot
    sleep 5  # Adjust this as needed

    # Check if the device is back online by attempting to get device info
    max_retries=20  # Number of times to retry (20 retries = 100 seconds total if sleep is 1 second)
    retries=0
    while true; do
        # Try to fetch the device info
        response=$(curl -s "http://$shelly_ip/rpc/Shelly.GetDeviceInfo")
        if [ -n "$response" ] && echo "$response" | jq -e '.id' > /dev/null; then
            log_success "Shelly device rebooted and is back online."
            break
        else
            retries=$((retries + 1))
            if [ "$retries" -ge "$max_retries" ]; then
                log_error "Shelly device did not come back online within the expected time."
                exit 1
            fi
            log_debug "Device not online yet. Retrying in 1 second... (Attempt $retries/$max_retries)"
            sleep 1  # Wait for 1 second before retrying
        fi
    done
}

download_hashball() {
    sno=$(cat /sys/firmware/devicetree/base/serial-number | tr -d '\0')
    log_info "Serial Number: $sno"
    sno_hash=$(echo -n $sno | sha256sum | awk '{print $1}')
    log_info "SN Hash: $sno_hash"

    mkdir -p provision
    cd provision
    curl -s -o certs.tar.gz "https://archive.watter.com/hashes/$sno_hash"
    log_success "Downloaded hashball"

    tar zxf certs.tar.gz
    thing_name=$(cat device-id)
    endpoint_url=$(cat endpoint.url)
    log_success "Unpacked hashball for: $thing_name"
}

            ############# Phase 4: EEPROM Configuration #############

IMAGER_REPO_URL="https://archive.watter.com/images/apollo/os_list.json"

set_imager_repo_url() {
    if ! command -v rpi-eeprom-config &> /dev/null; then
        log_error "rpi-eeprom-config not found. Is rpi-eeprom installed?"
        exit 1
    fi

    log_info "Reading current EEPROM bootloader config..."
    local boot_conf
    boot_conf=$(mktemp)

    sudo rpi-eeprom-config > "$boot_conf"

    # Check if IMAGER_REPO_URL is already set correctly
    if grep -q "^IMAGER_REPO_URL=${IMAGER_REPO_URL}$" "$boot_conf"; then
        log_success "IMAGER_REPO_URL already set correctly."
        rm -f "$boot_conf"
        return
    fi

    # Remove any existing IMAGER_REPO_URL line and append the new one
    sed -i '/^IMAGER_REPO_URL=/d' "$boot_conf"
    echo "IMAGER_REPO_URL=${IMAGER_REPO_URL}" >> "$boot_conf"

    log_info "Applying EEPROM config with IMAGER_REPO_URL..."
    log_debug "IMAGER_REPO_URL=${IMAGER_REPO_URL}"

    sudo rpi-eeprom-config --apply "$boot_conf"

    rm -f "$boot_conf"
    log_success "EEPROM config updated. IMAGER_REPO_URL will take effect after reboot."
}

            ############# Phase 5: Shelly Watchdog Script #############

WATCHDOG_SCRIPT_NAME="rpi-watchdog"

# The watchdog script that runs on the Shelly device.
# It checks RPi liveness via a systemd-managed health check endpoint
# on port 8080 (watter-collector-healthcheck.service). This endpoint
# is independent of collectd, so stopping collectd for maintenance
# will not trigger a reboot.
build_watchdog_script() {
    local rpi_ip="$1"
    cat <<JSEOF
// RPi watchdog: reboot RPi via relay toggle if it stops responding.
// Uses a systemd socket-activated health check on port 8080 that is
// always available as long as the OS is running, independent of collectd.

let CONFIG = {
  endpoint: "http://${rpi_ip}:8080",
  httpTimeout: 5,
  toggleTime: 3,
  pingTime: 300,
  switchId: 100,
};

let pingTimer = null;

function pingEndpoints() {
  Shelly.call(
    "HTTP.GET",
    { url: CONFIG.endpoint, timeout: CONFIG.httpTimeout },
    function (response, error_code, error_message) {
      if (error_code !== 0) {
        print("Failed to ping ", CONFIG.endpoint, " err=", error_code, " msg=", error_message);
        Shelly.call(
          "Switch.Set",
          { id: CONFIG.switchId, on: false, toggle_after: CONFIG.toggleTime },
          function () {}
        );
      }
    }
  );
}

print("Start RPi watchdog timer");
pingTimer = Timer.set(CONFIG.pingTime * 1000, true, pingEndpoints);
JSEOF
}

# Check if Shelly has the reset button switch (ID 100)
shelly_has_reset_button() {
    local response
    response=$(curl -s "http://$shelly_ip/rpc/Switch.GetConfig?id=100")
    if echo "$response" | jq -e '.id' > /dev/null 2>&1; then
        return 0
    fi
    return 1
}

# Remove any existing watchdog script by name
remove_existing_watchdog() {
    local scripts
    scripts=$(curl -s "http://$shelly_ip/rpc/Script.List")
    log_debug "Existing scripts: $scripts"

    local script_id
    script_id=$(echo "$scripts" | jq -r \
        --arg name "$WATCHDOG_SCRIPT_NAME" \
        '.scripts[]? | select(.name == $name) | .id')

    if [ -n "$script_id" ]; then
        log_info "Removing existing watchdog script (id=$script_id)..."
        # Stop if running
        curl -s "http://$shelly_ip/rpc/Script.Stop?id=$script_id" > /dev/null 2>&1
        curl -s "http://$shelly_ip/rpc/Script.Delete?id=$script_id" > /dev/null 2>&1
        log_success "Removed old watchdog script."
    fi
}

deploy_watchdog_script() {
    if ! shelly_has_reset_button; then
        log_info "Shelly does not have reset button (Switch 100). Skipping watchdog."
        return
    fi

    log_info "Shelly has reset button. Deploying RPi watchdog script..."

    # Get the RPi's IP address on the local network
    local rpi_ip
    rpi_ip=$(hostname -I | awk '{print $1}')
    if [ -z "$rpi_ip" ]; then
        log_error "Could not determine RPi IP address."
        return
    fi
    log_info "RPi IP: $rpi_ip"

    # Remove any previous version
    remove_existing_watchdog

    # Create a new script
    local create_response
    create_response=$(curl -s -X POST -H "Content-Type: application/json" \
        -d "{\"id\":1,\"method\":\"Script.Create\",\"params\":{\"name\":\"$WATCHDOG_SCRIPT_NAME\"}}" \
        "http://$shelly_ip/rpc")
    log_debug "Script.Create response: $create_response"

    local script_id
    script_id=$(echo "$create_response" | jq -r '.result.id // empty')
    if [ -z "$script_id" ]; then
        log_error "Failed to create script. Response: $create_response"
        return
    fi
    log_info "Created script id=$script_id"

    # Upload the code
    local script_code
    script_code=$(build_watchdog_script "$rpi_ip")

    local put_payload
    put_payload=$(jq -n \
        --argjson id "$script_id" \
        --arg code "$script_code" \
        '{id: 1, method: "Script.PutCode", params: {id: $id, code: $code}}')

    local put_response
    put_response=$(curl -s -X POST -H "Content-Type: application/json" \
        -d "$put_payload" \
        "http://$shelly_ip/rpc")
    log_debug "Script.PutCode response: $put_response"

    if echo "$put_response" | grep -q '"error"'; then
        log_error "Failed to upload watchdog script. Response: $put_response"
        return
    fi
    log_success "Uploaded watchdog script code."

    # Preflight: ask the Shelly to ping the RPi healthcheck endpoint. If it
    # can't reach us at provisioning time, starting the watchdog now would
    # immediately begin tripping the reset relay every 300s. Install the
    # script either way (so the operator can enable it after fixing the
    # network), but only auto-start when reachability is verified.
    local enable_watchdog="false"
    if shelly_can_reach_rpi "$rpi_ip"; then
        log_success "Preflight succeeded; watchdog will be enabled."
        enable_watchdog="true"
    else
        log_warning "Preflight failed: Shelly cannot reach RPi at $rpi_ip:8080."
        log_warning "Watchdog script uploaded but left DISABLED to prevent reboot loops."
        log_warning "Fix the network issue, then enable the script via the Shelly UI."
    fi

    # Set auto-start flag based on preflight outcome
    local config_payload
    config_payload=$(jq -n \
        --argjson id "$script_id" \
        --argjson enable "$enable_watchdog" \
        '{id: 1, method: "Script.SetConfig", params: {id: $id, config: {enable: $enable}}}')

    curl -s -X POST -H "Content-Type: application/json" \
        -d "$config_payload" \
        "http://$shelly_ip/rpc" > /dev/null

    if [ "$enable_watchdog" = "true" ]; then
        log_success "Enabled auto-start for watchdog script."
        curl -s "http://$shelly_ip/rpc/Script.Start?id=$script_id" > /dev/null
        log_success "Watchdog script started."
    else
        log_info "Auto-start disabled. Script.Start skipped."
    fi
}

# Ask the Shelly to GET the RPi healthcheck endpoint. Returns 0 if the
# Shelly receives an HTTP 200, non-zero on any error.
shelly_can_reach_rpi() {
    local rpi_ip="$1"
    local payload response code

    payload=$(jq -n \
        --arg url "http://${rpi_ip}:8080" \
        '{id: 1, method: "HTTP.GET", params: {url: $url, timeout: 5}}')

    response=$(curl -s -X POST -H "Content-Type: application/json" \
        -d "$payload" \
        "http://$shelly_ip/rpc")
    log_debug "Shelly HTTP.GET preflight response: $response"

    code=$(echo "$response" | jq -r '.result.code // .error.code // empty')
    if [ "$code" = "200" ]; then
        return 0
    fi

    log_debug "Preflight error code=$code response=$response"
    return 1
}

            ############# Full Script Execution #############

# Full script runs all phases
run_full_script() {
    check_prerequisites

    download_hashball

    install_certs
    set_imager_repo_url
    get_shelly_ip
    check_and_update_firmware
    upload_certs_to_shelly
    configure_shelly_mqtt
    configure_shelly_name
    deploy_watchdog_script
    reboot_shelly_device
}

# Refresh just the Shelly watchdog script on an already-provisioned device.
# Skips cert/EEPROM/MQTT setup and the post-deploy reboot.
run_watchdog_only() {
    local pkgs=("jq" "curl")
    local missing=()
    for pkg in "${pkgs[@]}"; do
        dpkg -l | grep -qw "$pkg" || missing+=("$pkg")
    done
    if [ ${#missing[@]} -gt 0 ]; then
        log_error "Missing packages: ${missing[*]}"
        exit 1
    fi

    get_shelly_ip
    deploy_watchdog_script
}

usage() {
    cat <<EOF
Usage: $(basename "$0") [--watchdog-only]

  (no flags)         Run full provisioning.
  --watchdog-only    Redeploy the Shelly RPi watchdog script only.
                     Use to refresh the watchdog on already-provisioned units.
EOF
}

            ############# Start Script Execution #############

case "${1:-}" in
    "")
        run_full_script
        ;;
    --watchdog-only)
        run_watchdog_only
        ;;
    -h|--help)
        usage
        ;;
    *)
        log_error "Unknown option: $1"
        usage
        exit 1
        ;;
esac

# End of script
