#include "thread.h"
#include "heap.h"
#include "../drivers/serial.h" // Подставь свой путь к ком-порту

#define STACK_SIZE 1024

static thread_t* current_thread = NULL;
static thread_t* main_thread = NULL;
static uint32_t next_thread_id = 1;

// Объявляем функцию, которая скомпилирована в switch.asm
extern void asm_switch_context(uint32_t** old_esp, uint32_t* new_esp);

void kthread_init(void) {
    kprint_serial("THREADS: Initializing Round-Robin Scheduler...\n");

    main_thread = (thread_t*)kmalloc(sizeof(thread_t));
    main_thread->id = 0;
    main_thread->state = THREAD_RUNNING;
    main_thread->stack_limit = NULL; 
    main_thread->esp = NULL;         
    main_thread->next = main_thread; // Зацикливаем на себя

    current_thread = main_thread;
    kprint_serial("THREADS: Main kernel thread (kmain) registered.\n");
}

thread_t* kthread_create(void (*entry_point)(void)) {
    thread_t* new_thread = (thread_t*)kmalloc(sizeof(thread_t));
    void* stack_mem = kmalloc(STACK_SIZE);

    if (!new_thread || !stack_mem) {
        kprint_serial("THREADS ERROR: NO MEMORY FOR NEW THREAD!\n");
        return NULL; // или вешай здесь cpu_halt()
    }

    new_thread->id = next_thread_id++;
    new_thread->stack_limit = stack_mem;
    new_thread->state = THREAD_READY;

    // Стек растет вниз, встаем на самый конец выделенного куска
    uint32_t* stack_top = (uint32_t*)((uint32_t)stack_mem + STACK_SIZE);

    // Когда поток вызовет `ret` внутри asm_switch_context, 
    // он должен извлечь из стека адрес точки входа.
    stack_top--; 
    *stack_top = (uint32_t)entry_point; // Адрес возврата для ret

    // Теперь имитируем push четырех регистров, которые сохраняет asm_switch_context:
    // edi, esi, ebp, ebx
    stack_top--; *stack_top = 0; // edi
    stack_top--; *stack_top = 0; // esi
    stack_top--; *stack_top = 0; // ebp
    stack_top--; *stack_top = 0; // ebx

    new_thread->esp = stack_top;

    new_thread->esp = stack_top;

    // Круглый список: вставляем новый поток сразу после текущего
    new_thread->next = current_thread->next;
    current_thread->next = new_thread;

    kprint_serial("THREADS: New thread created successfully.\n");
    return new_thread;
}

void kthread_yield(void) {
    thread_t* old_thread = current_thread;
    thread_t* new_thread = current_thread->next;

    // Ищем живой поток по кругу
    while (new_thread->state != THREAD_READY && new_thread->state != THREAD_RUNNING) {
        new_thread = new_thread->next;
    }

    if (new_thread == old_thread) return; 

    current_thread = new_thread;
    if (old_thread->state == THREAD_RUNNING) {
        old_thread->state = THREAD_READY;
    }
    new_thread->state = THREAD_RUNNING;

    // Сама магия смены контекста
    asm_switch_context(&(old_thread->esp), new_thread->esp);
}