#include <kernel/printk.h>
#include <drivers/tty.h>
#include <drivers/serial.h>
#include <lib/printf.h>
#include <lib/string.h>
#include <kernel/sync.h>

#define DMESG_BUFFER_SIZE 65536 // 64 KB Ring Buffer

static char dmesg_buffer[DMESG_BUFFER_SIZE];
static size_t dmesg_head = 0;
static size_t dmesg_total = 0;
static spinlock_t printk_lock = SPINLOCK_INIT;

void printk_init(void) {
    memset(dmesg_buffer, 0, sizeof(dmesg_buffer));
    dmesg_head = 0;
    dmesg_total = 0;
}

static void dmesg_append(const char* str, size_t len) {
    for (size_t i = 0; i < len; i++) {
        dmesg_buffer[dmesg_head] = str[i];
        dmesg_head = (dmesg_head + 1) % DMESG_BUFFER_SIZE;
        if (dmesg_total < DMESG_BUFFER_SIZE) {
            dmesg_total++;
        }
    }
}

int vprintk(const char* fmt, va_list args) {
    char buffer[1024];
    int len = vsnprintf(buffer, sizeof(buffer), fmt, args);
    if (len <= 0) return 0;

    spin_lock(&printk_lock);

    // Skip loglevel prefix if present (e.g. "\0016")
    const char* display_ptr = buffer;
    if (buffer[0] == '\001' && buffer[1] >= '0' && buffer[1] <= '7') {
        display_ptr += 2;
        len -= 2;
    }

    // 1. Output to active TTY console
    tty_write(display_ptr, (size_t)len);

    // 2. Output to COM1 serial port for host debugging
    for (int i = 0; i < len; i++) {
        serial_write_char(display_ptr[i]);
    }

    // 3. Store in kernel dmesg ring buffer
    dmesg_append(display_ptr, (size_t)len);

    spin_unlock(&printk_lock);
    return len;
}

int printk(const char* fmt, ...) {
    va_list args;
    va_start(args, fmt);
    int ret = vprintk(fmt, args);
    va_end(args);
    return ret;
}

int kprintf(const char* fmt, ...) {
    va_list args;
    va_start(args, fmt);
    int ret = vprintk(fmt, args);
    va_end(args);
    return ret;
}

void dmesg_dump(void) {
    spin_lock(&printk_lock);
    if (dmesg_total < DMESG_BUFFER_SIZE) {
        tty_write(dmesg_buffer, dmesg_total);
    } else {
        // Output from head to end, then 0 to head
        tty_write(dmesg_buffer + dmesg_head, DMESG_BUFFER_SIZE - dmesg_head);
        tty_write(dmesg_buffer, dmesg_head);
    }
    spin_unlock(&printk_lock);
}

const char* dmesg_get_buffer(size_t* size_out) {
    if (size_out) *size_out = dmesg_total;
    return dmesg_buffer;
}