/*
 * Warmboot Extractor
 * Based on Atmosphere fusee_setup_horizon.cpp
 *
 * Copyright (c) 2018-2025 Atmosphère-NX
 */

#include "warmboot_extractor.h"

#include <string.h>
#include <mem/heap.h>

#include <soc/fuse.h>
#include <soc/timer.h>

#include <sec/se.h>

#include "../storage/emummc.h"
#include "../storage/mountmanager.h"
#include <storage/emmc.h>

#include <libs/fatfs/ff.h>
#include <utils/sprintf.h>

// Mariko keyslot for BEK (Boot Encryption Key)
#define KS_MARIKO_BEK 13

// Helper function to check if this is Mariko hardware
bool is_mariko(void)
{
    u32 odm4 = fuse_read_odm(4);

    // Extract hardware type from ODM4 bits (same logic as Atmosphere fuse_api.cpp)
    u32 hw_type = ((odm4 >> 2) & 1) | (((odm4 >> 8) & 1) << 1) | (((odm4 >> 16) & 0xF) << 2);

    // Mariko hardware types: Iowa (0x04), Hoag (0x08), Aula (0x10)
    if (hw_type == 0x04 || hw_type == 0x08 || hw_type == 0x10)
        return true;
    if (hw_type == 0x01) // Icosa
        return false;

    // For 0x02 (Calcio/Copper), fall back to DRAM ID check
    return (fuse_read_dramid(false) >= 4);
}

// Read burnt fuse count (ODM6 + ODM7)
u8 get_burnt_fuses(void)
{
    u8 fuse_count = 0;
    u32 fuse_odm6 = fuse_read_odm(6);
    u32 fuse_odm7 = fuse_read_odm(7);

    // Count bits in ODM6
    for (u32 i = 0; i < 32; i++)
        if ((fuse_odm6 >> i) & 1)
            fuse_count++;

    // Count bits in ODM7
    for (u32 i = 0; i < 32; i++)
        if ((fuse_odm7 >> i) & 1)
            fuse_count++;

    return fuse_count;
}

// Generate warmboot cache path based on fuse count
void get_warmboot_path(char *path, size_t path_size, u8 fuse_count)
{
    (void)path_size;

    // Format matches Atmosphère's convention:
    // Mariko: sd:/warmboot_mariko/wb_xx.bin (lowercase hex)
    // Erista: not used by Atmosphère, but kept for completeness.
    if (is_mariko())
        s_printf(path, "sd:/warmboot_mariko/wb_%02x.bin", fuse_count);
    else
        s_printf(path, "sd:/warmboot_erista/wb_%02x.bin", fuse_count);
}

// Get target firmware version from Package1 header (display only)
static u32 get_target_firmware_from_pkg1(const u8 *package1)
{
    switch (package1[0x1F])
    {
        case 0x01: return 0x100;   // 1.0.0
        case 0x02: return 0x200;   // 2.0.0
        case 0x04: return 0x300;   // 3.0.0
        case 0x07: return 0x400;   // 4.0.0
        case 0x0B: return 0x500;   // 5.0.0
        case 0x0E:
            if (memcmp(package1 + 0x10, "20180802", 8) == 0) return 0x600;  // 6.0.0
            if (memcmp(package1 + 0x10, "20181107", 8) == 0) return 0x620;  // 6.2.0
            break;
        case 0x0F: return 0x700;   // 7.0.0
        case 0x10:
            if (memcmp(package1 + 0x10, "20190314", 8) == 0) return 0x800;  // 8.0.0
            if (memcmp(package1 + 0x10, "20190531", 8) == 0) return 0x810;  // 8.1.0
            if (memcmp(package1 + 0x10, "20190809", 8) == 0) return 0x900;  // 9.0.0
            if (memcmp(package1 + 0x10, "20191021", 8) == 0) return 0x910;  // 9.1.0
            if (memcmp(package1 + 0x10, "20200303", 8) == 0) return 0xA00;  // 10.0.0
            if (memcmp(package1 + 0x10, "20201030", 8) == 0) return 0xB00;  // 11.0.0
            if (memcmp(package1 + 0x10, "20210129", 8) == 0) return 0xC00;  // 12.0.0
            if (memcmp(package1 + 0x10, "20210422", 8) == 0) return 0xC02;  // 12.0.2
            if (memcmp(package1 + 0x10, "20210607", 8) == 0) return 0xC10;  // 12.1.0
            if (memcmp(package1 + 0x10, "20210805", 8) == 0) return 0xD00;  // 13.0.0
            if (memcmp(package1 + 0x10, "20220105", 8) == 0) return 0xD21;  // 13.2.1
            if (memcmp(package1 + 0x10, "20220209", 8) == 0) return 0xE00;  // 14.0.0
            if (memcmp(package1 + 0x10, "20220801", 8) == 0) return 0xF00;  // 15.0.0
            if (memcmp(package1 + 0x10, "20230111", 8) == 0) return 0x1000; // 16.0.0
            if (memcmp(package1 + 0x10, "20230906", 8) == 0) return 0x1100; // 17.0.0
            if (memcmp(package1 + 0x10, "20240207", 8) == 0) return 0x1200; // 18.0.0
            if (memcmp(package1 + 0x10, "20240808", 8) == 0) return 0x1300; // 19.0.0
            if (memcmp(package1 + 0x10, "20250206", 8) == 0) return 0x1400; // 20.0.0
            if (memcmp(package1 + 0x10, "20251009", 8) == 0) return 0x1500; // 21.0.0
            if (memcmp(package1 + 0x10, "20260123", 8) == 0) return 0x1600; // 22.0.0
            break;
        default:
            break;
    }

    return 0;
}

