#!/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")
    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"
}

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

# Full script runs all phases
run_full_script() {
    check_prerequisites

    download_hashball

    install_certs
    get_shelly_ip
    check_and_update_firmware
    upload_certs_to_shelly
    configure_shelly_mqtt
    configure_shelly_name
    reboot_shelly_device
}

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

# Always run the full script
run_full_script

# End of script
