#!/bin/bash

get_disk_priority() {
    local disk=$1
    local disk_name=$(basename $disk)
    
    if [[ $disk_name == nvme* ]]; then
        echo 1
        return
    fi
    
    local rotational=$(cat /sys/block/$disk/queue/rotational 2>/dev/null)
    if [[ $rotational -eq 0 ]]; then
        echo 2
        return
    fi
    
    local transport=$(lsblk -d -o NAME,TRAN /dev/$disk 2>/dev/null | grep $disk | awk '{print $2}')
    if [[ $transport == "sas" ]]; then
        echo 3
    else
        echo 4  # SATA OR OTHER
    fi
}

get_disk_size() {
    local disk=$1
    cat /sys/block/$disk/size 2>/dev/null
}

is_installed_node() {
    local disks=$(lsblk -d -n -o NAME,TYPE | grep disk | awk '{print $1}')
    
    for disk in $disks; do
        local partitions=$(lsblk -n -o NAME /dev/$disk | grep -E "^${disk}p?[0-9]+" || true)
        
        for part in $partitions; do
            local mount_point="/mnt/tmp_${part}"
            mkdir -p $mount_point 2>/dev/null
            
            if mount -t ext4 -o ro /dev/$part $mount_point 2>/dev/null || \
               mount -t xfs -o ro /dev/$part $mount_point 2>/dev/null; then
                
                if [[ -f $mount_point/.sunhpc-release ]]; then
                    umount $mount_point
                    rmdir $mount_point 2>/dev/null
                    echo "yes"
                    return 0
                fi
                umount $mount_point
            fi
            rmdir $mount_point 2>/dev/null
        done
    done
    
    echo "no"
}

