#include <GUIConstantsEx.au3>
#include <WinAPI.au3>
#include <WinAPIFiles.au3>
#include <MsgBoxConstants.au3>
#include <FileConstants.au3>
#include <ButtonConstants.au3>
#include <StaticConstants.au3>
#include <WindowsConstants.au3>

; Add COM error handler
Local $oMyError = ObjEvent("AutoIt.Error", "ErrFunc")

Global Enum $DISKMBR_ERR_NONE, _
    $DISKMBR_ERR_SOURCE, _
    $DISKMBR_ERR_OPEN, _
    $DISKMBR_ERR_POINTER, _
    $DISKMBR_ERR_STRUCT, _
    $DISKMBR_ERR_READ, _
    $DISKMBR_ERR_WRITE, _
    $DISKMBR_ERR_LOCK

Global $nTotalSectors = 0, $nBytesPerSector = 0
Global $lblProgress, $lblProgressPercent, $lblETA
Global $sSaveLocation = ""

; Create GUI
$hMainGUI = GUICreate("Drive Imaging Tool", 800, 600)
$lstDrives = GUICtrlCreateListView("Index|Model|Size|Path|InterfaceType|Partitions|TotalSectors|BytesPerSector", 10, 10, 780, 300)
$btnCreateImage = GUICtrlCreateButton("Create Image", 10, 320, 200, 40)
$btnSelectSaveLocation = GUICtrlCreateButton("Select Save Location", 220, 320, 200, 40)
$lblSaveLocation = GUICtrlCreateLabel("Save Location: Not Selected", 10, 370, 780, 20)
$btnRestoreImage = GUICtrlCreateButton("Restore Image", 10, 400, 200, 40)
$lblProgress = GUICtrlCreateLabel("Progress: 0 / 0 sectors", 10, 450, 780, 20)
$lblProgressPercent = GUICtrlCreateLabel("Progress: 0%", 10, 480, 780, 20)
$lblETA = GUICtrlCreateLabel("ETA: --:--:--", 10, 510, 780, 20)
GUISetState(@SW_SHOW)

; Populate drive list
Local $oWMI = ObjGet("winmgmts:root\CIMV2")
If Not IsObj($oWMI) Then
    MsgBox($MB_ICONERROR, "Error", "Failed to create WMI object.")
    Exit
EndIf

Local $oDisks = $oWMI.ExecQuery("SELECT * FROM Win32_DiskDrive")
If Not IsObj($oDisks) Then
    MsgBox($MB_ICONERROR, "Error", "Failed to retrieve disk information.")
    Exit
EndIf

For $oDisk In $oDisks
    GUICtrlCreateListViewItem($oDisk.Index & "|" & $oDisk.Model & "|" & Round($oDisk.Size / 1024^3, 2) & " GB|" & "\\.\PhysicalDrive" & $oDisk.Index & "|" & $oDisk.InterfaceType & "|" & $oDisk.Partitions & "|" & $oDisk.TotalSectors & "|" & $oDisk.BytesPerSector, $lstDrives)
Next

Func _GetSelectedDrive()
    Local $sSelected = GUICtrlRead(GUICtrlRead($lstDrives))
    If $sSelected = "" Then Return ""
    Local $aSelected = StringSplit($sSelected, "|")
    $nTotalSectors = $aSelected[7]
    $nBytesPerSector = $aSelected[8]
    Return $aSelected[4]
EndFunc

Func _LockVolume($hDrive)
    Local $aResult = DllCall("kernel32.dll", "bool", "DeviceIoControl", _
        "handle", $hDrive, _
        "dword", 0x00090018, _ ; FSCTL_LOCK_VOLUME
        "ptr", 0, _
        "dword", 0, _
        "ptr", 0, _
        "dword", 0, _
        "dword*", 0, _
        "ptr", 0)

    If @error Or Not $aResult[0] Then
        ConsoleWrite("Failed to lock volume" & @CRLF)
        ConsoleWrite("WinAPI Last Error = " & _WinAPI_GetLastError() & ": " & _WinAPI_GetLastErrorMessage() & @CRLF)
        Return SetError($DISKMBR_ERR_LOCK, @error, False)
    EndIf

    ConsoleWrite("Successfully locked volume" & @CRLF)
    Return True
