#include <stdarg.h>
#include <stdbool.h>
#include "stdio.h"
#include "debug.h"
#include "video/render.h"
#include "arch/x86_64/e9.h"
#include "mp/spinlock.h"
#include <stddef.h>
#include <limits.h>

static const uint32_t g_LogSeverityColors[] =
{
    [LVL_DEBUG]    = 0xAAAAAA, // light gray
    [LVL_INFO]     = 0xFFFFFF, // white
    [LVL_WARN]     = 0xFFFF00, // yellow
    [LVL_ERROR]    = 0xFF0000, // red
    [LVL_CRITICAL] = 0xFFFFFF, // white (can do red background separately if you want)
};

static spinlock_t s_printf_lock = SPINLOCK_INIT;

typedef struct {
    char* buf;
    size_t size;
    size_t pos;
} snprintf_ctx_t;

static void buf_putc(snprintf_ctx_t* ctx, char c)
{
    if (ctx->pos + 1 < ctx->size) {
        ctx->buf[ctx->pos] = c;
    }
    ctx->pos++;
}

static void buf_puts(snprintf_ctx_t* ctx, const char* str)
{
    while (*str) {
        buf_putc(ctx, *str++);
    }
}

void fputc(char c)
{
    putchar(c);
    e9_putc(c);
}

void fputs(const char* str)
{
    const char* s = str;
    while (*s) fputc(*s++);
}

void fputs_colored(const char* str, uint32_t color) {
    while (*str) putchar_colored(*str++, color);
    while (*str) fputc(*str++);
}

const char g_HexChars[] = "0123456789abcdef";

static void buf_print_unsigned(snprintf_ctx_t* ctx, uint64_t number, int radix, int width, int zero_pad)
{
    char buffer[32];
    int pos = 0;

    do {
        buffer[pos++] = g_HexChars[number % radix];
        number /= radix;
    } while (number > 0);

    // Apply zero-padding / width
    if (zero_pad) {
        while (pos < width)
            buffer[pos++] = '0';
    }

    while (--pos >= 0)
        buf_putc(ctx, buffer[pos]);
}

static void buf_print_signed(snprintf_ctx_t* ctx, int64_t number, int radix, int width, int zero_pad)
{
    if (number < 0) {
        buf_putc(ctx, '-');
        buf_print_unsigned(ctx, (uint64_t)(-number), radix, width, zero_pad);
    } else {
        buf_print_unsigned(ctx, (uint64_t)number, radix, width, zero_pad);
    }
}



int snprintf(char* buf, size_t size, const char* fmt, ...)
{
    va_list args;
    va_start(args, fmt);
    int ret = vsnprintf(buf, size, fmt, args);
    va_end(args);
    return ret;
}

static void fprintf_unsigned(uint64_t number, int radix, int width, int zero_pad)
{
    char buffer[32];
    int pos = 0;

    do {
        buffer[pos++] = g_HexChars[number % radix];
        number /= radix;
    } while (number > 0);

    if (zero_pad) {
        while (pos < width)
            buffer[pos++] = '0';
    }

    while (--pos >= 0)
        fputc(buffer[pos]);
}

static void fprintf_signed(int64_t number, int radix, int width, int zero_pad)
{
    if (number < 0) {
        fputc('-');
        fprintf_unsigned((uint64_t)(-number), radix, width, zero_pad);
    } else {
        fprintf_unsigned((uint64_t)number, radix, width, zero_pad);
    }
}


void vfprintf(const char* fmt, va_list args)
{
    while (*fmt) {
        if (*fmt != '%') {
            fputc(*fmt++);
            continue;
        }
        fmt++; // skip %

        /* flags */
        int zero_pad = 0;
        if (*fmt == '0') {
            zero_pad = 1;
            fmt++;
        }

        /* width */
        int width = 0;
        while (*fmt >= '0' && *fmt <= '9') {
            width = width * 10 + (*fmt++ - '0');
        }

        /* precision */
        int precision = -1;
        if (*fmt == '.') {
            fmt++;
            precision = 0;
            while (*fmt >= '0' && *fmt <= '9') {
                precision = precision * 10 + (*fmt++ - '0');
            }
        }

        /* length modifier */
        int is_long_long = 0;
        if (*fmt == 'l') {
            fmt++;
            if (*fmt == 'l') {
                is_long_long = 1;
                fmt++;
            }
        } else if (*fmt == 'z') {           // %zu, %zx
            is_long_long = 1;
            fmt++;
        } else if (*fmt == 'h') {
            fmt++; // ignore h/hh
        }

        /* specifier */
        switch (*fmt++) {
            case 'c':
                fputc((char)va_arg(args, int));
                break;

            case 's': {
                const char* s = va_arg(args, const char*);
                if (!s) s = "(null)";
                if (precision >= 0) {
                    while (*s && precision--) fputc(*s++);
                } else {
                    fputs(s);
                }
                break;
            }

            case 'd':
            case 'i':
                if (is_long_long)
                    fprintf_signed(va_arg(args, long long), 10, width, zero_pad);
                else
                    fprintf_signed(va_arg(args, long), 10, width, zero_pad);
                break;

            case 'u':
                if (is_long_long)
                    fprintf_unsigned(va_arg(args, unsigned long long), 10, width, zero_pad);
                else
                    fprintf_unsigned(va_arg(args, unsigned long), 10, width, zero_pad);
                break;

            case 'x':
            case 'X':
                if (is_long_long)
                    fprintf_unsigned(va_arg(args, unsigned long long), 16, width, zero_pad);
                else
                    fprintf_unsigned(va_arg(args, unsigned long), 16, width, zero_pad);
                break;

            case 'p':
                fputs("0x");
                fprintf_unsigned(va_arg(args, uint64_t), 16, 16, 1); // always full 64-bit
                break;

            case 'o':
                if (is_long_long)
                    fprintf_unsigned(va_arg(args, unsigned long long), 8, width, zero_pad);
                else
                    fprintf_unsigned(va_arg(args, unsigned long), 8, width, zero_pad);
                break;

            case '%':
                fputc('%');
                break;

            default:
                fputc('%');
                fputc(*(fmt - 1));
                break;
        }
    }
}

