#!/usr/bin/env bash
# =============================================================================
# dragon-q8b-thermal - Thermal Governor & Cooling Management for Dragon Q8B
# =============================================================================
# Reproduces the thermal governor configuration functionality of Radxa OS
# (rsetup). Default cooling is fanless (power_allocator). step_wise is used
# only on zones that list a pwm-fan cooling device (Heatsink 6845B overlay).
# =============================================================================
set -Eeuo pipefail

DEFAULT_GOVERNOR="power_allocator"
CONFIG_FILE="/etc/dragon-q8b/thermal.conf"
PACKAGED_TMPFILES="/usr/lib/tmpfiles.d/dragon-q8b-thermal.conf"
SYSFS_THERMAL_DIR="/sys/class/thermal"

usage() {
    cat <<'EOF'
Usage: dragon-q8b-thermal [OPTIONS]

Configure and monitor thermal governor policies and fan cooling on the
Radxa Dragon Q8B (Qualcomm SC8280XP).

Options:
  -m, --mode POLICY        Set thermal governor policy (default: power_allocator)
                           Supported: step_wise, power_allocator, user_space,
                                      fair_share, bang_bang
  -s, --status             Display status of all thermal zones and cooling devices
  -p, --persist            Persist the selected governor across system reboots
      --apply-config       Read /etc/dragon-q8b/thermal.conf and apply policy
  -f, --fan-speed LEVEL    Set cooling device state (0 to max_state; requires user_space)
  -i, --menu               Open interactive thermal configuration menu
  -b, --batch              Run in non-interactive batch mode (default when options given)
  -n, --dry-run            Show actions without writing to sysfs or disk
  -q, --quiet              Suppress informational output
  -h, --help               Show this help message

Examples:
  dragon-q8b-thermal                          # Apply default fanless power_allocator
  dragon-q8b-thermal --status                 # View current temperatures & governors
  dragon-q8b-thermal -m power_allocator       # Switch to passive/fanless cooling mode
  dragon-q8b-thermal -m user_space -f 3       # Switch to manual control and set fan level
  dragon-q8b-thermal --menu                   # Open interactive rsetup-style menu
  dragon-q8b-overlay enable pwm-fan           # Heatsink 6845B overlay, then reboot
EOF
}

log() {
    if [[ "$QUIET" -eq 0 ]]; then
        printf '[dragon-q8b-thermal] %s\n' "$*"
    fi
}

log_err() {
    printf '[dragon-q8b-thermal] ERROR: %s\n' "$*" >&2
}

log_warn() {
    printf '[dragon-q8b-thermal] WARNING: %s\n' "$*" >&2
}

# Defaults
MODE=""
ACTION=""
PERSIST=0
FAN_SPEED=""
DRY_RUN=0
QUIET=0
INTERACTIVE=0

# Parse arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
        -m|--mode|--policy)
            MODE=${2:?missing policy name after $1}
            ACTION="set"
            shift 2
            ;;
        -s|--status)
            ACTION="status"
            shift
            ;;
        -p|--persist)
            PERSIST=1
            shift
            ;;
        --apply-config)
            ACTION="apply_config"
            shift
            ;;
        -f|--fan-speed|--fan-level)
            FAN_SPEED=${2:?missing fan speed level after $1}
            ACTION="fan_speed"
            shift 2
            ;;
        -i|--menu|--interactive)
            INTERACTIVE=1
            ACTION="menu"
            shift
            ;;
        -b|--batch|--non-interactive)
            INTERACTIVE=0
            shift
            ;;
        -n|--dry-run)
            DRY_RUN=1
            shift
            ;;
        -q|--quiet)
            QUIET=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            log_err "Unknown argument: $1"
            usage >&2
            exit 2
            ;;
    esac
done

require_root() {
    if [[ "$DRY_RUN" -eq 1 ]]; then
        return 0
    fi
    if [[ $EUID -ne 0 ]]; then
        log_err "Root privileges required to modify thermal settings. Please run with sudo."
        exit 1
    fi
}

get_thermal_zones() {
    local zones=()
    if [[ -d "$SYSFS_THERMAL_DIR" ]]; then
        for z in "$SYSFS_THERMAL_DIR"/thermal_zone*; do
            [[ -d "$z" ]] && zones+=("$z")
        done
    fi
    printf '%s\n' "${zones[@]:-}"
}