EndFunc

Func _ReadAllSectors($sDrive, $nTotalSectors, $nBytesPerSector, $sOutputFile)
    Local $errRet = $DISKMBR_ERR_NONE
    Local $extRet = 0
    Local $myRet = False
    Const $BUFFER_SIZE = 1024 * 1024 ; 1 MB buffer

    ConsoleWrite("Starting _ReadAllSectors for " & $sDrive & @CRLF)
    ConsoleWrite("Total Sectors: " & $nTotalSectors & ", Bytes Per Sector: " & $nBytesPerSector & @CRLF)

    ; Open the output file
    Local $hFile = FileOpen($sOutputFile, $FO_BINARY + $FO_OVERWRITE)
    If $hFile == -1 Then
        ConsoleWrite("Error: Unable to create output file." & @CRLF)
        Return SetError($DISKMBR_ERR_OPEN, @error, $myRet)
    EndIf

    ; Open the physical drive
    Local $hDrive = _WinAPI_CreateFileEx($sDrive, $OPEN_EXISTING, $GENERIC_READ, $FILE_SHARE_READ, $FILE_FLAG_NO_BUFFERING + $FILE_FLAG_RANDOM_ACCESS)
    If $hDrive == -1 Then
        $extRet = @error
        $errRet = $DISKMBR_ERR_OPEN
        ConsoleWrite("Unable to access " & $sDrive & " [@error = " & $extRet & "]" & @CRLF)
        ConsoleWrite("WinAPI Last Error = " & _WinAPI_GetLastError() & ": " & _WinAPI_GetLastErrorMessage() & @CRLF)
        FileClose($hFile)
        Return SetError($errRet, $extRet, $myRet)
    EndIf

    ConsoleWrite("Successfully opened drive handle" & @CRLF)

    ; Create a buffer for reading
    Local $tBuffer = DllStructCreate("byte[" & $BUFFER_SIZE & "]")
    If @error Then
        ConsoleWrite("Error creating buffer structure: " & @error & @CRLF)
        _WinAPI_CloseHandle($hDrive)
        FileClose($hFile)
        Return SetError($DISKMBR_ERR_STRUCT, @error, $myRet)
    EndIf
    Local $pBuffer = DllStructGetPtr($tBuffer)

    ; Variables for progress tracking
    Local $iCurrentSector = 0
    Local $iLastPercentage = 0
    Local $iStartTime = TimerInit()

    ; Read all sectors and write to file
    While $iCurrentSector < $nTotalSectors
        Local $iSectorsToRead = ($nTotalSectors - $iCurrentSector) > ($BUFFER_SIZE / $nBytesPerSector) ? ($BUFFER_SIZE / $nBytesPerSector) : ($nTotalSectors - $iCurrentSector)
        Local $iBytesToRead = $iSectorsToRead * $nBytesPerSector

        _WinAPI_SetFilePointerEx($hDrive, $iCurrentSector * $nBytesPerSector, $FILE_BEGIN)
        If @error Then
            ConsoleWrite("Error setting file pointer for sector " & $iCurrentSector & ": " & @error & @CRLF)
            ConsoleWrite("WinAPI Last Error: " & _WinAPI_GetLastError() & @CRLF)
        EndIf

        Local $aResult = DllCall("kernel32.dll", "bool", "ReadFile", "handle", $hDrive, "ptr", $pBuffer, "dword", $iBytesToRead, "dword*", 0, "ptr", 0)
        If @error Or Not $aResult[0] Then
            ConsoleWrite("Failed to read sector " & $iCurrentSector & @CRLF)
            ConsoleWrite("Error: " & @error & @CRLF)
            ConsoleWrite("WinAPI Last Error: " & _WinAPI_GetLastError() & @CRLF)
            ExitLoop
        EndIf

        ; Write buffer to file
        FileWrite($hFile, DllStructGetData($tBuffer, 1, $iBytesToRead))
        $iCurrentSector += $iSectorsToRead

        ; Update progress
        Local $iPercentage = Floor(($iCurrentSector / $nTotalSectors) * 100)
        If $iPercentage > $iLastPercentage Then
            ConsoleWrite("Progress: " & $iCurrentSector & " / " & $nTotalSectors & " sectors (" & $iPercentage & "%)" & @CRLF)
            GUICtrlSetData($lblProgress, "Progress: " & $iCurrentSector & " / " & $nTotalSectors & " sectors")
            GUICtrlSetData($lblProgressPercent, "Progress: " & $iPercentage & "%")
            Local $iElapsedTime = TimerDiff($iStartTime)
            Local $iETA = Floor($iElapsedTime / $iCurrentSector * ($nTotalSectors - $iCurrentSector) / 1000)
            GUICtrlSetData($lblETA, "ETA: " & _FormatTime($iETA))
            $iLastPercentage = $iPercentage
        EndIf
    WEnd

    ConsoleWrite("Finished reading sectors" & @CRLF)
    FileClose($hFile)
    _WinAPI_CloseHandle($hDrive)
    Return True