int vsnprintf(char* buf, size_t size, const char* fmt, va_list args)
{
    snprintf_ctx_t ctx = { .buf = buf, .size = size, .pos = 0 };

    while (*fmt) {
        if (*fmt != '%') {
            buf_putc(&ctx, *fmt++);
            continue;
        }
        fmt++;

        int zero_pad = 0;
        if (*fmt == '0') { zero_pad = 1; fmt++; }

        int width = 0;
        while (*fmt >= '0' && *fmt <= '9') {
            width = width * 10 + (*fmt++ - '0');
        }

        int precision = -1;
        if (*fmt == '.') {
            fmt++;
            precision = 0;
            while (*fmt >= '0' && *fmt <= '9') {
                precision = precision * 10 + (*fmt++ - '0');
            }
        }

        int is_long_long = 0;
        if (*fmt == 'l') {
            fmt++;
            if (*fmt == 'l') { is_long_long = 1; fmt++; }
        } else if (*fmt == 'z') {
            is_long_long = 1;
            fmt++;
        } else if (*fmt == 'h') {
            fmt++;
        }

        switch (*fmt++) {
            case 'c':
                buf_putc(&ctx, (char)va_arg(args, int));
                break;

            case 's': {
                const char* s = va_arg(args, const char*);
                if (!s) s = "(null)";
                int max = (precision >= 0) ? precision : INT_MAX;
                while (*s && max--) buf_putc(&ctx, *s++);
                break;
            }

            case 'd':
            case 'i':
                if (is_long_long)
                    buf_print_signed(&ctx, va_arg(args, long long), 10, width, zero_pad);
                else
                    buf_print_signed(&ctx, va_arg(args, long), 10, width, zero_pad);
                break;

            case 'u':
                if (is_long_long)
                    buf_print_unsigned(&ctx, va_arg(args, unsigned long long), 10, width, zero_pad);
                else
                    buf_print_unsigned(&ctx, va_arg(args, unsigned long), 10, width, zero_pad);
                break;

            case 'x':
            case 'X':
                if (is_long_long)
                    buf_print_unsigned(&ctx, va_arg(args, unsigned long long), 16, width, zero_pad);
                else
                    buf_print_unsigned(&ctx, va_arg(args, unsigned long), 16, width, zero_pad);
                break;

            case 'p':
                buf_puts(&ctx, "0x");
                buf_print_unsigned(&ctx, va_arg(args, uint64_t), 16, 16, 1);
                break;

            case 'o':
                if (is_long_long)
                    buf_print_unsigned(&ctx, va_arg(args, unsigned long long), 8, width, zero_pad);
                else
                    buf_print_unsigned(&ctx, va_arg(args, unsigned long), 8, width, zero_pad);
                break;

            case '%':
                buf_putc(&ctx, '%');
                break;

            default:
                buf_putc(&ctx, '%');
                buf_putc(&ctx, *(fmt - 1));
                break;
        }
    }

    if (ctx.size > 0) {
        if (ctx.pos < ctx.size)
            ctx.buf[ctx.pos] = '\0';
        else
            ctx.buf[ctx.size - 1] = '\0';
    }

    return (int)ctx.pos;
}

void printf(const char* fmt, ...)
{
    uint64_t flags;
    spinlock_acquire_irqsave(&s_printf_lock, &flags);

    va_list args;
    va_start(args, fmt);

    vfprintf(fmt, args);

    va_end(args);

    spinlock_release_irqrestore(&s_printf_lock, flags);
}

void fprintf(const char* fmt, ...)
{
    va_list args;
    va_start(args, fmt);
    vfprintf(fmt, args);
    va_end(args);
}