/*
 * Copyright (C) The Android Open Source Project
 * Copyright (C) 2025 oxmc / PawletOS
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include "BootControl.h"

#include <algorithm>
#include <vector>

#include <fcntl.h>
#include <linux/fs.h>
#include <sys/ioctl.h>
#include <unistd.h>

#include <zlib.h>

#include <android-base/logging.h>
#include <android-base/properties.h>
#include <android-base/unique_fd.h>

namespace aidl::android::hardware::boot {

static constexpr size_t kCopyBlockSize = 1u * 1024u * 1024u;  // 1 MiB

// ---------------------------------------------------------------------------
// Constructor
// ---------------------------------------------------------------------------

BootControl::BootControl(const char* misc,
                         const char* boot_a,
                         const char* boot_b)
    : misc_device_(misc),
      boot_a_device_(boot_a),
      boot_b_device_(boot_b) {}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

/* static */
uint32_t BootControl::computeCrc32(const void* data, size_t len) {
    return ::crc32(0, reinterpret_cast<const Bytef*>(data), len);
}

/* static */
bool BootControl::isValidSlot(int32_t slot) {
    return slot == 0 || slot == 1;
}

void BootControl::initDefaultMisc(PawletBootControl* ctrl) {
    memset(ctrl, 0, sizeof(*ctrl));
    ctrl->magic        = PAWLET_MAGIC;
    ctrl->version      = PAWLET_VERSION;
    ctrl->active_slot  = 0;
    ctrl->merge_status = static_cast<uint8_t>(MergeStatus::NONE);
    // Slot A: high priority, fresh attempts, not yet confirmed successful.
    ctrl->slots[0] = {15, static_cast<uint8_t>(PAWLET_MAX_BOOT_ATTEMPTS), 0, 0};
    // Slot B: low priority until activated.
    ctrl->slots[1] = {0,  static_cast<uint8_t>(PAWLET_MAX_BOOT_ATTEMPTS), 0, 0};
}

bool BootControl::readMisc(PawletBootControl* ctrl) {
    android::base::unique_fd fd(
            open(misc_device_.c_str(), O_RDONLY | O_CLOEXEC));
    if (fd.get() < 0) {
        PLOG(ERROR) << "Cannot open misc for reading: " << misc_device_;
        return false;
    }
    if (lseek(fd.get(), PAWLET_MISC_OFFSET, SEEK_SET) < 0) {
        PLOG(ERROR) << "lseek failed on " << misc_device_;
        return false;
    }
    if (read(fd.get(), ctrl, sizeof(*ctrl)) != static_cast<ssize_t>(sizeof(*ctrl))) {
        PLOG(ERROR) << "Short read from " << misc_device_;
        initDefaultMisc(ctrl);
        return true;  // Return safe defaults; writeMisc will persist later.
    }

    if (ctrl->magic != PAWLET_MAGIC || ctrl->version != PAWLET_VERSION) {
        LOG(WARNING) << "misc: invalid magic/version — initialising defaults";
        initDefaultMisc(ctrl);
        return true;
    }

    // Verify CRC32 over all bytes except the crc32 field itself.
    uint32_t stored_crc = ctrl->crc32;
    ctrl->crc32 = 0;
    uint32_t computed = computeCrc32(ctrl, sizeof(*ctrl) - sizeof(uint32_t));
    ctrl->crc32 = stored_crc;

    if (stored_crc != computed) {
        LOG(WARNING) << "misc: CRC mismatch (stored=0x" << std::hex << stored_crc
                     << " computed=0x" << computed << ") — using defaults";
        initDefaultMisc(ctrl);
    }
    return true;
}

bool BootControl::writeMisc(const PawletBootControl& in) {
    PawletBootControl ctrl = in;
    ctrl.crc32 = 0;
    ctrl.crc32 = computeCrc32(&ctrl, sizeof(ctrl) - sizeof(uint32_t));

    android::base::unique_fd fd(
            open(misc_device_.c_str(), O_WRONLY | O_CLOEXEC));
    if (fd.get() < 0) {
        PLOG(ERROR) << "Cannot open misc for writing: " << misc_device_;
        return false;
    }
    if (lseek(fd.get(), PAWLET_MISC_OFFSET, SEEK_SET) < 0) {
        PLOG(ERROR) << "lseek failed on " << misc_device_;
        return false;
    }
    if (write(fd.get(), &ctrl, sizeof(ctrl)) != static_cast<ssize_t>(sizeof(ctrl))) {
        PLOG(ERROR) << "Short write to " << misc_device_;
        return false;
    }
    fsync(fd.get());
    return true;
}

bool BootControl::copyBootPartition(const std::string& src,
                                    const std::string& dst) {
    android::base::unique_fd src_fd(open(src.c_str(), O_RDONLY | O_CLOEXEC));
    if (src_fd.get() < 0) {
        PLOG(ERROR) << "Cannot open source boot partition: " << src;
        return false;
    }
    android::base::unique_fd dst_fd(open(dst.c_str(), O_WRONLY | O_CLOEXEC));
    if (dst_fd.get() < 0) {
        PLOG(ERROR) << "Cannot open dest boot partition: " << dst;
        return false;
    }

    uint64_t size = 0;
    if (ioctl(src_fd.get(), BLKGETSIZE64, &size) < 0) {
        PLOG(ERROR) << "BLKGETSIZE64 failed on " << src;
        return false;
    }

    LOG(INFO) << "Syncing boot partition " << src << " -> " << dst
              << " (" << size << " bytes)";

    std::vector<uint8_t> buf(kCopyBlockSize);
    uint64_t copied = 0;
    while (copied < size) {
        size_t to_read = std::min(static_cast<uint64_t>(buf.size()), size - copied);
        ssize_t n = read(src_fd.get(), buf.data(), to_read);
        if (n <= 0) {
            PLOG(ERROR) << "Read error at offset " << copied << " on " << src;
            return false;
        }
        ssize_t w = write(dst_fd.get(), buf.data(), static_cast<size_t>(n));
        if (w != n) {
            PLOG(ERROR) << "Write error at offset " << copied << " on " << dst;
            return false;
        }
        copied += static_cast<uint64_t>(n);
    }
    fsync(dst_fd.get());
    LOG(INFO) << "Boot partition sync complete (" << copied << " bytes)";
    return true;
}

