/*
    * wallpaper.hpp
    * MontaukOS Desktop - JPEG wallpaper loading, scaling, and directory scanning
    * Copyright (c) 2026 Daniel Hammer
*/

#pragma once

#include <montauk/syscall.h>
#include <montauk/string.h>
#include <montauk/heap.h>
#include <montauk/thread.h>
#include <gui/gui.hpp>
#include <gui/desktop.hpp>

// Forward-declare stb_image functions (implementation in libjpeg.a).
// We avoid including stb_image.h directly because its declaration section
// pulls in <stdlib.h> which is unavailable in the desktop's freestanding build.
extern "C" {
    unsigned char* stbi_load_from_memory(const unsigned char* buffer, int len,
                                         int* x, int* y, int* channels_in_file,
                                         int desired_channels);
    void stbi_image_free(void* retval_from_stbi_load);
    const char* stbi_failure_reason(void);
}

namespace gui {

// ============================================================================
// Wallpaper loading
// ============================================================================

// Decode `path` and scale it to cover (screen_w x screen_h). Returns a newly
// montauk::malloc'd ARGB buffer the caller owns, or nullptr on failure. This
// is the heavy lifting; it touches no shared state and is safe to run from a
// worker thread.
inline uint32_t* wallpaper_decode_scaled(const char* path,
                                         int screen_w, int screen_h) {
    int fd = montauk::open(path);
    if (fd < 0) return nullptr;

    uint64_t size = montauk::getsize(fd);
    if (size == 0 || size > 16 * 1024 * 1024) {
        montauk::close(fd);
        return nullptr;
    }

    uint8_t* filedata = (uint8_t*)montauk::malloc(size);
    if (!filedata) { montauk::close(fd); return nullptr; }

    int bytes_read = montauk::read(fd, filedata, 0, size);
    montauk::close(fd);
    if (bytes_read <= 0) { montauk::mfree(filedata); return nullptr; }

    int img_w, img_h, channels;
    unsigned char* rgb = stbi_load_from_memory(filedata, bytes_read,
                                               &img_w, &img_h, &channels, 3);
    montauk::mfree(filedata);
    if (!rgb) return nullptr;

    int dst_w = screen_w;
    int dst_h = screen_h;
    uint32_t* scaled = (uint32_t*)montauk::malloc((uint64_t)dst_w * dst_h * 4);
    if (!scaled) { stbi_image_free(rgb); return nullptr; }

    int src_crop_w, src_crop_h, src_x0, src_y0;
    if ((int64_t)img_w * dst_h > (int64_t)img_h * dst_w) {
        src_crop_h = img_h;
        src_crop_w = (int)((int64_t)img_h * dst_w / dst_h);
        src_x0 = (img_w - src_crop_w) / 2;
        src_y0 = 0;
    } else {
        src_crop_w = img_w;
        src_crop_h = (int)((int64_t)img_w * dst_h / dst_w);
        src_x0 = 0;
        src_y0 = (img_h - src_crop_h) / 2;
    }

    for (int y = 0; y < dst_h; y++) {
        int sy = src_y0 + (int)((int64_t)y * src_crop_h / dst_h);
        if (sy < 0) sy = 0;
        if (sy >= img_h) sy = img_h - 1;
        for (int x = 0; x < dst_w; x++) {
            int sx = src_x0 + (int)((int64_t)x * src_crop_w / dst_w);
            if (sx < 0) sx = 0;
            if (sx >= img_w) sx = img_w - 1;
            int si = (sy * img_w + sx) * 3;
            scaled[y * dst_w + x] = 0xFF000000u
                | ((uint32_t)rgb[si] << 16)
                | ((uint32_t)rgb[si + 1] << 8)
                | (uint32_t)rgb[si + 2];
        }
    }

    stbi_image_free(rgb);
    return scaled;
}

// Synchronous loader retained for code paths that genuinely want to block
// (e.g. desktop boot, where there is no UI to keep responsive yet).
inline bool wallpaper_load(DesktopSettings* s, const char* path,
                           int screen_w, int screen_h) {
    uint32_t* scaled = wallpaper_decode_scaled(path, screen_w, screen_h);
    if (!scaled) return false;

    if (s->bg_wallpaper) montauk::mfree(s->bg_wallpaper);
    s->bg_wallpaper   = scaled;
    s->bg_wallpaper_w = screen_w;
    s->bg_wallpaper_h = screen_h;
    montauk::strncpy(s->bg_image_path, path, 127);
    s->bg_image = true;
    s->bg_gradient = false;
    return true;
}

// ============================================================================
// Async wallpaper loader
//
// Decoding a screen-sized JPEG blocks the desktop event loop. The async
// loader spawns a worker thread per click; rapid cycling is tolerated via a
// generation counter -- only the highest-gen worker's output is published.
// The main thread calls wallpaper_poll() each frame to swap in any new
// buffer (it owns the actual handoff and the free of the previous one).
// ============================================================================

struct WallpaperPublishSlot {
    montauk::Mutex lock;
    uint64_t latest_gen;       // bumped on every load request
    uint64_t pending_gen;      // gen of pending_pixels (0 = empty)
    uint64_t applied_gen;      // highest gen consumed by the main thread
    uint32_t* pending_pixels;  // worker-produced ARGB buffer awaiting handoff
    int pending_w, pending_h;
    char pending_path[128];
};

inline WallpaperPublishSlot& wallpaper_slot() {
    static WallpaperPublishSlot s{};
    return s;
}

struct WallpaperJob {
    uint64_t gen;
    int screen_w, screen_h;
    char path[256];
};

inline int wallpaper_worker(void* raw) {
    auto* job = (WallpaperJob*)raw;
    uint32_t* scaled = wallpaper_decode_scaled(job->path, job->screen_w, job->screen_h);

    auto& slot = wallpaper_slot();
    slot.lock.lock();
    bool publish = scaled != nullptr
        && job->gen > slot.applied_gen
        && job->gen > slot.pending_gen;
    if (publish) {
        if (slot.pending_pixels) {
            // A stale (lower-gen) pending buffer was never consumed -- drop it.
            montauk::mfree(slot.pending_pixels);
        }
        slot.pending_pixels = scaled;
        slot.pending_w = job->screen_w;
        slot.pending_h = job->screen_h;
        montauk::strncpy(slot.pending_path, job->path, 127);
        slot.pending_gen = job->gen;
    }
    slot.lock.unlock();

    if (!publish && scaled) montauk::mfree(scaled);
    montauk::mfree(job);
    return 0;
}

// Kick off a background decode. Returns true if the worker thread was
// spawned; the actual wallpaper swap happens later when wallpaper_poll
// observes the published buffer.
inline bool wallpaper_load_async(const char* path,
                                 int screen_w, int screen_h) {
    auto* job = (WallpaperJob*)montauk::malloc(sizeof(WallpaperJob));
    if (!job) return false;
    montauk::strncpy(job->path, path, 255);
    job->screen_w = screen_w;
    job->screen_h = screen_h;

    auto& slot = wallpaper_slot();
    slot.lock.lock();
    job->gen = ++slot.latest_gen;
    slot.lock.unlock();

    int tid = montauk::thread_spawn(&wallpaper_worker, job);
    if (tid < 0) {
        montauk::mfree(job);
        return false;
    }
    return true;
}

// If a worker has produced a newer buffer than the main thread has applied,
// swap it into `s` and free the previous wallpaper. Returns true if a new
// wallpaper was installed (caller should mark the background dirty).
inline bool wallpaper_poll(DesktopSettings* s) {
    auto& slot = wallpaper_slot();
    slot.lock.lock();
    if (slot.pending_gen == 0 || slot.pending_gen <= slot.applied_gen) {
        slot.lock.unlock();
        return false;
    }
    uint32_t* new_pixels = slot.pending_pixels;
    int new_w = slot.pending_w;
    int new_h = slot.pending_h;
    char new_path[128];
    montauk::strncpy(new_path, slot.pending_path, 127);
    new_path[127] = '\0';
    slot.applied_gen = slot.pending_gen;
    slot.pending_pixels = nullptr;
    slot.pending_gen = 0;
    slot.lock.unlock();

    uint32_t* old = s->bg_wallpaper;
    s->bg_wallpaper = new_pixels;
    s->bg_wallpaper_w = new_w;
    s->bg_wallpaper_h = new_h;
    montauk::strncpy(s->bg_image_path, new_path, 127);
    s->bg_image = true;
    s->bg_gradient = false;
    if (old) montauk::mfree(old);
    return true;
}

inline void wallpaper_free(DesktopSettings* s) {
    if (s->bg_wallpaper) {
        montauk::mfree(s->bg_wallpaper);
        s->bg_wallpaper = nullptr;
        s->bg_wallpaper_w = 0;
        s->bg_wallpaper_h = 0;
    }
    s->bg_image_path[0] = '\0';
    s->bg_image = false;
}

// ============================================================================
// Directory scanning for JPEG files
// ============================================================================

static constexpr int WALLPAPER_MAX_FILES = 16;

struct WallpaperFileList {
    char names[WALLPAPER_MAX_FILES][64];
    int count;
};

// Scan dir_path for JPEG files and append to list.
// name_prefix is prepended to each stored filename (e.g. "Pictures/").
// Caller must zero list->count before the first call if starting fresh.
inline void wallpaper_scan_dir(const char* dir_path, WallpaperFileList* list,
                               const char* name_prefix = nullptr) {
    const char* raw_names[64];
    int total = montauk::readdir(dir_path, raw_names, 64);
    if (total <= 0) return;

    // Retain support for older ramdisks that returned paths from the VFS root.
    const char* after_drive = dir_path;
    for (int k = 0; after_drive[k]; k++) {
        if (after_drive[k] == ':' && after_drive[k + 1] == '/') {
            after_drive += k + 2;
            break;
        }
    }
    char prefix[256] = {0};
    int prefix_len = 0;
    if (after_drive[0] != '\0') {
        montauk::strcpy(prefix, after_drive);
        prefix_len = montauk::slen(prefix);
        if (prefix_len > 0 && prefix[prefix_len - 1] != '/') {
            prefix[prefix_len++] = '/';
            prefix[prefix_len] = '\0';
        }
    }

    int np_len = name_prefix ? montauk::slen(name_prefix) : 0;

    auto to_lower = [](char c) -> char {
        return (c >= 'A' && c <= 'Z') ? (char)(c + 32) : c;
    };

    for (int i = 0; i < total && list->count < WALLPAPER_MAX_FILES; i++) {
        const char* name = raw_names[i];

        // Strip prefix
        if (prefix_len > 0) {
            bool match = true;
            for (int k = 0; k < prefix_len; k++) {
                if (name[k] != prefix[k]) { match = false; break; }
            }
            if (match) name += prefix_len;
        }

        int nlen = montauk::slen(name);

        // Skip directories
        if (nlen > 0 && name[nlen - 1] == '/') continue;

        // Check for .jpg
        bool is_jpeg = false;
        if (nlen >= 4 &&
            to_lower(name[nlen - 4]) == '.' &&
            to_lower(name[nlen - 3]) == 'j' &&
            to_lower(name[nlen - 2]) == 'p' &&
            to_lower(name[nlen - 1]) == 'g') {
            is_jpeg = true;
        }
        // Check for .jpeg
        if (!is_jpeg && nlen >= 5 &&
            to_lower(name[nlen - 5]) == '.' &&
            to_lower(name[nlen - 4]) == 'j' &&
            to_lower(name[nlen - 3]) == 'p' &&
            to_lower(name[nlen - 2]) == 'e' &&
            to_lower(name[nlen - 1]) == 'g') {
            is_jpeg = true;
        }

        if (is_jpeg) {
            if (np_len > 0) {
                montauk::strncpy(list->names[list->count], name_prefix, 63);
                int cur = np_len < 63 ? np_len : 63;
                montauk::strncpy(list->names[list->count] + cur, name,
                                 63 - cur);
            } else {
                montauk::strncpy(list->names[list->count], name, 63);
            }
            list->count++;
        }
    }
}

// Scan Pictures/ subdirectory first, then home directory.
inline void wallpaper_scan_home(const char* home_dir, WallpaperFileList* list) {
    list->count = 0;

    char pictures[256];
    montauk::strcpy(pictures, home_dir);
    int hlen = montauk::slen(pictures);
    if (hlen > 0 && pictures[hlen - 1] != '/') {
        pictures[hlen++] = '/';
        pictures[hlen] = '\0';
    }
    montauk::strncpy(pictures + hlen, "Pictures", 256 - hlen);

    wallpaper_scan_dir(pictures, list, "Pictures/");
    wallpaper_scan_dir(home_dir, list);
}

} // namespace gui