#include "../include/task.h"
#include "../include/stdio.h"
#include "../include/paging.h"
#include "../include/string.h"
#include "../include/liballoc.h"
#include "../include/gdt.h"
#include "../include/isr.h"

#include <stddef.h>

Process* current = 0;
Process* queue = 0;

uint32_t irq_disable_counter = 0;
uint32_t pid_counter;
uint32_t task_count;

extern void trapret(void);
extern void switchProcess(Process* next);

Process* task_create(uint32_t func, uint32_t user_esp, uint32_t ring, uint32_t* pagedir)
{
    void *p_physical = alloc_page();
    Process* p = (Process*)((uint32_t)p_physical+0xC0000000);
    map_page(p_physical, (void*)p, PAGE_PRESENT | PAGE_RW);
    memset(p, 0, sizeof(Process));
    p->pid = ++pid_counter;
    p->state = Ready;
    p->ring = ring;

    // Выделяем стек ядра
    void *kstack_physical = alloc_page();
    p->kstack = (char*)((uint32_t)kstack_physical+0xC0000000);
    map_page(kstack_physical, p->kstack, PAGE_PRESENT | PAGE_RW);
    debug_log("kstack: %X!!!!!!!!!!!!\n", p->kstack);

    p->kstack_top = (uint32_t)p->kstack + KSTACKSIZE;

    uint8_t* sp = (uint8_t*)(p->kstack + KSTACKSIZE);

    sp -= sizeof(TrapFrame);
    p->tf = (TrapFrame*)sp;
    memset(p->tf, 0, sizeof(TrapFrame));

    p->tf->cs = ring == 3 ? (SEG_UCODE | DPL_USER) : (SEG_KCODE);
    p->tf->ds = ring == 3 ? (SEG_UDATA | DPL_USER) : (SEG_KDATA);
    p->tf->es = p->tf->ds;
    p->tf->fs = p->tf->ds;
    p->tf->gs = p->tf->ds;

    p->tf->usermode_ss = ring == 3 ? (SEG_UDATA | DPL_USER) : (SEG_KDATA);

    p->tf->eflags = FL_IF;
    p->tf->eip = (uint32_t)func;

    p->pagedir = pagedir;

    if(ring == 3)
    {
        p->tf->usermode_esp = user_esp;

    }

    sp -= sizeof(Context);
    p->context = (Context*)sp;
    memset(p->context, 0, sizeof(Context));

    p->context->eip = (uint32_t)trapret;
    p->kesp = (uint32_t)sp;

    p->wake_up_time = 0;

    p->next = 0;

    if (!queue) {
        queue = p;
        current = queue;
    } else {
        Process* curr = queue;
        while (curr->next)
            curr = curr->next;
        curr->next = p;
    }
    return p;
}

// Остальные функции остаются без изменений
void scheduler_lock()
{
    asm volatile("cli");
    irq_disable_counter++;
}

void scheduler_unlock()
{
    irq_disable_counter--;
    if (irq_disable_counter == 0)
        asm volatile("sti");
}

//TODO: there's an issue with this. If there's only one task in the queue (or is it there really?),
//the scheduler doesn't switch the context to it, but just does nothing
void schedule() {
    scheduler_lock();

    Process* prev = NULL;
    Process* p = queue;
    while (p) {
        if (p->state == Terminated) {
            Process* next_proc = p->next;
            // Unlink from the queue
            if (prev) {
                prev->next = next_proc;
            } else {
                queue = next_proc;
            }

            // If the terminated process is the current one, we must not switch to it.
            if (p == current) {
                current = NULL;
            }

            destroy_page_dir(p->pagedir);
            free_page((void*)virt_to_phys(p->kstack));
            free_page((void*)virt_to_phys(p));

            p = next_proc;
        } else {
            prev = p;
            p = p->next;
        }
    }

    // If 'current' was terminated or this is the first run, find a new process.
    if (!current || current->state != Ready) {
        current = queue; // Start search from the beginning
    }

    if(current->wake_up_time > timer_ticks)
    {
        current = queue;
    }

    // Find the next ready process to run
    Process* next = current ? current->next : NULL;
    if (!next) next = queue;

    // Iterate through the list to find a process that is ready to run
    Process* start_node = next;
    while (next && next->state != Ready) {
        next = next->next;
        if (!next) next = queue; // Wrap around
        if (next == start_node) { // Full circle, no ready process
            scheduler_unlock();
            return;
        }
    }

    if (next && next != current) {
        tss.esp0 = next->kstack_top;
        switchProcess(next);
    }

    scheduler_unlock();
}

void idle()
{
    while (1) {}
}

void task1()
{
    while (1) {
        printf("task1\n");
    }
}

void task2()
{
	int a = 0;
    while (1) { a++; printf("task2 %X\n", a); }
}

void scheduler_init()
{
    uint32_t* kernel_tasks_pagedir = (uint32_t*)create_page_dir();

    task_create((uint32_t)idle, 0, 0, kernel_tasks_pagedir);
}

void task_kill(Process* proc) {
    //TODO: that proc != current is sus
    if (!proc || proc != current) return;

    scheduler_lock();
    proc->state = Terminated;

    // Remove from process queue
    if (queue == proc) {
        queue = proc->next;
    } else {
        Process* prev = queue;
        while (prev && prev->next != proc) prev = prev->next;
        if (prev) prev->next = proc->next;
    }

    // Schedule next process
    Process* next = queue;
    while (next && next->state != Ready) next = next->next;
    if (!next) next = queue;  // Fallback to idle if needed
    
    current = next;
    set_page_dir(virt_to_phys(next->pagedir));
    
    // Free resources safely from new context
    destroy_page_dir(proc->pagedir);
    free_page((void*)virt_to_phys(proc->kstack));
    free_page((void*)virt_to_phys(proc));
    
    scheduler_unlock();
    switchProcess(next);
}