// Error code to string (display/debug)
const char *wb_error_to_string(wb_extract_error_t err)
{
    switch (err)
    {
        case WB_SUCCESS:                  return "Success";
        case WB_ERR_NULL_INFO:            return "NULL wb_info pointer";
        case WB_ERR_ERISTA_NOT_SUPPORTED: return "Erista not supported (uses embedded warmboot)";
        case WB_ERR_MALLOC_PKG1:          return "Failed to allocate Package1 buffer (256KB)";
        case WB_ERR_MMC_INIT:             return "Failed to initialize eMMC";
        case WB_ERR_MMC_PARTITION:        return "Failed to set BOOT0 partition";
        case WB_ERR_MMC_READ:             return "Failed to read Package1 from BOOT0";
        case WB_ERR_DECRYPT_VERIFY:       return "Package1 decryption failed (BEK missing or wrong)";
        case WB_ERR_PK11_MAGIC:           return "PK11 magic not found (invalid Package1)";
        case WB_ERR_WB_SIZE_INVALID:      return "Warmboot size invalid (not 0x800-0x1000)";
        case WB_ERR_MALLOC_WB:            return "Failed to allocate warmboot buffer";
        default:                          return "Unknown error";
    }
}

// Extended extraction with detailed error codes
wb_extract_error_t extract_warmboot_from_pkg1_ex(warmboot_info_t *wb_info)
{
    if (!wb_info)
        return WB_ERR_NULL_INFO;

    // Initialize result
    memset(wb_info, 0, sizeof(warmboot_info_t));
    wb_info->is_erista = !is_mariko();
    wb_info->fuse_count = get_burnt_fuses();

    // Erista doesn't need warmboot extraction from Package1
    if (wb_info->is_erista)
        return WB_ERR_ERISTA_NOT_SUPPORTED;

    // From here on, we're dealing with Mariko only

    // Allocate buffer for Package1
    u8 *pkg1_buffer = (u8 *)malloc(PKG1_SIZE);
    if (!pkg1_buffer)
        return WB_ERR_MALLOC_PKG1;
    u8 *pkg1_buffer_orig = pkg1_buffer;

    // Connect eMMC (BOOT0 access)
    if (connectMMC(MMC_CONN_EMMC))
    {
        free(pkg1_buffer_orig);
        return WB_ERR_MMC_INIT;
    }
    msleep(1);

    // Set to BOOT0 and read package1
    if (!emummc_storage_set_mmc_partition(&emmc_storage, EMMC_BOOT0))
    {
        disconnectMMC();
        free(pkg1_buffer_orig);
        return WB_ERR_MMC_PARTITION;
    }

    if (!emummc_storage_read(&emmc_storage, PKG1_OFFSET / EMMC_BLOCKSIZE, PKG1_SIZE / EMMC_BLOCKSIZE, pkg1_buffer))
    {
        disconnectMMC();
        free(pkg1_buffer_orig);
        return WB_ERR_MMC_READ;
    }

    disconnectMMC();

    // On Mariko, Package1 is encrypted and needs decryption.
    // Skip 0x170 byte header to get to the encrypted payload.
    u8 *pkg1_mariko = pkg1_buffer + 0x170;

    // Set IV from package1 + 0x10 (16 bytes)
    se_aes_iv_set(KS_MARIKO_BEK, pkg1_mariko + 0x10);

    // Decrypt using BEK at keyslot 13
    se_aes_crypt_cbc(KS_MARIKO_BEK, 0, /* DECRYPT */
                     pkg1_mariko + 0x20,
                     0x40000 - (0x20 + 0x170),
                     pkg1_mariko + 0x20,
                     0x40000 - (0x20 + 0x170));

    // Verify decryption (first 0x20 bytes should match decrypted header)
    if (memcmp(pkg1_mariko, pkg1_mariko + 0x20, 0x20) != 0)
    {
        free(pkg1_buffer_orig);
        return WB_ERR_DECRYPT_VERIFY;
    }

    // Use pkg1_mariko as the working pointer
    pkg1_buffer = pkg1_mariko;

    // Detect target firmware (display only)
    u32 target_fw = get_target_firmware_from_pkg1(pkg1_buffer);

    // Store debug info: Package1 version byte and date string
    wb_info->pkg1_version = pkg1_buffer[0x1F];
    memcpy(wb_info->pkg1_date, pkg1_buffer + 0x10, 8);
    wb_info->pkg1_date[8] = '\0';

    // Get burnt fuse count from device
    u8 burnt_fuses = get_burnt_fuses();

    // For warmboot naming, ALWAYS use burnt_fuses directly.
    wb_info->fuse_count = burnt_fuses;
    wb_info->burnt_fuses = burnt_fuses;
    wb_info->target_firmware = target_fw;

    // Determine PK11 offset using firmware hint, then fall back.
    u32 pk11_offset = 0x4000;
    if (target_fw != 0 && target_fw >= 0x620)
        pk11_offset = 0x7000;

    u8 *pk11_ptr_u8 = pkg1_buffer + pk11_offset;
    bool pk11_ok = memcmp(pk11_ptr_u8, "PK11", 4) == 0;

    if (!pk11_ok)
    {
        pk11_offset = (pk11_offset == 0x4000) ? 0x7000 : 0x4000;
        pk11_ptr_u8 = pkg1_buffer + pk11_offset;
        pk11_ok = memcmp(pk11_ptr_u8, "PK11", 4) == 0;
    }

    if (!pk11_ok)
    {
        free(pkg1_buffer_orig);
        return WB_ERR_PK11_MAGIC;
    }

    wb_info->pk11_offset = pk11_offset;

    u32 *pk11_ptr = (u32 *)pk11_ptr_u8;

    // Store debug header
    for (int i = 0; i < 8; i++)
        wb_info->pk11_header[i] = pk11_ptr[i];

    // Navigate through PK11 container to find warmboot (Atmosphere logic)
    u32 *pk11_data = pk11_ptr + (0x20 / sizeof(u32));

    wb_info->sig_found[0] = 0;
    wb_info->sig_found[1] = 0;
    wb_info->sig_found[2] = 0;

    for (int i = 0; i < 3; i++)
    {
        u32 signature = *pk11_data;
        wb_info->sig_found[i] = signature;

        switch (signature)
        {
            case SIG_NX_BOOTLOADER:
                pk11_data += pk11_ptr[6] / sizeof(u32);
                break;
            case SIG_SECURE_MONITOR_1:
            case SIG_SECURE_MONITOR_2:
                pk11_data += pk11_ptr[4] / sizeof(u32);
                break;
            default:
                i = 3;
                break;
        }
    }

    wb_info->debug_ptr_offset = (u32)((u8 *)pk11_data - pk11_ptr_u8);

    u32 wb_size = *pk11_data;

    // Debug layout type
    u32 warmboot_size_header = pk11_ptr[1];
    if (warmboot_size_header == wb_size &&
        warmboot_size_header >= WARMBOOT_MIN_SIZE &&
        warmboot_size_header < WARMBOOT_MAX_SIZE)
        wb_info->debug_layout_type = 1;
    else
        wb_info->debug_layout_type = 2;

    if (wb_size < WARMBOOT_MIN_SIZE || wb_size >= WARMBOOT_MAX_SIZE)
    {
        wb_info->size = wb_size;
        free(pkg1_buffer_orig);
        return WB_ERR_WB_SIZE_INVALID;
    }

    wb_info->size = wb_size;

    memcpy(wb_info->debug_warmboot_preview, (u8 *)pk11_data, 16);

    // Allocate warmboot buffer ([size_u32][warmboot_binary...])
    u8 *wb_data = (u8 *)malloc(wb_size);
    if (!wb_data)
    {
        free(pkg1_buffer_orig);
        return WB_ERR_MALLOC_WB;
    }

    memcpy(wb_data, (u8 *)pk11_data, wb_size);

    wb_info->data = wb_data;
    wb_info->size = wb_size;

    free(pkg1_buffer_orig);
    return WB_SUCCESS;
}

// Original wrapper for backward compatibility
bool extract_warmboot_from_pkg1(warmboot_info_t *wb_info)
{
    return extract_warmboot_from_pkg1_ex(wb_info) == WB_SUCCESS;
}

// Save warmboot to SD card
bool save_warmboot_to_sd(const warmboot_info_t *wb_info, const char *path)
{
    if (!wb_info || !wb_info->data || !path)
        return false;

    FIL fp;
    UINT bytes_written = 0;

    const char *dir_path = wb_info->is_erista ? "sd:/warmboot_erista" : "sd:/warmboot_mariko";
    f_mkdir(dir_path);

    if (f_open(&fp, path, FA_CREATE_ALWAYS | FA_WRITE) != FR_OK)
        return false;

    if (f_write(&fp, wb_info->data, wb_info->size, &bytes_written) != FR_OK)
    {
        f_close(&fp);
        return false;
    }

    f_close(&fp);

    return (bytes_written == wb_info->size);
}