EndFunc


Func _FormatTime($iSeconds)
    Local $iHours = Floor($iSeconds / 3600)
    Local $iMinutes = Floor(($iSeconds - $iHours * 3600) / 60)
    Local $iSecs = Floor($iSeconds - $iHours * 3600 - $iMinutes * 60)
    Return StringFormat("%02d:%02d:%02d", $iHours, $iMinutes, $iSecs)
EndFunc


Func _RestoreFromImage($sDrive, $sImageFile)
    Local $errRet = $DISKMBR_ERR_NONE
    Local $extRet = 0
    Local $myRet = False
    Const $BUFFER_SIZE = 1024 * 1024 ; 1 MB buffer

    ConsoleWrite("Starting _RestoreFromImage for " & $sDrive & @CRLF)

    ; Open the image file
    Local $hFile = FileOpen($sImageFile, $FO_BINARY)
    If $hFile == -1 Then
        ConsoleWrite("Error: Unable to open image file." & @CRLF)
        Return SetError($DISKMBR_ERR_OPEN, @error, $myRet)
    EndIf

    ; Open the physical drive
    Local $hDrive = _WinAPI_CreateFileEx($sDrive, $OPEN_EXISTING, $GENERIC_WRITE, $FILE_SHARE_WRITE, $FILE_FLAG_NO_BUFFERING + $FILE_FLAG_RANDOM_ACCESS)
    If $hDrive == -1 Then
        $extRet = @error
        $errRet = $DISKMBR_ERR_OPEN
        ConsoleWrite("Unable to access " & $sDrive & " [@error = " & $extRet & "]" & @CRLF)
        ConsoleWrite("WinAPI Last Error = " & _WinAPI_GetLastError() & ": " & _WinAPI_GetLastErrorMessage() & @CRLF)
        FileClose($hFile)
        Return SetError($errRet, $extRet, $myRet)
    EndIf

    ConsoleWrite("Successfully opened drive handle" & @CRLF)

    ; Lock the volume
    If Not _LockVolume($hDrive) Then
        FileClose($hFile)
        _WinAPI_CloseHandle($hDrive)
        Return SetError($DISKMBR_ERR_LOCK, @error, $myRet)
    EndIf

    ; Create a buffer for reading
    Local $tBuffer = DllStructCreate("byte[" & $BUFFER_SIZE & "]")
    If @error Then
        ConsoleWrite("Error creating buffer structure: " & @error & @CRLF)
        _WinAPI_CloseHandle($hDrive)
        FileClose($hFile)
        Return SetError($DISKMBR_ERR_STRUCT, @error, $myRet)
    EndIf
    Local $pBuffer = DllStructGetPtr($tBuffer)

    ; Variables for progress tracking
    Local $iCurrentSector = 0
    Local $iLastPercentage = 0
    Local $iStartTime = TimerInit()

    ; Read all sectors from the image file and write to the drive
    While Not @error
        Local $iSectorsToRead = ($BUFFER_SIZE / $nBytesPerSector)
        Local $iBytesToRead = $iSectorsToRead * $nBytesPerSector
        Local $iReadBytes = FileRead($hFile, $iBytesToRead)

        If @error Or $iReadBytes == "" Then
            ConsoleWrite("Error reading from image file or end of file reached." & @CRLF)
            ExitLoop
        EndIf

        ; Write buffer to the drive
        _WinAPI_SetFilePointerEx($hDrive, $iCurrentSector * $nBytesPerSector, $FILE_BEGIN)
        If @error Then
            ConsoleWrite("Error setting file pointer for sector " & $iCurrentSector & ": " & @error & @CRLF)
            ExitLoop
        EndIf

        Local $aResult = DllCall("kernel32.dll", "bool", "WriteFile", "handle", $hDrive, "ptr", $pBuffer, "dword", $iBytesToRead, "dword*", 0, "ptr", 0)
        If @error Or Not $aResult[0] Then
            ConsoleWrite("Failed to write sector " & $iCurrentSector & @CRLF)
            ConsoleWrite("Error: " & @error & @CRLF)
            ConsoleWrite("WinAPI Last Error: " & _WinAPI_GetLastError() & @CRLF)
            ExitLoop
        EndIf

        $iCurrentSector += $iSectorsToRead

        ; Update progress
        Local $iPercentage = Floor(($iCurrentSector / $nTotalSectors) * 100)
        If $iPercentage > $iLastPercentage Then
            ConsoleWrite("Progress: " & $iCurrentSector & " / " & $nTotalSectors & " sectors (" & $iPercentage & "%)" & @CRLF)
            GUICtrlSetData($lblProgress, "Progress: " & $iCurrentSector & " / " & $nTotalSectors & " sectors")
            GUICtrlSetData($lblProgressPercent, "Progress: " & $iPercentage & "%")
            Local $iElapsedTime = TimerDiff($iStartTime)
            Local $iETA = Floor($iElapsedTime / $iCurrentSector * ($nTotalSectors - $iCurrentSector) / 1000)
            GUICtrlSetData($lblETA, "ETA: " & _FormatTime($iETA))
            $iLastPercentage = $iPercentage
        EndIf
    WEnd

    ConsoleWrite("Finished restoring sectors" & @CRLF)
    FileClose($hFile)
    _WinAPI_CloseHandle($hDrive)
    Return True
