/* context-switch-preemption.cc
 * This program shows an example of how we can make two threads and
 * switch between them when a timer fires.  This version is bug free,
 * it enables interrupts at the start of a thread's execution.
 */

#include <iostream>
#include "timer.hh"

using namespace std;

// A thread is a stack plus the current top of the stack
typedef struct Thread {
    char stack[8192];
    char *saved_rsp;
} Thread;


/* This sets up a thread so that the first time it is switched to,
 * it "resumes" by starting to run the specified func.
 */
static const int kNumRegistersToSave = 6;
Thread create_thread(void (*func)()) {
    Thread new_thread;
    void *stack_top = new_thread.stack + sizeof(new_thread.stack);

    // Make it look like this thread was about to start func,
    // and then context switched.  In other words, once
    // the saved registers are popped, ret should take us to
    // the start of func.
    *(void **)((char *)stack_top - sizeof(void *)) = (void *)func;

    // Move the stack pointer downwards by 6 registers to make it appear like
    // we pushed 6 registers on when we previously context switched
    new_thread.saved_rsp = (char *)stack_top - sizeof(void *) - kNumRegistersToSave * sizeof(long);
    return new_thread;
}

/* This function context_switch is written in assembly!  It switches from running 'current'
 * to running 'next'.  To do this, tt first pushes 'current's registers onto the stack
 * to save them, then updates 'current' with the new stack pointer.  Then it loads 'next's
 * stack pointer and from then on we are operating on 'next's stack, not current's!  So we
 * proceed to pop off all the registers we left on 'next's' stack from before, and then
 * return, and this will return back to the function that 'next' had been executing. 
 * 'current' will be resumed when another thread calls context_switch to switch back to
 * its stack.
 */
extern "C" void context_switch(Thread& current, Thread& next);

asm(R"(.text
	.global context_switch
context_switch:
	pushq %rbp
	pushq %rbx
	pushq %r12
	pushq %r13
	pushq %r14
	pushq %r15
	movq %rsp,0x2000(%rdi)	    # saved current's stack pointer
    movq 0x2000(%rsi),%rsp	    # load next's stack pointer
	popq %r15
	popq %r14
	popq %r13
	popq %r12
	popq %rbx
	popq %rbp
	ret     # return back to the function **this stack** was running
)");


Thread main_thread;
Thread other_thread;

Thread *current_thread;
Thread *nonrunning_thread;

void other_func() {
    intr_enable(true);
    while (true) {
        cout << "Other thread here!  Hello." << endl;
    }
}

void timer_interrupt_handler() {
    // We want to run the waiting thread and pause the current thread
    Thread *temp = current_thread;
    current_thread = nonrunning_thread;
    nonrunning_thread = temp;

    context_switch(*nonrunning_thread, *current_thread);
}

int main(int argc, char *argv[]) {
    // Initialize other_thread to run other_func
    other_thread = create_thread(other_func);

    // Keep track of which thread is running and not
    current_thread = &main_thread;
    nonrunning_thread = &other_thread;

    // Fire the timer every 500,000 microseconds to context switch
    timer_init(500000, timer_interrupt_handler);

    while (true) {
        cout << "I am the main thread" << endl;
    }
}