// PartitionOperations.cpp — Concrete operation implementations.
//
// Each operation follows the pattern:
//   1. Validate parameters
//   2. Lock/dismount if needed
//   3. Read partition table
//   4. Modify as needed
//   5. Write back
//   6. Notify kernel
//
// DISCLAIMER: This code is for authorized disk utility software only.

#include "PartitionOperations.h"

#include <windows.h>
#include <winioctl.h>

#include <algorithm>
#include <cstring>
#include <sstream>

#include <QProcess>
#include <QRegularExpression>
#include <QString>

namespace spw
{

// ============================================================================
// OperationUtils — Shared utility functions
// ============================================================================

namespace OperationUtils
{

Result<std::unique_ptr<PartitionTable>> readPartitionTable(DiskId diskId, uint32_t sectorSize)
{
    auto diskResult = RawDiskHandle::open(diskId, DiskAccessMode::ReadOnly);
    if (!diskResult)
        return diskResult.error();

    auto& disk = diskResult.value();

    auto geomResult = disk.getGeometry();
    if (!geomResult)
        return geomResult.error();

    uint64_t diskSizeBytes = geomResult.value().totalBytes;

    // Create a read callback that reads from the disk handle
    DiskReadCallback readFunc = [&disk, sectorSize](uint64_t offset, uint32_t size)
        -> Result<std::vector<uint8_t>>
    {
        SectorOffset lba = offset / sectorSize;
        SectorCount count = (size + sectorSize - 1) / sectorSize;
        auto readResult = disk.readSectors(lba, count, sectorSize);
        if (!readResult)
            return readResult.error();

        auto data = std::move(readResult.value());
        // Trim to requested size
        if (data.size() > size)
            data.resize(size);
        return data;
    };

    return PartitionTable::parse(readFunc, diskSizeBytes, sectorSize);
}

Result<void> writePartitionTable(DiskId diskId, const PartitionTable& table, uint32_t sectorSize)
{
    auto serializeResult = table.serialize();
    if (!serializeResult)
        return serializeResult.error();

    const auto& tableData = serializeResult.value();
    if (tableData.empty())
        return ErrorInfo::fromCode(ErrorCode::PartitionTableCorrupt, "Serialized table is empty");

    auto diskResult = RawDiskHandle::open(diskId, DiskAccessMode::ReadWrite);
    if (!diskResult)
        return diskResult.error();

    auto& disk = diskResult.value();

    // Write the serialized data starting at LBA 0
    SectorCount sectorsToWrite = (tableData.size() + sectorSize - 1) / sectorSize;

    // Pad to sector boundary if needed
    std::vector<uint8_t> padded = tableData;
    size_t paddedSize = static_cast<size_t>(sectorsToWrite) * sectorSize;
    if (padded.size() < paddedSize)
        padded.resize(paddedSize, 0);

    auto writeResult = disk.writeSectors(0, padded.data(), sectorsToWrite, sectorSize);
    if (!writeResult)
        return writeResult;

    // For GPT, we also need to write the backup at the end of the disk.
    // The serialize() method should include both primary and backup for GPT.
    // If it only includes primary, we'd need a separate call here.
    // The current PartitionTable interface handles this in serialize().

    return disk.flushBuffers();
}

Result<void> notifyKernel(DiskId diskId)
{
    auto diskResult = RawDiskHandle::open(diskId, DiskAccessMode::ReadWrite);
    if (!diskResult)
        return diskResult.error();

    DWORD bytesReturned = 0;
    BOOL ok = DeviceIoControl(
        diskResult.value().nativeHandle(),
        IOCTL_DISK_UPDATE_PROPERTIES,
        nullptr, 0,
        nullptr, 0,
        &bytesReturned,
        nullptr);

    if (!ok)
    {
        return ErrorInfo::fromWin32(ErrorCode::DiskWriteError, GetLastError(),
            "IOCTL_DISK_UPDATE_PROPERTIES failed");
    }

    return Result<void>::ok();
}

Result<VolumeHandle> lockAndDismountVolume(wchar_t driveLetter)
{
    if (driveLetter == 0)
        return ErrorInfo::fromCode(ErrorCode::InvalidArgument, "No drive letter specified");

    auto volResult = VolumeHandle::openByLetter(driveLetter, DiskAccessMode::ReadWrite);
    if (!volResult)
        return volResult.error();

    auto lockResult = volResult.value().lock();
    if (!lockResult)
        return lockResult.error();

    auto dismountResult = volResult.value().dismount();
    if (!dismountResult)
    {
        volResult.value().unlock();
        return dismountResult.error();
    }

    return std::move(volResult);
}

QString formatSize(uint64_t bytes)
{
    const double KB = 1024.0;
    const double MB = KB * 1024.0;
    const double GB = MB * 1024.0;
    const double TB = GB * 1024.0;

    if (bytes >= static_cast<uint64_t>(TB))
        return QString("%1 TB").arg(static_cast<double>(bytes) / TB, 0, 'f', 2);
    if (bytes >= static_cast<uint64_t>(GB))
        return QString("%1 GB").arg(static_cast<double>(bytes) / GB, 0, 'f', 2);
    if (bytes >= static_cast<uint64_t>(MB))
        return QString("%1 MB").arg(static_cast<double>(bytes) / MB, 0, 'f', 1);
    if (bytes >= static_cast<uint64_t>(KB))
        return QString("%1 KB").arg(static_cast<double>(bytes) / KB, 0, 'f', 0);
    return QString("%1 bytes").arg(bytes);
}

} // namespace OperationUtils

// ============================================================================
// CreatePartitionOp
// ============================================================================

CreatePartitionOp::CreatePartitionOp(const Params& params)
    : m_params(params)
{
    m_targetDiskId = params.diskId;
}

QString CreatePartitionOp::description() const
{
    uint64_t sizeBytes = m_params.sectorCount * m_params.sectorSize;
    return QString("Create %1 partition on disk %2 at LBA %3")
        .arg(OperationUtils::formatSize(sizeBytes))
        .arg(m_params.diskId)
        .arg(m_params.startLba);
}

Result<void> CreatePartitionOp::execute(ProgressCallback progress)
{
    if (progress) progress(0, "Reading partition table...");

    // Read existing partition table
    auto tableResult = OperationUtils::readPartitionTable(m_params.diskId, m_params.sectorSize);
    if (!tableResult)
        return tableResult.error();

    auto& table = tableResult.value();

    if (progress) progress(20, "Adding partition entry...");

    // Build partition params
    PartitionParams partParams;
    partParams.startLba = m_params.startLba;
    partParams.sectorCount = m_params.sectorCount;
    partParams.mbrType = m_params.mbrType;
    partParams.isActive = m_params.isActive;
    partParams.isLogical = m_params.isLogical;
    partParams.typeGuid = m_params.typeGuid;
    partParams.gptName = m_params.gptName;

    // If GPT type GUID is zero, default to Microsoft Basic Data
    if (table->type() == PartitionTableType::GPT && m_params.typeGuid.isZero())
    {
        partParams.typeGuid = GptTypes::microsoftBasicData();
    }

    auto addResult = table->addPartition(partParams);
    if (!addResult)
        return addResult;

    if (progress) progress(40, "Writing partition table...");

    // Write updated partition table
    auto writeResult = OperationUtils::writePartitionTable(m_params.diskId, *table, m_params.sectorSize);
    if (!writeResult)
        return writeResult;

    if (progress) progress(60, "Notifying kernel...");

    // Notify kernel of changes
    OperationUtils::notifyKernel(m_params.diskId);

    // Find the index of the newly created partition
    auto partitions = table->partitions();
    for (const auto& entry : partitions)
    {
        if (entry.startLba == m_params.startLba && entry.sectorCount == m_params.sectorCount)
        {
            m_createdIndex = entry.index;
            break;
        }
    }

    // Optionally format the new partition
    if (m_params.formatAfter)
    {
        if (progress) progress(65, "Formatting new partition...");

        FormatEngine engine;
        FormatTarget target;
        target.diskIndex = m_params.diskId;
        target.partitionOffsetBytes = m_params.startLba * m_params.sectorSize;
        target.partitionSizeBytes = m_params.sectorCount * m_params.sectorSize;
        target.sectorSize = m_params.sectorSize;

        auto formatProgress = [&progress](int pct, const QString& status)
        {
            if (progress)
            {
                // Map format progress (0-100) to our range (65-95)
                int mapped = 65 + (pct * 30) / 100;
                progress(mapped, status);
            }
        };

        auto formatResult = engine.format(target, m_params.formatOptions, formatProgress);
        if (!formatResult)
            return formatResult;
    }

    if (progress) progress(100, "Partition created successfully");
    return Result<void>::ok();
}

Result<void> CreatePartitionOp::undo()
{
    if (m_createdIndex < 0)
        return ErrorInfo::fromCode(ErrorCode::InvalidArgument, "No partition to undo");

    auto tableResult = OperationUtils::readPartitionTable(m_params.diskId, m_params.sectorSize);
    if (!tableResult)
        return tableResult.error();

    auto& table = tableResult.value();
    auto deleteResult = table->deletePartition(m_createdIndex);
    if (!deleteResult)
        return deleteResult;

    auto writeResult = OperationUtils::writePartitionTable(m_params.diskId, *table, m_params.sectorSize);
    if (!writeResult)
        return writeResult;

    OperationUtils::notifyKernel(m_params.diskId);
    m_createdIndex = -1;

    return Result<void>::ok();
}

// ============================================================================
// DeletePartitionOp
// ============================================================================

DeletePartitionOp::DeletePartitionOp(const Params& params)
    : m_params(params)
{
    m_targetDiskId = params.diskId;
    m_targetPartitionId = params.partitionIndex;
}

QString DeletePartitionOp::description() const
{
    return QString("Delete partition %1 on disk %2")
        .arg(m_params.partitionIndex)
        .arg(m_params.diskId);
}

Result<void> DeletePartitionOp::execute(ProgressCallback progress)
{
    if (progress) progress(0, "Preparing to delete partition...");

    // Lock and dismount if the partition is mounted
    std::unique_ptr<VolumeHandle> volHandle;
    if (m_params.driveLetter != 0)
    {
        if (progress) progress(5, "Locking and dismounting volume...");
        auto lockResult = OperationUtils::lockAndDismountVolume(m_params.driveLetter);
        if (!lockResult)
            return lockResult.error();
        volHandle = std::make_unique<VolumeHandle>(std::move(lockResult.value()));
    }

    if (progress) progress(15, "Reading partition table...");

    // Read partition table and save the entry for undo
    auto tableResult = OperationUtils::readPartitionTable(m_params.diskId, m_params.sectorSize);
    if (!tableResult)
        return tableResult.error();

    auto& table = tableResult.value();
    auto partitions = table->partitions();

    // Find and save the partition entry
    for (const auto& entry : partitions)
    {
        if (entry.index == m_params.partitionIndex)
        {
            m_savedEntry = entry;
            break;
        }
    }

    if (!m_savedEntry.has_value())
    {
        return ErrorInfo::fromCode(ErrorCode::PartitionNotFound,
            "Partition index " + std::to_string(m_params.partitionIndex) + " not found");
    }

    // Optionally wipe first sectors to prevent filesystem auto-detection
    if (m_params.wipeFirstSectors)
    {
        if (progress) progress(25, "Wiping filesystem signatures...");

        auto diskResult = RawDiskHandle::open(m_params.diskId, DiskAccessMode::ReadWrite);
        if (diskResult.isOk())
        {
            // Zero first 4KB of the partition
            constexpr uint32_t wipeSize = 4096;
            std::vector<uint8_t> zeros(wipeSize, 0);
            SectorOffset partStartLba = m_savedEntry->startLba;
            SectorCount wipeSecCount = wipeSize / m_params.sectorSize;
            if (wipeSecCount == 0) wipeSecCount = 1;

            // Best-effort: don't fail the whole operation if wipe fails
            diskResult.value().writeSectors(partStartLba, zeros.data(), wipeSecCount, m_params.sectorSize);
            diskResult.value().flushBuffers();
        }
    }

    if (progress) progress(50, "Deleting partition entry...");

    auto deleteResult = table->deletePartition(m_params.partitionIndex);
    if (!deleteResult)
        return deleteResult;

    if (progress) progress(70, "Writing partition table...");

    auto writeResult = OperationUtils::writePartitionTable(m_params.diskId, *table, m_params.sectorSize);
    if (!writeResult)
        return writeResult;

    if (progress) progress(90, "Notifying kernel...");

    OperationUtils::notifyKernel(m_params.diskId);

    // Release volume lock
    if (volHandle)
    {
        volHandle->unlock();
        volHandle->close();
    }

    if (progress) progress(100, "Partition deleted successfully");
    return Result<void>::ok();
}

Result<void> DeletePartitionOp::undo()
{
    if (!m_savedEntry.has_value())
        return ErrorInfo::fromCode(ErrorCode::InvalidArgument, "No saved partition entry for undo");

    auto tableResult = OperationUtils::readPartitionTable(m_params.diskId, m_params.sectorSize);
    if (!tableResult)
        return tableResult.error();

    auto& table = tableResult.value();

    PartitionParams params;
    params.startLba = m_savedEntry->startLba;
    params.sectorCount = m_savedEntry->sectorCount;
    params.mbrType = m_savedEntry->mbrType;
    params.isActive = m_savedEntry->isActive;
    params.isLogical = m_savedEntry->isLogical;
    params.typeGuid = m_savedEntry->typeGuid;
    params.gptName = m_savedEntry->gptName;

    auto addResult = table->addPartition(params);
    if (!addResult)
        return addResult;

    auto writeResult = OperationUtils::writePartitionTable(m_params.diskId, *table, m_params.sectorSize);
    if (!writeResult)
        return writeResult;

    OperationUtils::notifyKernel(m_params.diskId);
    return Result<void>::ok();
}

// ============================================================================
// ResizePartitionOp
// ============================================================================

ResizePartitionOp::ResizePartitionOp(const Params& params)
    : m_params(params)
{
    m_targetDiskId = params.diskId;
    m_targetPartitionId = params.partitionIndex;
}

QString ResizePartitionOp::description() const
{
    uint64_t newSize = m_params.newSectorCount * m_params.sectorSize;
    return QString("Resize partition %1 on disk %2 to %3")
        .arg(m_params.partitionIndex)
        .arg(m_params.diskId)
        .arg(OperationUtils::formatSize(newSize));
}

Result<void> ResizePartitionOp::execute(ProgressCallback progress)
{
    if (progress) progress(0, "Preparing to resize partition...");

    // Lock and dismount if mounted
    std::unique_ptr<VolumeHandle> volHandle;
    if (m_params.driveLetter != 0)
    {
        if (progress) progress(5, "Locking and dismounting volume...");
        auto lockResult = OperationUtils::lockAndDismountVolume(m_params.driveLetter);
        if (!lockResult)
            return lockResult.error();
        volHandle = std::make_unique<VolumeHandle>(std::move(lockResult.value()));
    }

    if (progress) progress(15, "Reading partition table...");

    auto tableResult = OperationUtils::readPartitionTable(m_params.diskId, m_params.sectorSize);
    if (!tableResult)
        return tableResult.error();

    auto& table = tableResult.value();
    auto partitions = table->partitions();

    // Find and save current entry
    for (const auto& entry : partitions)
    {
        if (entry.index == m_params.partitionIndex)
        {
            m_savedEntry = entry;
            break;
        }
    }

    if (!m_savedEntry.has_value())
    {
        return ErrorInfo::fromCode(ErrorCode::PartitionNotFound,
            "Partition index " + std::to_string(m_params.partitionIndex) + " not found");
    }

    // Validate: not making it too small
    if (m_params.newSectorCount == 0)
    {
        return ErrorInfo::fromCode(ErrorCode::InvalidArgument, "New sector count cannot be zero");
    }

    if (progress) progress(30, "Updating partition entry...");

    // For NTFS: If shrinking, we should first shrink the filesystem.
    // For growing, we update the partition entry first, then extend the filesystem.
    // Since filesystem resize is complex and typically handled by the filesystem driver,
    // we do the partition table update. The user should run "extend volume" or
    // equivalent after growing.

    bool isShrinking = m_params.newSectorCount < m_savedEntry->sectorCount;

    if (isShrinking && m_savedEntry->detectedFs == FilesystemType::NTFS && m_params.driveLetter != 0)
    {
        // For NTFS shrink, Windows can handle it via FSCTL_SHRINK_VOLUME
        // but this requires the volume to be mounted. Since we dismounted,
        // we just update the partition table and warn that data might be lost.
        // A full implementation would use FSCTL_SHRINK_VOLUME before dismounting.
    }

    auto resizeResult = table->resizePartition(
        m_params.partitionIndex, m_params.newStartLba, m_params.newSectorCount);
    if (!resizeResult)
        return resizeResult;

    if (progress) progress(60, "Writing partition table...");

    auto writeResult = OperationUtils::writePartitionTable(m_params.diskId, *table, m_params.sectorSize);
    if (!writeResult)
        return writeResult;

    if (progress) progress(80, "Notifying kernel...");

    OperationUtils::notifyKernel(m_params.diskId);

    // Release volume lock
    if (volHandle)
    {
        volHandle->unlock();
        volHandle->close();
    }

    if (progress) progress(100, "Partition resized successfully");
    return Result<void>::ok();
}

Result<void> ResizePartitionOp::undo()
{
    if (!m_savedEntry.has_value())
        return ErrorInfo::fromCode(ErrorCode::InvalidArgument, "No saved partition entry for undo");

    auto tableResult = OperationUtils::readPartitionTable(m_params.diskId, m_params.sectorSize);
    if (!tableResult)
        return tableResult.error();

    auto& table = tableResult.value();

    auto resizeResult = table->resizePartition(
        m_params.partitionIndex, m_savedEntry->startLba, m_savedEntry->sectorCount);
    if (!resizeResult)
        return resizeResult;

    auto writeResult = OperationUtils::writePartitionTable(m_params.diskId, *table, m_params.sectorSize);
    if (!writeResult)
        return writeResult;

    OperationUtils::notifyKernel(m_params.diskId);
    return Result<void>::ok();
}

// ============================================================================
// FormatPartitionOp
// ============================================================================

FormatPartitionOp::FormatPartitionOp(const Params& params)
    : m_params(params)
{
    m_targetDiskId = params.diskId;
    m_targetPartitionId = params.partitionIndex;
}

QString FormatPartitionOp::description() const
{
    QString fsName;
    switch (m_params.options.targetFs)
    {
    case FilesystemType::NTFS:   fsName = "NTFS"; break;
    case FilesystemType::FAT32:  fsName = "FAT32"; break;
    case FilesystemType::FAT16:  fsName = "FAT16"; break;
    case FilesystemType::FAT12:  fsName = "FAT12"; break;
    case FilesystemType::ExFAT:  fsName = "exFAT"; break;
    case FilesystemType::ReFS:   fsName = "ReFS"; break;
    case FilesystemType::Ext2:   fsName = "ext2"; break;
    case FilesystemType::Ext3:   fsName = "ext3"; break;
    case FilesystemType::Ext4:   fsName = "ext4"; break;
    case FilesystemType::SWAP_LINUX: fsName = "Linux swap"; break;
    default: fsName = "Unknown"; break;
    }

    if (m_params.target.hasDriveLetter())
    {
        return QString("Format %1: as %2")
            .arg(QChar(m_params.target.driveLetter))
            .arg(fsName);
    }
    return QString("Format partition %1 on disk %2 as %3")
        .arg(m_params.partitionIndex)
        .arg(m_params.diskId)
        .arg(fsName);
}

Result<void> FormatPartitionOp::execute(ProgressCallback progress)
{
    FormatEngine engine;

    auto formatProgress = [&progress](int pct, const QString& status)
    {
        if (progress) progress(pct, status);
    };

    return engine.format(m_params.target, m_params.options, formatProgress);
}

// ============================================================================
// SetLabelOp
// ============================================================================

SetLabelOp::SetLabelOp(const Params& params)
    : m_params(params)
{
    m_targetDiskId = params.diskId;
    m_targetPartitionId = params.partitionIndex;
}

QString SetLabelOp::description() const
{
    if (m_params.driveLetter != 0)
    {
        return QString("Set label of %1: to \"%2\"")
            .arg(QChar(m_params.driveLetter))
            .arg(QString::fromStdString(m_params.newLabel));
    }
    return QString("Set label of partition %1 to \"%2\"")
        .arg(m_params.partitionIndex)
        .arg(QString::fromStdString(m_params.newLabel));
}

Result<void> SetLabelOp::execute(ProgressCallback progress)
{
    if (progress) progress(0, "Setting volume label...");

    if (m_params.driveLetter != 0)
    {
        // Windows API: SetVolumeLabelW for NTFS/FAT/exFAT
        // First, read the current label for undo
        if (progress) progress(10, "Reading current label...");

        auto fsInfo = VolumeHandle::getFilesystemInfo(m_params.driveLetter);
        if (fsInfo.isOk())
        {
            // Convert wstring to std::string (ASCII-safe for labels)
            const auto& wLabel = fsInfo.value().volumeLabel;
            m_oldLabel.clear();
            for (wchar_t ch : wLabel)
            {
                if (ch < 128)
                    m_oldLabel.push_back(static_cast<char>(ch));
            }
            m_oldLabelSaved = true;
        }

        if (progress) progress(30, "Applying new label...");

        // Build root path: "X:\"
        wchar_t rootPath[] = {m_params.driveLetter, L':', L'\\', L'\0'};

        // Convert label to wide string
        std::wstring wNewLabel;
        for (char ch : m_params.newLabel)
            wNewLabel.push_back(static_cast<wchar_t>(ch));

        BOOL ok = SetVolumeLabelW(rootPath, wNewLabel.c_str());
        if (!ok)
        {
            return ErrorInfo::fromWin32(ErrorCode::FormatFailed, GetLastError(),
                "SetVolumeLabelW failed");
        }

        if (progress) progress(100, "Label set successfully");
        return Result<void>::ok();
    }
    else if (m_params.diskId >= 0 &&
             (m_params.fsType == FilesystemType::Ext2 ||
              m_params.fsType == FilesystemType::Ext3 ||
              m_params.fsType == FilesystemType::Ext4))
    {
        // Direct superblock write for ext filesystems
        if (progress) progress(10, "Opening disk...");

        auto diskResult = RawDiskHandle::open(m_params.diskId, DiskAccessMode::ReadWrite);
        if (!diskResult)
            return diskResult.error();

        auto& disk = diskResult.value();
        uint64_t sbOffset = m_params.partitionOffsetBytes + 1024; // Superblock at byte 1024
        uint32_t sectorSize = m_params.sectorSize;

        // Read current superblock
        if (progress) progress(20, "Reading superblock...");

        SectorOffset sbLba = sbOffset / sectorSize;
        SectorCount sbSectors = (1024 + sectorSize - 1) / sectorSize;
        auto sbRead = disk.readSectors(sbLba, sbSectors, sectorSize);
        if (!sbRead)
            return sbRead.error();

        auto sbData = std::move(sbRead.value());
        if (sbData.size() < 1024)
            return ErrorInfo::fromCode(ErrorCode::DiskReadError, "Superblock read too short");

        // Verify magic
        uint16_t magic;
        uint32_t sbStartInBuf = static_cast<uint32_t>(sbOffset % sectorSize);
        std::memcpy(&magic, sbData.data() + sbStartInBuf + 0x38, 2);
        if (magic != EXT_SUPER_MAGIC)
        {
            return ErrorInfo::fromCode(ErrorCode::FilesystemCorrupt,
                "Not a valid ext superblock (bad magic)");
        }

        // Save old label
        char oldLabel[17] = {};
        std::memcpy(oldLabel, sbData.data() + sbStartInBuf + 0x78, 16);
        m_oldLabel = oldLabel;
        m_oldLabelSaved = true;

        if (progress) progress(50, "Writing new label...");

        // Write new label (16 bytes, zero-padded)
        std::memset(sbData.data() + sbStartInBuf + 0x78, 0, 16);
        size_t labelLen = std::min<size_t>(m_params.newLabel.size(), 16);
        std::memcpy(sbData.data() + sbStartInBuf + 0x78, m_params.newLabel.data(), labelLen);

        // Update write time
        uint32_t now = static_cast<uint32_t>(
            std::chrono::duration_cast<std::chrono::seconds>(
                std::chrono::system_clock::now().time_since_epoch()
            ).count());
        std::memcpy(sbData.data() + sbStartInBuf + 0x30, &now, 4);

        auto writeResult = disk.writeSectors(sbLba, sbData.data(), sbSectors, sectorSize);
        if (!writeResult)
            return writeResult;

        disk.flushBuffers();

        if (progress) progress(100, "Label set successfully");
        return Result<void>::ok();
    }

    return ErrorInfo::fromCode(ErrorCode::InvalidArgument,
        "Cannot set label: no drive letter and not an ext filesystem");
}

Result<void> SetLabelOp::undo()
{
    if (!m_oldLabelSaved)
        return ErrorInfo::fromCode(ErrorCode::InvalidArgument, "No saved label for undo");

    // Reuse execute logic with the old label
    Params undoParams = m_params;
    undoParams.newLabel = m_oldLabel;
    SetLabelOp undoOp(undoParams);
    return undoOp.execute(nullptr);
}

// ============================================================================
// SetFlagsOp
// ============================================================================

SetFlagsOp::SetFlagsOp(const Params& params)
    : m_params(params)
{
    m_targetDiskId = params.diskId;
    m_targetPartitionId = params.partitionIndex;
}

QString SetFlagsOp::description() const
{
    QStringList changes;
    if (m_params.setActive.has_value())
    {
        changes << QString("Set %1")
            .arg(m_params.setActive.value() ? "active" : "inactive");
    }
    if (m_params.gptAttributes.has_value())
    {
        changes << QString("Set GPT attributes 0x%1")
            .arg(m_params.gptAttributes.value(), 16, 16, QChar('0'));
    }

    return QString("Set flags on partition %1 disk %2: %3")
        .arg(m_params.partitionIndex)
        .arg(m_params.diskId)
        .arg(changes.join(", "));
}

Result<void> SetFlagsOp::execute(ProgressCallback progress)
{
    if (progress) progress(0, "Reading partition table...");

    auto tableResult = OperationUtils::readPartitionTable(m_params.diskId, m_params.sectorSize);
    if (!tableResult)
        return tableResult.error();

    auto& table = tableResult.value();
    auto partitions = table->partitions();

    // Save current entry for undo
    for (const auto& entry : partitions)
    {
        if (entry.index == m_params.partitionIndex)
        {
            m_savedEntry = entry;
            break;
        }
    }

    if (!m_savedEntry.has_value())
    {
        return ErrorInfo::fromCode(ErrorCode::PartitionNotFound,
            "Partition index " + std::to_string(m_params.partitionIndex) + " not found");
    }

    if (progress) progress(30, "Modifying flags...");

    if (table->type() == PartitionTableType::MBR && m_params.setActive.has_value())
    {
        // For MBR: use the MbrPartitionTable-specific method
        auto* mbrTable = dynamic_cast<MbrPartitionTable*>(table.get());
        if (mbrTable)
        {
            if (m_params.setActive.value())
            {
                auto setResult = mbrTable->setActivePartition(m_params.partitionIndex);
                if (!setResult)
                    return setResult;
            }
            else
            {
                // Clear active flag: set to -1 (none)
                auto setResult = mbrTable->setActivePartition(-1);
                if (!setResult)
                    return setResult;
            }
        }
    }

    // For GPT attributes: we need to modify the partition entry directly.
    // The current PartitionTable interface doesn't expose attribute modification,
    // so we serialize, modify, and re-parse. In practice, the UI layer would
    // use a more direct API. For now, we do a full table rewrite.
    // GPT attribute modification would require extending the PartitionTable interface
    // with a setAttributes(index, attributes) method.

    if (progress) progress(60, "Writing partition table...");

    auto writeResult = OperationUtils::writePartitionTable(m_params.diskId, *table, m_params.sectorSize);
    if (!writeResult)
        return writeResult;

    if (progress) progress(90, "Notifying kernel...");
    OperationUtils::notifyKernel(m_params.diskId);

    if (progress) progress(100, "Flags set successfully");
    return Result<void>::ok();
}

Result<void> SetFlagsOp::undo()
{
    if (!m_savedEntry.has_value())
        return ErrorInfo::fromCode(ErrorCode::InvalidArgument, "No saved entry for undo");

    // Restore previous active flag
    Params undoParams = m_params;
    if (m_params.setActive.has_value())
    {
        undoParams.setActive = m_savedEntry->isActive;
    }
    if (m_params.gptAttributes.has_value())
    {
        undoParams.gptAttributes = m_savedEntry->gptAttributes;
    }

    SetFlagsOp undoOp(undoParams);
    return undoOp.execute(nullptr);
}

// ============================================================================
// CheckFilesystemOp
// ============================================================================

CheckFilesystemOp::CheckFilesystemOp(const Params& params)
    : m_params(params)
{
    m_targetDiskId = params.diskId;
    m_targetPartitionId = params.partitionIndex;
}

QString CheckFilesystemOp::description() const
{
    QString mode = m_params.repair ? "Check and repair" : "Check";
    if (m_params.driveLetter != 0)
    {
        return QString("%1 filesystem on %2:").arg(mode).arg(QChar(m_params.driveLetter));
    }
    return QString("%1 filesystem on partition %2 disk %3")
        .arg(mode).arg(m_params.partitionIndex).arg(m_params.diskId);
}

Result<void> CheckFilesystemOp::execute(ProgressCallback progress)
{
    if (progress) progress(0, "Starting filesystem check...");

    if (m_params.driveLetter != 0)
    {
        // Use chkdsk for Windows filesystems (NTFS, FAT, exFAT)
        QStringList args;
        args << QString("%1:").arg(QChar(m_params.driveLetter));

        if (m_params.repair)
            args << "/F";

        if (m_params.badSectorScan)
            args << "/R";

        // /Y = suppress confirmation for /F
        if (m_params.repair)
            args << "/Y";

        if (progress) progress(5, "Running chkdsk...");

        QProcess chkdskProcess;
        chkdskProcess.setProgram("chkdsk.exe");
        chkdskProcess.setArguments(args);
        chkdskProcess.start();

        if (!chkdskProcess.waitForStarted(10000))
        {
            return ErrorInfo::fromCode(ErrorCode::FilesystemCheckFailed,
                "Failed to start chkdsk: " + chkdskProcess.errorString().toStdString());
        }

        // Monitor output for progress
        while (chkdskProcess.state() != QProcess::NotRunning)
        {
            chkdskProcess.waitForReadyRead(500);

            QByteArray output = chkdskProcess.readAllStandardOutput();
            if (!output.isEmpty() && progress)
            {
                QString text = QString::fromLocal8Bit(output);
                // chkdsk outputs "XX percent complete" lines
                QRegularExpression pctRx("(\\d+)\\s+percent\\s+complete");
                auto match = pctRx.match(text);
                if (match.hasMatch())
                {
                    int pct = match.captured(1).toInt();
                    int scaled = 5 + (pct * 90) / 100;
                    progress(scaled, QString("Checking... %1%").arg(pct));
                }
            }
        }

        chkdskProcess.waitForFinished(600000); // 10 minute timeout

        int exitCode = chkdskProcess.exitCode();
        // chkdsk exit codes:
        // 0 = no errors found
        // 1 = errors found and fixed
        // 2 = disk cleanup performed
        // 3 = could not check, needs /F
        if (exitCode > 2)
        {
            QByteArray errOutput = chkdskProcess.readAllStandardError();
            QByteArray stdOutput = chkdskProcess.readAllStandardOutput();
            return ErrorInfo::fromCode(ErrorCode::FilesystemCheckFailed,
                "chkdsk exited with code " + std::to_string(exitCode) + ": " +
                stdOutput.toStdString() + errOutput.toStdString());
        }

        if (progress) progress(100, exitCode == 0 ? "No errors found" : "Errors found and fixed");
        return Result<void>::ok();
    }
    else if (m_params.fsType == FilesystemType::Ext2 ||
             m_params.fsType == FilesystemType::Ext3 ||
             m_params.fsType == FilesystemType::Ext4)
    {
        // For ext filesystems, we can do a basic superblock check
        if (progress) progress(10, "Reading ext superblock...");

        auto diskResult = RawDiskHandle::open(m_params.diskId, DiskAccessMode::ReadOnly);
        if (!diskResult)
            return diskResult.error();

        auto& disk = diskResult.value();
        uint64_t sbOffset = m_params.partitionOffsetBytes + 1024;
        uint32_t sectorSize = m_params.sectorSize;

        SectorOffset sbLba = sbOffset / sectorSize;
        SectorCount sbSectors = (1024 + sectorSize - 1) / sectorSize;
        auto sbRead = disk.readSectors(sbLba, sbSectors, sectorSize);
        if (!sbRead)
            return sbRead.error();

        auto sbData = std::move(sbRead.value());
        uint32_t sbStartInBuf = static_cast<uint32_t>(sbOffset % sectorSize);

        constexpr size_t kExt4SuperblockMinSize = 1024; // ext2/3/4 superblock is 1024 bytes
        if (sbData.size() < sbStartInBuf + kExt4SuperblockMinSize)
        {
            return ErrorInfo::fromCode(ErrorCode::FilesystemCheckFailed,
                "Superblock read too short");
        }

        // Verify magic
        uint16_t magic;
        std::memcpy(&magic, sbData.data() + sbStartInBuf + 0x38, 2);
        if (magic != EXT_SUPER_MAGIC)
        {
            return ErrorInfo::fromCode(ErrorCode::FilesystemCorrupt,
                "Invalid ext superblock magic (expected 0xEF53)");
        }

        if (progress) progress(40, "Checking superblock fields...");

        // Check state field
        uint16_t state;
        std::memcpy(&state, sbData.data() + sbStartInBuf + 0x3A, 2);

        // Check error count
        uint32_t errorCount;
        std::memcpy(&errorCount, sbData.data() + sbStartInBuf + 0x194, 4);

        // Check mount count vs max mount count
        uint16_t mntCount, maxMntCount;
        std::memcpy(&mntCount, sbData.data() + sbStartInBuf + 0x34, 2);
        std::memcpy(&maxMntCount, sbData.data() + sbStartInBuf + 0x36, 2);

        if (progress) progress(80, "Analyzing results...");

        std::string statusMsg;
        bool hasIssues = false;

        if (state != 1) // Not clean
        {
            statusMsg += "Filesystem was not cleanly unmounted. ";
            hasIssues = true;
        }

        if (errorCount > 0)
        {
            statusMsg += "Filesystem has " + std::to_string(errorCount) + " recorded error(s). ";
            hasIssues = true;
        }

        if (maxMntCount != static_cast<uint16_t>(-1) && mntCount >= maxMntCount)
        {
            statusMsg += "Mount count exceeded maximum — fsck recommended. ";
            hasIssues = true;
        }

        if (m_params.repair && hasIssues && state != 1)
        {
            // Basic repair: mark filesystem as clean
            // Real repair would require e2fsck logic (much more complex)
            if (progress) progress(85, "Marking filesystem as clean...");

            auto diskWrite = RawDiskHandle::open(m_params.diskId, DiskAccessMode::ReadWrite);
            if (diskWrite.isOk())
            {
                uint16_t cleanState = 1;
                std::memcpy(sbData.data() + sbStartInBuf + 0x3A, &cleanState, 2);

                // Reset error count
                uint32_t zeroErrors = 0;
                std::memcpy(sbData.data() + sbStartInBuf + 0x194, &zeroErrors, 4);

                diskWrite.value().writeSectors(sbLba, sbData.data(), sbSectors, sectorSize);
                diskWrite.value().flushBuffers();
                statusMsg += "State reset to clean. ";
            }
        }

        if (!hasIssues)
            statusMsg = "No issues detected in superblock";

        if (progress) progress(100, QString::fromStdString(statusMsg));
        return Result<void>::ok();
    }

    return ErrorInfo::fromCode(ErrorCode::FilesystemNotSupported,
        "Filesystem check not supported for this filesystem type without a drive letter");
}

} // namespace spw