#include <zephyr/fs/littlefs.h>
#include <zephyr/drivers/flash.h>
#include <zephyr/dfu/flash_img.h>
#include <zephyr/dfu/mcuboot.h>
#include <zephyr/pm/device.h>
#include <fs.h>
LOG_MODULE_REGISTER(buzz_fs, LOG_LEVEL_INF);

#define STORAGE_PARTITION_ID FIXED_PARTITION_ID(littlefs_storage)
#define SLOT1_ID FIXED_PARTITION_ID(slot1_partition)

FS_LITTLEFS_DECLARE_DEFAULT_CONFIG(fs_storage_data);

#define QSPI_FLASH_NODE DT_ALIAS(qspi_flash)
#if !DT_NODE_EXISTS(QSPI_FLASH_NODE)
#error "QSPI Flash alias not defined in devicetree"
#endif

static const struct device *flash_dev = DEVICE_DT_GET(QSPI_FLASH_NODE);

static volatile uint32_t open_count = 0;
static struct k_mutex flash_pm_lock;
static struct slot_info_t slot1_info;
static struct flash_img_context flash_ctx;

static struct fs_mount_t fs_storage_mnt = {
    .type = FS_LITTLEFS,
    .fs_data = &fs_storage_data,
    .storage_dev = (void *)STORAGE_PARTITION_ID,
    .mnt_point = "/lfs",
};

int fs_init(void) {
    int rc = fs_mount(&fs_storage_mnt);
    if (rc < 0) {
        LOG_ERR("Error mounting filesystem: %d", rc);
        return rc;
    }
    k_mutex_init(&flash_pm_lock);

    const struct flash_area *fa;
    rc = flash_area_open(SLOT1_ID, &fa);
    if (rc < 0) {
        LOG_ERR("Error opening flash area for slot 1: %d", rc);
        return rc;
    }
    slot1_info.start_addr = fa->fa_off;
    slot1_info.size = fa->fa_size;
    flash_area_close(fa);

    LOG_DBG("Filesystem mounted successfully");
    return 0;
}

int fs_pm_flash_suspend(void)
{
    if (!device_is_ready(flash_dev)) {
        return -ENODEV;
    }

    k_mutex_lock(&flash_pm_lock, K_FOREVER);

    if (open_count > 0) {
        open_count--;
        if (open_count == 0) {
            int rc = pm_device_action_run(flash_dev, PM_DEVICE_ACTION_SUSPEND);
            if (rc < 0) {
                LOG_WRN("Could not suspend flash: %d", rc);
            } else {
                LOG_DBG("Flash entered deep power-down");
            }
        }
    }

    k_mutex_unlock(&flash_pm_lock);
    return 0;
}

int fs_pm_flash_resume(void)
{
    if (!device_is_ready(flash_dev)) return -ENODEV;

    k_mutex_lock(&flash_pm_lock, K_FOREVER);

    if (open_count == 0) {
        int rc = pm_device_action_run(flash_dev, PM_DEVICE_ACTION_RESUME);
        if (rc == 0) {
            k_busy_wait(50); // t-exit-dpd
            LOG_DBG("Flash resumed");
        }
    }
    open_count++;

    k_mutex_unlock(&flash_pm_lock);
    return 0;
}

int fs_pm_open(struct fs_file_t *file, const char *path, fs_mode_t mode)
{
    LOG_DBG("PM Opening file '%s' with mode 0x%02x", path, mode);
    fs_pm_flash_resume();
    int rc = fs_open(file, path, mode);
    if (rc < 0)
    {
        fs_pm_flash_suspend();
    } 
    return rc;
}

int fs_pm_close(struct fs_file_t *file)
{
    LOG_DBG("PM Closing file");
    int rc = fs_close(file);
    if (rc == 0)
    {
        fs_pm_flash_suspend();
    }
    return rc;
}

int fs_pm_opendir(struct fs_dir_t *dirp, const char *path)
{
    LOG_DBG("PM Opening directory '%s'", path);
    fs_pm_flash_resume();
    int rc = fs_opendir(dirp, path);
    if (rc < 0)
    {
        fs_pm_flash_suspend();
    }
    return rc;
}

int fs_pm_closedir(struct fs_dir_t *dirp)
{
    LOG_DBG("PM Closing directory");
    int rc = fs_closedir(dirp);
    if (rc == 0)
    {
        fs_pm_flash_suspend();
    }
    return rc;
}

int fs_pm_unlink(const char *path)
{
    LOG_DBG("PM Unlinking file '%s'", path);
    fs_pm_flash_resume();
    int rc = fs_unlink(path);
    fs_pm_flash_suspend();
    return rc;
}

int fs_pm_statvfs(const char *path, struct fs_statvfs *stat)
{
    LOG_DBG("PM Getting filesystem stats for '%s'", path);
    fs_pm_flash_resume();
    int rc = fs_statvfs(path, stat);
    fs_pm_flash_suspend();
    return rc;
}

int fs_pm_mkdir(const char *path)
{
    LOG_DBG("PM Creating directory '%s'", path);
    fs_pm_flash_resume();
    int rc = fs_mkdir(path);
    fs_pm_flash_suspend();
    return rc;
}
int flash_get_slot_info(slot_info_t *info) {
    if (slot1_info.size != 0) {
        *info = slot1_info;
        return 0;
    }
    return -ENODEV;
}

int flash_init_firmware_upload(void) {
    int rc;

    rc = flash_img_init_id(&flash_ctx, SLOT1_ID);
    if (rc != 0) {
        printk("Error initializing flash image: %d\n", rc);
        return rc;
    }
    return 0;
}

int flash_write_firmware_block(const uint8_t *buffer, size_t length, bool is_last_block) {
    int rc;
    rc = flash_img_buffered_write(&flash_ctx, buffer, length, is_last_block);
    
    if (rc != 0) {
        printk("Error writing flash image: %d\n", rc);
        return rc;
    }
    return 0;
}