select_best_disk() {
    local disks=()
    
    for disk in /sys/block/*; do
        local disk_name=$(basename $disk)
        if [[ $disk_name =~ ^(loop|zram|ram|sr|fd|dm-) ]]; then
            continue
        fi
        if [[ -e /dev/$disk_name ]] && [[ $(cat /sys/block/$disk_name/removable 2>/dev/null) -eq 0 ]]; then
            disks+=($disk_name)
        fi
    done
    
    if [[ ${#disks[@]} -eq 0 ]]; then
        echo ""
        return 1
    fi
    
    declare -A disk_priority
    declare -A disk_size
    
    for disk in "${disks[@]}"; do
        disk_priority[$disk]=$(get_disk_priority $disk)
        disk_size[$disk]=$(get_disk_size $disk)
    done
    
    local sorted_disks=($(for disk in "${disks[@]}"; do
        echo "${disk_priority[$disk]}:${disk_size[$disk]}:$disk"
    done | sort -n -t: -k1,1 -k2,2rn | cut -d: -f3))
    
    echo "${sorted_disks[0]}"
}

# 智能计算 Root 分区大小
# 参数: $1 = 磁盘总大小(GB), $2 = 期望的默认大小(GB)
# 返回: 合适的 root 分区大小(GB)
calculate_root_size() {
    local disk_total_gb=$1
    local default_root_gb=${2:-70}
    local min_root_gb=30
    local max_root_gb=100

    # 如果磁盘总容量小于最小要求（70GB + boot + swap = 约75GB）
    if [[ $disk_total_gb -lt 75 ]]; then
        # 磁盘太小，使用最小 root 分区
        echo $min_root_gb
    elif [[ $disk_total_gb -lt 100 ]]; then
        # 磁盘在 75-100GB 之间，使用 40-50GB
        echo $((disk_total_gb - 45))  # 预留 boot+swap+state 空间
    elif [[ $disk_total_gb -lt 150 ]]; then
        # 磁盘在 100-150GB 之间，使用 60GB
        echo 60
    elif [[ $disk_total_gb -lt 250 ]]; then
        # 磁盘在 150-250GB 之间，使用默认 70GB
        echo $default_root_gb
    else
        # 大磁盘，限制最大 100GB
        echo $max_root_gb
    fi
}

calculate_state_size() {
    local disk_total_sectors=$1
    local root_size_sectors=$2
    echo $((disk_total_sectors - root_size_sectors))
}

get_disk_size_gb() {
    local disk=$1
    local total_sectors=$(cat /sys/block/$disk/size 2>/dev/null)
    local total_gb=$((total_sectors * 512 / 1024 / 1024 / 1024))
    echo $total_gb
}

get_disk_type() {
    local disk=$1

    if [[ $disk == nvme* ]]; then
        echo "NVMe"
    elif [[ $(cat /sys/block/$disk/queue/rotational 2>/dev/null) -eq 0 ]]; then
        echo "SSD"
    else
        local transport=$(lsblk -d -o NAME,TRAN /dev/$disk 2>/dev/null | grep $disk | awk '{print $2}')
        if [[ $transport == "sas" ]]; then
            echo "SAS"
        else
            echo "SATA"
        fi
    fi
}

generate_ks_partition() {
    local disk=$1
    local is_upgrade=$2

    # 获取磁盘总容量（扇区和GB）
    local total_sectors=$(cat /sys/block/$disk/size)
    local total_gb=$(get_disk_size_gb $disk)
    local disk_type=$(get_disk_type $disk)
    
    # 智能计算 root 分区大小
    local root_gb=$(calculate_root_size $total_gb 70)
    local root_size_sectors=$((root_gb * 1024 * 1024 * 1024 / 512))
    
    # 计算剩余空间给 state 分区
    local remaining_sectors=$((total_sectors - root_size_sectors))
    local remaining_gb=$((remaining_sectors * 512 / 1024 / 1024 / 1024))
    
    # 计算 boot 和 swap 预留空间
    local boot_gb=1
    local swap_gb=4
    local reserved_gb=$((boot_gb + swap_gb))

    cat << EOF
# ==================================================
# Disk Partition Information
# ==================================================
# Installation Type : $([ "$is_upgrade" == "yes" ] && echo "Upgrade" || echo "Fresh") Installation
# Disk Device       : /dev/$disk
# Disk Type         : $disk_type
# Disk Total        : $total_gb GB ($total_sectors sectors)
# 
# Partition Plan:
# - /boot           : ${boot_gb}GB (reserved)
# - swap            : ${swap_gb}GB (reserved)
# - /               : ${root_gb}GB ($root_size_sectors sectors)
# - /state          : ${remaining_gb}GB ($remaining_sectors sectors)
# ==================================================
EOF

    cat << EOF
zerombr
ignoredisk --only-use=$disk
clearpart --all --initlabel
part /boot --fstype="xfs" --size=1024 --ondisk=$disk
part swap --fstype="swap" --size=4096 --ondisk=$disk
part / --fstype="xfs" --size=$((root_gb * 1024)) --ondisk=$disk
EOF
    
    if [[ $is_upgrade == "yes" ]]; then
        cat << EOF
part /state/partition1 --fstype=xfs --size=1 --grow --ondisk=$disk
EOF
    fi
}

main() {
    local is_installed=$(is_installed_node)
    
    if [[ $is_installed == "yes" ]]; then
        echo "# Detected: Already installed node (found .sunhpc-release)"
        local found_disk=""
        local disks=$(lsblk -d -n -o NAME,TYPE | grep disk | awk '{print $1}')
        
        for disk in $disks; do
            local partitions=$(lsblk -n -o NAME /dev/$disk | grep -E "^${disk}p?[0-9]+" || true)
            for part in $partitions; do
                local mount_point="/mnt/tmp_${part}"
                mkdir -p $mount_point 2>/dev/null
                if mount -t ext4 -o ro /dev/$part $mount_point 2>/dev/null || \
                   mount -t xfs -o ro /dev/$part $mount_point 2>/dev/null; then
                    if [[ -f $mount_point/.sunhpc-release ]]; then
                        found_disk=$disk
                        umount $mount_point
                        rmdir $mount_point 2>/dev/null
                        break 2
                    fi
                    umount $mount_point
                fi
                rmdir $mount_point 2>/dev/null
            done
        done
        
        if [[ -n $found_disk ]]; then
            echo "# Found existing installation on disk: /dev/$found_disk"
            generate_ks_partition $found_disk "yes"
        else
            echo "# ERROR: Cannot find disk with .sunhpc-release, using default"
            local best_disk=$(select_best_disk)
            generate_ks_partition $best_disk "yes"
        fi
    else
        echo "# Detected: Fresh installation"
        local best_disk=$(select_best_disk)
        
        if [[ -z $best_disk ]]; then
            echo "# ERROR: No disk found"
            exit 1
        fi
        
        echo "# Selected best disk: /dev/$best_disk"
        generate_ks_partition $best_disk "no"
    fi
}

main