// ---------------------------------------------------------------------------
// IBootControl implementation
// ---------------------------------------------------------------------------

ndk::ScopedAStatus BootControl::getNumberSlots(int32_t* out) {
    *out = 2;
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::getCurrentSlot(int32_t* out) {
    // The RPi firmware writes the booted slot into the kernel cmdline via
    // androidboot.slot_suffix, which init propagates to ro.boot.slot_suffix.
    const std::string suffix =
            android::base::GetProperty("ro.boot.slot_suffix", "");
    if (suffix == "_b") {
        *out = 1;
    } else {
        // Default to slot A (covers both "_a" and the uninitialised case).
        *out = 0;
    }
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::getActiveBootSlot(int32_t* out) {
    PawletBootControl ctrl{};
    if (!readMisc(&ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }
    *out = static_cast<int32_t>(ctrl.active_slot);
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::getSuffix(int32_t slot, std::string* out) {
    if (!isValidSlot(slot)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
    }
    *out = (slot == 0) ? "_a" : "_b";
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::markBootSuccessful() {
    int32_t current_slot = 0;
    getCurrentSlot(&current_slot);

    PawletBootControl ctrl{};
    if (!readMisc(&ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }

    ctrl.slots[current_slot].successful_boot = 1;
    ctrl.slots[current_slot].tries_remaining =
            static_cast<uint8_t>(PAWLET_MAX_BOOT_ATTEMPTS);
    ctrl.active_slot = static_cast<uint8_t>(current_slot);

    if (!writeMisc(ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }

    // RPi-specific: after a successful slot-B boot, copy boot_b -> boot_a so
    // that subsequent normal reboots (which always start from partition 1 /
    // boot_a) also run the new software.
    if (current_slot == 1) {
        LOG(INFO) << "Slot B verified — syncing boot_b -> boot_a";
        if (!copyBootPartition(boot_b_device_, boot_a_device_)) {
            // Non-fatal: the update is still live; we just can't guarantee
            // persistence after a non-tryboot reboot.
            LOG(ERROR) << "boot_b -> boot_a sync failed; update may be lost on "
                          "next normal reboot";
        }
    }

    LOG(INFO) << "Slot " << current_slot << " marked as successful";
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::setActiveBootSlot(int32_t slot) {
    if (!isValidSlot(slot)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
    }

    PawletBootControl ctrl{};
    if (!readMisc(&ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }

    ctrl.active_slot = static_cast<uint8_t>(slot);

    // Activate the target slot with full priority and attempts.
    ctrl.slots[slot].priority       = 15;
    ctrl.slots[slot].tries_remaining =
            static_cast<uint8_t>(PAWLET_MAX_BOOT_ATTEMPTS);
    ctrl.slots[slot].successful_boot = 0;

    // Demote the other slot so it is still bootable but lower priority.
    int32_t other = slot ^ 1;
    if (ctrl.slots[other].priority > 7) {
        ctrl.slots[other].priority = 7;
    }

    if (!writeMisc(ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }

    LOG(INFO) << "Active boot slot set to " << slot
              << " (suffix: " << (slot == 0 ? "_a" : "_b") << ")";
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::setSlotAsUnbootable(int32_t slot) {
    if (!isValidSlot(slot)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
    }

    PawletBootControl ctrl{};
    if (!readMisc(&ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }

    ctrl.slots[slot].priority        = 0;
    ctrl.slots[slot].tries_remaining = 0;
    ctrl.slots[slot].successful_boot = 0;

    if (!writeMisc(ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }

    LOG(WARNING) << "Slot " << slot << " marked as unbootable";
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::isSlotBootable(int32_t slot, bool* out) {
    if (!isValidSlot(slot)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
    }
    PawletBootControl ctrl{};
    if (!readMisc(&ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }
    *out = (ctrl.slots[slot].priority > 0) &&
           (ctrl.slots[slot].tries_remaining > 0);
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::isSlotMarkedSuccessful(int32_t slot, bool* out) {
    if (!isValidSlot(slot)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
    }
    PawletBootControl ctrl{};
    if (!readMisc(&ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }
    *out = (ctrl.slots[slot].successful_boot == 1);
    return ndk::ScopedAStatus::ok();
}

// RPi does not use Virtual A/B snapshots, but we must implement these
// methods so update_engine can track merge state generically.

ndk::ScopedAStatus BootControl::getSnapshotMergeStatus(MergeStatus* out) {
    PawletBootControl ctrl{};
    if (!readMisc(&ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }
    *out = static_cast<MergeStatus>(ctrl.merge_status);
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus BootControl::setSnapshotMergeStatus(MergeStatus status) {
    PawletBootControl ctrl{};
    if (!readMisc(&ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }
    ctrl.merge_status = static_cast<uint8_t>(status);
    if (!writeMisc(ctrl)) {
        return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
    }
    return ndk::ScopedAStatus::ok();
}

}  // namespace aidl::android::hardware::boot