get_cooling_devices() {
    local cdevs=()
    if [[ -d "$SYSFS_THERMAL_DIR" ]]; then
        for c in "$SYSFS_THERMAL_DIR"/cooling_device*; do
            [[ -d "$c" ]] && cdevs+=("$c")
        done
    fi
    printf '%s\n' "${cdevs[@]:-}"
}

pwm_fan_present() {
    local c
    while IFS= read -r c; do
        [[ -n "$c" ]] || continue
        [[ "$(cat "$c/type" 2>/dev/null || true)" == pwm-fan ]] && return 0
    done < <(get_cooling_devices)
    return 1
}

zone_has_pwm_fan() {
    local z=$1 link target base
    for link in "$z"/cdev[0-9] "$z"/cdev[0-9][0-9]; do
        [[ -e "$link" || -L "$link" ]] || continue
        base=$(basename "$link")
        [[ "$base" =~ ^cdev[0-9]+$ ]] || continue
        target=$(readlink -f "$link" 2>/dev/null || echo "$link")
        [[ -n "$target" ]] || continue
        [[ "$(cat "$target/type" 2>/dev/null || true)" == pwm-fan ]] && return 0
    done
    return 1
}

read_conf_value() {
    local key=$1 default=$2 file=${3:-$CONFIG_FILE}
    local value=$default
    if [[ -f "$file" ]]; then
        value=$(awk -F= -v key="$key" '
            $0 ~ "^[[:space:]]*" key "[[:space:]]*=" {
                gsub(/[[:space:]"\x27]/, "", $2)
                print $2
                exit
            }
        ' "$file" || true)
        [[ -n "$value" ]] || value=$default
    fi
    printf '%s\n' "$value"
}

show_status() {
    echo "================================================================================"
    echo "                   Radxa Dragon Q8B Thermal & Fan Status"
    echo "================================================================================"

    local zone_found=0
    printf "%-16s %-22s %-12s %-16s %s\n" "ZONE" "TYPE" "TEMP (°C)" "CURRENT POLICY" "AVAILABLE POLICIES"
    printf "%-16s %-22s %-12s %-16s %s\n" "----------------" "----------------------" "------------" "----------------" "------------------"

    while IFS= read -r z; do
        [[ -n "$z" ]] || continue
        zone_found=1
        local zid type temp_mc temp_c policy avail
        zid=$(basename "$z")
        type=$(cat "$z/type" 2>/dev/null || echo "unknown")
        temp_mc=$(cat "$z/temp" 2>/dev/null || echo "0")
        if [[ "$temp_mc" =~ ^-?[0-9]+$ ]]; then
            temp_c=$(awk "BEGIN {printf \"%.1f\", $temp_mc / 1000}")
        else
            temp_c="N/A"
        fi
        policy=$(cat "$z/policy" 2>/dev/null || echo "unknown")
        avail=$(cat "$z/available_policies" 2>/dev/null || echo "none")

        printf "%-16s %-22s %-12s %-16s %s\n" "$zid" "$type" "$temp_c" "$policy" "$avail"
    done < <(get_thermal_zones)

    if [[ $zone_found -eq 0 ]]; then
        echo "  (No sysfs thermal zones detected in $SYSFS_THERMAL_DIR)"
    fi

    echo
    echo "----------------------------- Cooling Devices ----------------------------------"
    local cdev_found=0
    printf "%-18s %-24s %-16s %s\n" "DEVICE" "TYPE" "CURRENT STATE" "MAX STATE"
    printf "%-18s %-24s %-16s %s\n" "------------------" "------------------------" "----------------" "---------"

    while IFS= read -r c; do
        [[ -n "$c" ]] || continue
        cdev_found=1
        local cid ctype cur max
        cid=$(basename "$c")
        ctype=$(cat "$c/type" 2>/dev/null || echo "unknown")
        cur=$(cat "$c/cur_state" 2>/dev/null || echo "N/A")
        max=$(cat "$c/max_state" 2>/dev/null || echo "N/A")
        printf "%-18s %-24s %-16s %s\n" "$cid" "$ctype" "$cur" "$max"
    done < <(get_cooling_devices)

    if [[ $cdev_found -eq 0 ]]; then
        echo "  (No cooling devices detected in $SYSFS_THERMAL_DIR)"
    fi

    echo
    echo "----------------------------- Saved Configuration ------------------------------"
    if [[ -f "$CONFIG_FILE" ]]; then
        echo "Config file: $CONFIG_FILE"
        grep -v '^#' "$CONFIG_FILE" | grep -v '^[[:space:]]*$' || echo "  (empty)"
    else
        echo "Config file: $CONFIG_FILE (not present, using system defaults)"
    fi
    if [[ -f "$PACKAGED_TMPFILES" ]]; then
        echo "Tmpfiles rule: $PACKAGED_TMPFILES"
        cat "$PACKAGED_TMPFILES"
    fi
    echo "================================================================================"
}

set_governor() {
    local target_policy=$1
    local persist_flag=$2
    local pwm_fan_only=${3:-0}

    require_root

    local updated=0
    local skipped=0
    local zones
    mapfile -t zones < <(get_thermal_zones)

    if [[ ${#zones[@]} -eq 0 ]]; then
        log_warn "No thermal zones found in $SYSFS_THERMAL_DIR."
    fi

    for z in "${zones[@]}"; do
        [[ -n "$z" ]] || continue
        local zid avail cur_policy type
        zid=$(basename "$z")
        type=$(cat "$z/type" 2>/dev/null || echo "unknown")
        avail=$(cat "$z/available_policies" 2>/dev/null || echo "")
        cur_policy=$(cat "$z/policy" 2>/dev/null || echo "")

        if [[ "$pwm_fan_only" -eq 1 ]] && ! zone_has_pwm_fan "$z"; then
            skipped=$((skipped + 1))
            continue
        fi

        if [[ -n "$avail" && " $avail " != *" $target_policy "* ]]; then
            log_warn "$zid ($type) does not support policy '$target_policy' (available: $avail); skipping."
            skipped=$((skipped + 1))
            continue
        fi

        if [[ "$cur_policy" == "$target_policy" ]]; then
            log "Zone $zid ($type) is already set to '$target_policy'."
            updated=$((updated + 1))
            continue
        fi

        if [[ "$DRY_RUN" -eq 1 ]]; then
            log "[DRY-RUN] Would set $zid/policy ($type) -> $target_policy"
            updated=$((updated + 1))
        else
            if echo "$target_policy" > "$z/policy" 2>/dev/null; then
                log "Set $zid ($type) policy to '$target_policy'."
                updated=$((updated + 1))
            else
                log_warn "Failed to set policy for $zid ($type)."
            fi
        fi
    done

    if [[ "$persist_flag" -eq 1 ]]; then
        save_persistence "$target_policy" 0
    fi

    log "Thermal governor update complete ($updated applied/matched, $skipped skipped)."
}

save_persistence() {
    local target_policy=$1
    local auto_pwm=${2:-0}

    require_root

    if [[ "$DRY_RUN" -eq 1 ]]; then
        log "[DRY-RUN] Would save configuration to $CONFIG_FILE (policy: $target_policy, AUTO_PWM_FAN=$auto_pwm)"
        return 0
    fi

    mkdir -p "$(dirname "$CONFIG_FILE")"
    cat > "$CONFIG_FILE" <<EOF
# Configuration for Radxa Dragon Q8B thermal governor
# Generated by dragon-q8b-thermal on $(date -u +'%Y-%m-%d %H:%M:%S UTC')
# Packaged tmpfiles live in $PACKAGED_TMPFILES; this tool does not write /etc/tmpfiles.d.
GOVERNOR=${target_policy}
AUTO_PWM_FAN=${auto_pwm}
EOF
    chmod 0644 "$CONFIG_FILE"
    log "Saved configuration to $CONFIG_FILE"
}

apply_saved_config() {
    local policy auto_pwm
    policy=$(read_conf_value GOVERNOR "$DEFAULT_GOVERNOR")
    auto_pwm=$(read_conf_value AUTO_PWM_FAN 1)

    require_root
    if [[ "$auto_pwm" != 1 ]]; then
        log "Applying explicit thermal governor: $policy"
        set_governor "$policy" 0
        return 0
    fi

    log "Applying default fanless governor '$DEFAULT_GOVERNOR'"
    set_governor "$DEFAULT_GOVERNOR" 0
    if pwm_fan_present; then
        log "pwm-fan cooling device present; using step_wise on zones that list it"
        set_governor step_wise 0 1
    fi
}

set_fan_speed() {
    local speed=$1
    require_root

    local cdevs
    mapfile -t cdevs < <(get_cooling_devices)

    if [[ ${#cdevs[@]} -eq 0 ]]; then
        log_err "No cooling devices found in $SYSFS_THERMAL_DIR."
        exit 1
    fi

    for c in "${cdevs[@]}"; do
        [[ -n "$c" ]] || continue
        local cid max_state
        cid=$(basename "$c")
        max_state=$(cat "$c/max_state" 2>/dev/null || echo "0")

        if (( speed < 0 || speed > max_state )); then
            log_err "Fan speed $speed is out of range for $cid (0 - $max_state)"
            exit 1
        fi

        if [[ "$DRY_RUN" -eq 1 ]]; then
            log "[DRY-RUN] Would set $cid/cur_state -> $speed (max: $max_state)"
        else
            echo "$speed" > "$c/cur_state"
            log "Set $cid cooling state to $speed (max: $max_state)."
        fi
    done
}

interactive_menu() {
    while true; do
        clear
        echo "================================================================================"
        echo "           Radxa Dragon Q8B Thermal Governor Setup (rsetup style)"
        echo "  Default is fanless power_allocator. Heatsink 6845B: overlay enable pwm-fan,"
        echo "  reboot, then step_wise is applied only on zones that list pwm-fan."
        echo "================================================================================"
        echo " 1) step_wise        - Active cooling (Heatsink 6845B: dragon-q8b-overlay enable pwm-fan then reboot)"
        echo " 2) power_allocator  - Passive cooling / fanless or fixed DC fan [Default]"
        echo " 3) user_space       - Manual fan speed & cooling device control"
        echo " 4) fair_share       - Distribute cooling load across devices proportionally"
        echo " 5) bang_bang        - Two-point threshold cooling"
        echo " 6) View Status      - Show thermal zone temperatures and cooling devices"
        echo " 7) Set Fan Speed    - Set manual fan speed state (requires user_space)"
        echo " 8) Exit"
        echo "================================================================================"
        printf "Select an option [1-8]: "
        read -r choice

        case "$choice" in
            1)
                echo
                set_governor "step_wise" 1
                printf "\nPress Enter to continue..."
                read -r _
                ;;
            2)
                echo
                set_governor "power_allocator" 1
                printf "\nPress Enter to continue..."
                read -r _
                ;;
            3)
                echo
                set_governor "user_space" 1
                printf "\nPress Enter to continue..."
                read -r _
                ;;
            4)
                echo
                set_governor "fair_share" 1
                printf "\nPress Enter to continue..."
                read -r _
                ;;
            5)
                echo
                set_governor "bang_bang" 1
                printf "\nPress Enter to continue..."
                read -r _
                ;;
            6)
                echo
                show_status
                printf "\nPress Enter to return to menu..."
                read -r _
                ;;
            7)
                echo
                printf "Enter desired fan speed level (integer): "
                read -r fspeed
                if [[ "$fspeed" =~ ^[0-9]+$ ]]; then
                    set_fan_speed "$fspeed"
                else
                    echo "Invalid number: $fspeed"
                fi
                printf "\nPress Enter to continue..."
                read -r _
                ;;
            8|q|Q)
                echo "Exiting."
                exit 0
                ;;
            *)
                echo "Invalid selection."
                sleep 1
                ;;
        esac
    done
}

# Main Execution Routing
if [[ "$ACTION" == "menu" || ($INTERACTIVE -eq 1 && -t 0 && -t 1 && -z "$ACTION") ]]; then
    interactive_menu
    exit 0
fi

case "$ACTION" in
    status)
        show_status
        ;;
    fan_speed)
        set_fan_speed "$FAN_SPEED"
        ;;
    apply_config)
        apply_saved_config
        ;;
    set)
        set_governor "$MODE" "$PERSIST"
        ;;
    *)
        log "No options specified; applying packaged thermal policy (power_allocator, step_wise if pwm-fan)."
        apply_saved_config
        ;;
esac
