#!/bin/bash

help(){
  echo -e "\n$0 
  info
  create\n    -vcpu\n    -ram \n    -disk \n    -ssh-keys-dir \n    -add-disk-block\n    -public-ssh-port
  manage\n    -start\n    -shutdown\n    -reboot\n    -kill\n    -erase 
  disk\n    -attach\n    -detach\n    -resize\n"
  exit 1
}

case "$1" in
  "create")
    if [[ $# -lt 2 ]]; then
      echo ""
      echo -e "$0 $1 [VM_NAME]"
      echo -e "default: -vcpu 8 -ram 8 -disk 60 --ssh-keys-dir keys/"
      echo -e "args[GB]:\n  -vcpu\n  -ram \n  -disk \n  -ssh-keys-dir \n  -add-disk-block\n  -public-ssh-port)"
      echo ""
      exit 1
    fi 

    VM_NAME=$2
    VCPU=8
    RAM_GB=8000 #memory is noted as mb in the virt-install program
    DISK_GB=64
    SSH_KEYS_DIR="keys/"
    DISK_BLOCK_GB=0
    PSP=${PSP:-0}

    while [[ $# -gt 0 ]]; do
      case "$1" in
        create)
          shift 2
          ;;
        -vcpu)
          VCPU=$2
          shift 2
          ;;
        -ram)
          RAM_GB=$(($2*1000))
          shift 2
          ;;
        -disk)
          DISK_GB=$2
          shift 2
          ;;
        -ssh-keys-dir)
          SSH_KEYS_DIR=$2
          shift 2
          ;;
        -add-disk-block)
          DISK_BLOCK_GB=$2
          shift 2
          ;;
        -public-ssh-port)
          PSP=$2
          shift 2
          ;;
        *)
          echo "Unknown argument: $1"
          exit 1
          ;;
      esac
    done
    download_url="https://download.rockylinux.org/pub/rocky/9.3/images/x86_64/Rocky-9-GenericCloud-LVM-9.3-20231113.0.x86_64.qcow2"
    dir_path="/var/lib/libvirt/images/"
    src_file="/var/lib/libvirt/images/Rocky9.3.qcow2"
    seed_iso="/var/lib/libvirt/images/seed.iso"
    packages=("nc" "htop" "wireguard-tools" "bind-utils" "tmux" "net-tools" "curl" "mlocate" "dnsmasq" "qemu-kvm" "libvirt" "libvirt-client" "bridge-utils" "virt-install" "virt-manager" "genisoimage")
    NEW_IMG_PATH="/var/lib/libvirt/images/${VM_NAME}.qcow2"
    XML_PATH="/tmp/${VM_NAME}.xml"
    DISK_BLOCK_GB_BASE_PATH="/var/lib/libvirt/images/${VM_NAME}-vm-disks/"
    DISK_BLOCK_GB_PATH="${DISK_BLOCK_GB_BASE_PATH}${VM_NAME}-vda-${DISK_BLOCK_GB}G.qcow2"
    ssh_private=$(cat /home/s22/man1)

    function check_package_installed() {
      local package_name=$1
      rpm -q "$package_name" &> /dev/null
    }

    for package in "${packages[@]}"; do
      if ! check_package_installed "$package"; then
        #echo "Downloading packages..."
        sudo dnf install -y "$package" &> /dev/null
      fi
    done

    [ ! -d $dir_path ] && mkdir -p $dir_path &> /dev/null
    [ ! -f "$src_file" ] && (wget -O "$src_file" "$download_url" &> /dev/null || { echo "Failed to download Rocky"; exit 1; })

    echo -e "#cloud-config\nusers:\n  - name: root\n    ssh-authorized-keys:" > user-data

    for key in "$SSH_KEYS_DIR"/*.pub; do
        echo "      - $(cat "$key")" >> user-data 
    done

    echo "    sudo: ['ALL=(ALL) NOPASSWD:ALL']
    groups: sudo
    shell: /bin/bash" >> user-data 

    if [ "$PSP" -ne 0 ]; then
      echo "write_files:" >> user-data
      echo "  - path: /root/man1" >> user-data
      echo "    content: |" >> user-data
      while read -r line; do
        echo "      $line" >> user-data
      done <<< "${ssh_private}"
    fi

    echo 'runcmd:' >> user-data 
    #echo '  - [ /usr/bin/wget, "http://example.com/file", -O, /tmp/examplefile ]' >> user-data 
    #echo '  - touch /root/test1.txt' >> user-data 
    if [ "$DISK_BLOCK_GB" -ne 0 ]; then
      echo "echo to disksetup in vm"
      echo "vm name 1 ${VM_NAME}"
      echo '  - echo "#!/bin/bash" > /root/disk-setup.sh' >> user-data
     # echo '  - echo "sleep 40;" >> /root/disk-setup.sh' >> user-data
      echo "  - echo \"DISK='/dev/vda'; MOUNT_PATH='/home/${VM_NAME}-vda-${DISK_BLOCK_GB}G';\" >> /root/disk-setup.sh" >> user-data
      echo "  - echo '[ ! -d \$MOUNT_PATH ] && mkdir -p \$MOUNT_PATH;' >> /root/disk-setup.sh" >> user-data
      echo "  - echo 'blkid | grep -q \$DISK || mkfs.ext4 \$DISK;' >> /root/disk-setup.sh" >> user-data
      echo "  - echo 'grep -q \$DISK /etc/fstab || echo \"\$DISK \$MOUNT_PATH ext4 defaults,nofail 0 0\" >> /etc/fstab;' >> /root/disk-setup.sh" >> user-data
      echo '  - echo "mount -a;" >> /root/disk-setup.sh' >> user-data
      echo '  - echo "systemctl daemon-reload" >> /root/disk-setup.sh' >> user-data
      echo '  - chmod +x /root/disk-setup.sh' >> user-data
      echo '  - /root/disk-setup.sh' >> user-data
      echo '  - rm -f /root/disk-setup.sh' >> user-data
    fi   

    if [ "$PSP" -ne 0 ]; then
      echo '  - chmod 600 /root/man1' >> user-data
      #kill ssh and reconnect every 4 hours
      ##echo "  - (echo \"0 */4 * * * PIDS=\\\$(pgrep -f \\\"ssh.*root@64.176.179.97\\\"); if [ -n \\\"\\\${PIDS}\\\" ]; then kill \\\${PIDS}; fi; /usr/bin/ssh -fN -i /root/man1 -R ${PSP}:localhost:22 -o StrictHostKeyChecking=no root@64.176.179.97\") | crontab -" >> user-data
      echo "  - (echo \"* * * * *  /root/initial.sh\") | crontab -" >> user-data
      echo "  - echo 'PIDS=\$(pgrep -f \"ssh.*root@64.176.179.97\"); if [ -z \"\${PIDS}\" ]; then /usr/bin/ssh -fN -i /root/man1 -R ${PSP}:localhost:22 -o StrictHostKeyChecking=no root@64.176.179.97; fi' > /root/initial.sh" >> user-data
      echo "  - chmod +x /root/initial.sh" >> user-data
      echo "  - /root/initial.sh" >> user-data
    fi

    #nofail is present in the fstab which means that boot will continue even if it fails to mount
    echo '  - growpart /dev/sda 4 ' >> user-data #do note that restart is required for the system to register the increased disk size 
    echo '  - sudo lvresize -l +100%FREE /dev/rocky/lvroot' >> user-data 
    echo '  - sudo dnf install -y epel-release dnf-utils' >> user-data
    echo '  - sudo dnf install -y nc xclip tmux htop tar tree wget curl mlocate nano vim unzip net-tools git python3 python3-pip make'>> user-data
    echo '  - touch /root/runcmd_done' >> user-data 

    echo -e "instance-id: iid-ihatecs\nlocal-hostname: cloudimg" > meta-data 

    genisoimage -output "$seed_iso" -volid cidata -joliet -rock user-data meta-data &> /dev/null || { echo "Failed to create seed.iso."; exit 1; }

    cp $src_file $NEW_IMG_PATH &> /dev/null || { echo "Failed to create a new image."; exit 1; }

    if virsh list --all | grep -q "$VM_NAME"; then
      echo -e "\n$VM_NAME already exist. Delete it using \n$0 manage -erase $VM_NAME\nExiting..."
      exit 1
    fi

    sudo virt-install --name $VM_NAME \
      --vcpus $VCPU \
      --ram $RAM_GB \
      --disk path=$NEW_IMG_PATH,size=$DISK_GB,format=qcow2 \
      --disk path=$seed_iso,device=cdrom \
      --os-type linux \
      --os-variant rhl9 \
      --virt-type kvm \
      --graphics none \
      --network bridge=virbr0,model=virtio \
      --print-xml > $XML_PATH || { echo "Failed to print XML."; exit 1; }

    sudo virsh define $XML_PATH &> /dev/null || { echo "Failed to define the new VM."; exit 1; }

    sudo qemu-img resize $NEW_IMG_PATH +$DISK_GB"G" #&> /dev/null

    virsh start $VM_NAME
    echo ""

    if [ $DISK_BLOCK_GB -ne 0 ]; then
      mkdir -p $DISK_BLOCK_GB_BASE_PATH
      qemu-img create -f qcow2 $DISK_BLOCK_GB_PATH "${DISK_BLOCK_GB}G" 
      virsh attach-disk $VM_NAME $DISK_BLOCK_GB_PATH vda --cache none --subdriver qcow2
    fi
    
    message="waiting 29s to begin finding ip address..."
    duration=29
    for ((i=$duration; i>=1; i--)); do
        printf "\r%s%2ds" "$message" $i
        sleep 1
    done
    echo

    while true; do
        OUTPUT=$($0 info | grep "$VM_NAME")
        if [[ "$OUTPUT" != "" ]]; then
            IP_ADDRESS=$(echo "$OUTPUT" | grep -oP '\d+\.\d+\.\d+\.\d+')
            # If IP address is found and is not empty
            if [[ ! -z "$IP_ADDRESS" ]]; then
                echo -e "\nIP address of $VM_NAME is:\n$IP_ADDRESS"
                break
            else
                echo "IP address for $VM_NAME not found. Retrying in 1 second..."
                sleep 1
            fi
        else
            echo "$VM_NAME not found. Exiting..."
            exit 1
        fi
    done
  
  #do take note that df -h will not reflect until reboot

    ;;
  "manage")
    if [[ $# -lt 2 ]]; then
      echo ""
      echo -e "$0 $1 arg [VM_1] [VM_2] [VM_3]  \n$0 $1 arg --all"
      echo -e "args:\n  -start\n  -shutdown\n  -kill\n  -erase "
      echo ""
      exit 1
    fi 
    
    arg=$2
    shift 2

    if [[ "$1" == "--all" ]]; then
      vms=$(virsh list --all --name)  # List all running VMs by name
      set -- $vms                     # Set the positional parameters to the VM names
    fi

    case $arg in
      "-shutdown")
        for vm in "$@"; do
          virsh shutdown "$vm" > /dev/null 2>&1
          printf "%-50s%10s\n" "Shutting down $vm..." $([[ $? -eq 0 ]] && echo "Successful" || echo "Failed")
        done
        ;;
      "-kill")
        for vm in "$@"; do
          virsh destroy "$vm" > /dev/null 2>&1
          printf "%-50s%10s\n" "Destroying $vm..." $([[ $? -eq 0 ]] && echo "Successful" || echo "Failed")
        done
        ;;
      "-erase")
        echo -e "\nWARNING: THIS IS IRREVERSIBLE. Sleeping for 10 seconds. Ctrl-C to stop. WARNING.\n"
        sleep 10
        echo -e "\nErasing image(s) from /var/lib/libvirt/images/\n"
        for vm in "$@"; do
          virsh destroy "$vm" > /dev/null 2>&1
          virsh undefine "$vm" --remove-all-storage > /dev/null 2>&1
          printf "%-50s%10s\n" "Removing $vm..." $([[ $? -eq 0 ]] && echo "Successful" || echo "Failed")
        done
       ;;
      "-start")
        for vm in "$@"; do
          virsh start "$vm" > /dev/null 2>&1
          printf "%-50s%10s\n" "Starting $vm..." $([[ $? -eq 0 ]] && echo "Successful" || echo "Failed")
        done
       ;;
      "-reboot")
        for vm in "$@"; do
          virsh shutdown "$vm" > /dev/null 2>&1
          virsh start "$vm" > /dev/null 2>&1
          printf "%-50s%10s\n" "Rebooting $vm..." $([[ $? -eq 0 ]] && echo "Successful" || echo "Failed")
        done
       ;;
      *)
        echo "Unknown argument: $arg" 
        exit 1
        ;;
    esac
    ;;
  "info")
    printf "%-4s %-10s %-15s %-10s %-10s %-6s %-8s %-12s\n" "Id" "Name" "IP" "State" "Network" "vCPUs" "RAM(GB)" "Disk(GB)"
    printf "%-4s %-10s %-15s %-10s %-10s %-6s %-8s %-12s\n" "----" "----------" "---------------" "----------" "----------" "-----" "-------" "-----------"
    
    vms=$(virsh list --name --all)

    id=1
    for vm in $vms; do
      # Get the MAC address of the VM
      mac=$(virsh dumpxml $vm | grep "mac address" | awk -F\' '{ print $2}')

      # Get the network name
      net=$(virsh dumpxml $vm | grep "<source network" | awk -F\' '{print $2}')

      # Assign "default" if no network name is found
      if [ -z "$net" ]
      then
        net="default"
      fi

      # Get the IP address of the VM
      ip=$(virsh net-dhcp-leases $net | grep $mac | awk '{print $5}' | cut -f1 -d'/')

      # Get the state of the VM
      state=$(virsh domstate $vm)

      # Get the vCPUs, RAM and Disk details of the VM
      vcpus=$(virsh dominfo $vm | grep "CPU(s)" | awk '{print $2}')
      
      ram=$(bc <<< "scale=2; $(virsh dominfo $vm | grep "Max memory" | awk '{print $3}')/1000000")

      # Use du to get the size of the disk file in GB
      disk=$(du -sk /var/lib/libvirt/images/${vm}.qcow2 | awk '{ printf "%.2f", $1/1024/1024 }' )

      printf "%-4s %-10s %-15s %-10s %-10s %-6s %-8s %-12s\n" "$id" "$vm" "$ip" "$state" "$net" "$vcpus" "$ram" "$disk"
      id=$((id+1))
    done
    exit 1
    ;;
  *)
    help
    exit 1
    ;;
esac