EndFunc

Func ErrFunc()
    MsgBox($MB_ICONERROR, "AutoIt Error", "AutoIt Error: " & @error & @CRLF & @extended)
EndFunc

While 1
    Switch GUIGetMsg()
        Case $GUI_EVENT_CLOSE
            Exit
        Case $btnCreateImage
            ; Handle image creation
            Local $sDrive = _GetSelectedDrive()
            If $sDrive <> "" Then
                Local $sFile = FileSaveDialog("Save Image As", "", "Disk Image Files (*.img)", 2)
                If $sFile <> "" Then
                    _ReadAllSectors($sDrive, $nTotalSectors, $nBytesPerSector, $sFile)
                EndIf
            EndIf
        Case $btnRestoreImage
            ; Handle restore from image
            Local $sDrive = _GetSelectedDrive()
            If $sDrive <> "" Then
                Local $sFile = FileOpenDialog("Open Image File", "", "Disk Image Files (*.img)", 1)
                If $sFile <> "" Then
                    _RestoreFromImage($sDrive, $sFile)
                EndIf
            EndIf
        Case $btnSelectSaveLocation
            ; Correct usage of FileSelectFolder with root directory
            Local $sSaveLocation = FileSelectFolder("Select Save Folder", "", 0, "")
            If $sSaveLocation <> "" Then
                GUICtrlSetData($lblSaveLocation, "Save Location: " & $sSaveLocation)
            EndIf
    EndSwitch
WEnd