// This is NEW CODE written by Zak, separated from the old proc.c code

#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "riscv.h"
#include "sched.h"
#include "defs.h"
#include "sched.h"
#include "proc.h"
#include "bitarray.h"
#include "fpu.h"
#include "kprintf.h"

struct bitarray *runnablearrays[NPRIO];
struct bitarray *exhaustedarrays[NPRIO];
struct bitarray *sleeping;
sched_core_t sched_cores[SCHED_CORE_MAX];
extern struct proc proc[NPROC];
int timeslice_max;
int timeslice_min;

int sched_timesliceforpriority(int priority) {
  if (priority < 0) priority = 0;
  if (priority >= NPRIO) priority = NPRIO - 1;
  return timeslice_min + ((timeslice_max - timeslice_min) / (priority + 1));
}

/* Attempts to restate a process. This should be avoided in situations where
 * you're already locking other processes, you may need to use more specific code in
 * some cases! This function assumes that the caller has locked the process.
 */
void sched_restate_alreadylocked(struct proc* p, sched_state_t s) {
  if (p->state == s) {
    return;
  }
  if (p->state == SCHED_STATE_RUNNABLE) {
    bitarray_set(runnablearrays[p->prio], (int) (p - proc), 0);
    bitarray_set(exhaustedarrays[p->prio], (int) (p - proc), 0);
  } else if (p->state == SCHED_STATE_SLEEPING) {
    bitarray_set(sleeping, (int) (p - proc), 0);
  }
  if (s == SCHED_STATE_RUNNABLE) {
    //if (p->timeslice > 0) {
      bitarray_set(runnablearrays[p->prio], (int) (p - proc), 1);
    //} else {
    //  bitarray_set(exhaustedarrays[p->prio], (int) (p - proc), 1);
    //}
    for (int i = 0; i < 64; i++) {
      unsigned long foo = 1 << i;
      if (p->affinitymask & foo) {
        sched_cores[i].preempted = 1;
      }
    }
  } else if (s == SCHED_STATE_SLEEPING) {
    bitarray_set(sleeping, (int) (p - proc), 1);
  }
  p->state = s;
}

// My simplified function to set processor affinity of the caller (TODO: Finish this...)
int affin(uint64 mask) {
  myproc()->affinitymask = mask;
  yield();
  return 0;
}

/* Wake any processes sleeping on the given pointer, ideally preempting
 * any lower priority processes that might be running.
 * 
 * NOTE: This can't be called while locking any processes.
 */
void sched_wake(void* pointer) {
  struct proc* me = myproc();
  for (int i = 0; i >= 0 && i < NPROC; i = bitarray_findlowest(sleeping, i+1)) {
    struct proc* pr = &proc[i];
    if (pr == me) {
      // Skip the calling process!
    } else {
      acquire(&pr->lock);
      if (pr->chan == pointer) {
        // If the pointer matches it will almost certainly be a candidate,
        // but make sure it's actually sleeping first!
        if (pr->state == SCHED_STATE_SLEEPING) {
          // Then set the state to SCHED_STATE_RUNNABLE and preempt any lower priority
          sched_restate_alreadylocked(pr, SCHED_STATE_RUNNABLE);
        }
      }
      release(&pr->lock);
    }
  }
}

/* Scans for the next SCHED_STATE_RUNNABLE-looking candidate, ideally WITHOUT locking each proc. */
struct proc* sched_nextcandidate(int prio, struct proc* p) {
  p++;
  if (p < &proc[NPROC]) {
    int idx = bitarray_findlowest(runnablearrays[prio], p->prio == prio ? (int) (p - proc) : 0);
    if (idx >= 0) return &proc[idx];
  }
  /*for(; p < &proc[NPROC]; p++) {
    if (p->state == SCHED_STATE_RUNNABLE) {
      return p;
    }
  }*/
  return 0ULL;
}

// Called when the scheduler is idle to wait for interrupts
void sched_waitforinterrupts() {
  intr_on(); // Ensure interrupts are actually on
#ifdef _ZCC
  // Use straightforward inline assembler syntax
  __asm {
    wfi
  };
#else
  asm volatile("wfi"); // GNU syntax
#endif
}

int sched_reviveexhausted(int prioidx) {
  int exhausted = 0;
  int count = 0;
  while ((exhausted = bitarray_poplowest(exhaustedarrays[prioidx], exhausted)) >= 0) {
    //printf("Reviving process #%d (index #%d)\n", proc[exhausted].pid, exhausted);
    proc[exhausted].timeslice = sched_timesliceforpriority(prioidx);
    bitarray_set(runnablearrays[proc[exhausted].prio], exhausted, 1);
    count++;
  }
  return count;
}

// The scheduler is entered on each CPU core after initialisation, it loops
// selecting processes to switch to and (ideally) powering down the CPU when
// none are found. The new scheduler handles CPU affinity as well as priority,
// but is not yet an optimal algorithm.
void scheduler() {
  sched_core_t* thiscpu = SCHED_CORE_THIS_NOINTERRUPTS();
  uint64 affinitymask = 1ULL << SCHED_CORE_THISNUMBER_NOINTERRUPTS();
  thiscpu->allowstarvation = 1; //(SCHED_CORE_THISNUMBER_NOINTERRUPTS() > 0 ? 1 : 0); // Default policy, run low-priority processes only on the main core
  thiscpu->process = (void*)0ULL;
  
  while (1) {
    int found_any = 0;
    int found_here = 0;
    int prioidx;
    intr_on(); // Enable interrupts in case they are disabled
    for (prioidx = 0; prioidx < NPRIO; prioidx++) {
      do {
        found_here = 0;
        struct proc* thisproc;
        //printf("PRIO%d\n", prioidx);
        for (thisproc = proc; thisproc != 0ULL; thisproc = sched_nextcandidate(prioidx, thisproc)) {
          acquire(&thisproc->lock);
          
          if ((thisproc->affinitymask & affinitymask) && thisproc->state == SCHED_STATE_RUNNABLE) {
            bitarray_set(runnablearrays[prioidx] /* check this is == runnablearrays[thisproc->prio] ? */, (int) (thisproc - proc), 0);
            bitarray_set(runnablearrays[thisproc->prio], (int) (thisproc - proc), 0);
            //enterprocess:
            thisproc->state = SCHED_STATE_RUNNING; // Set the processes state to running, then unset it in the runnable arrays:
            thiscpu->process = thisproc; // Set the CPU's process to this process
            sched_switchcontext(&thiscpu->registers, &thisproc->context); // Perform the actual context switch
            thisproc->timeslice -= 10000;
            //printf("Timeslice=%d\n", thisproc->timeslice);
            
            if (thisproc->state == SCHED_STATE_RUNNING || thisproc->state == SCHED_STATE_RUNNABLE) {
              if (thisproc->timeslice > 0) {
                /*
                thisproc->state = SCHED_STATE_RUNNABLE;
                bitarray_set(runnablearrays[thisproc->prio], (int) (thisproc - proc), 1);
                //goto enterprocess;
                found_any = found_here = 1;
                */
              } else {
                //thisproc->state = SCHED_STATE_RUNNABLE;
                //bitarray_set(runnablearrays[thisproc->prio], (int) (thisproc - proc), 0);
                //bitarray_set(exhaustedarrays[thisproc->prio], (int) (thisproc - proc), 1);
              }
            }
            
            // Check if the FPU status is dirty
            int fpustat = fpu_status_read();
            if (fpustat != 0) {
              printf("FPU Status %d, saving state\n", fpustat);
              fpu_save(&thiscpu->process->fpu_context);
              thiscpu->process->fpu_active = 0;
              thiscpu->process->fpu_saved = 1;
              fpu_status_write(0);
              printf("FPU Status %d after saving state\n", fpu_status_read());
            }
            thiscpu->process = (void*)0ULL; // Process is finished, immediately exit it's context
          }
          
          release(&thisproc->lock);
          
          if (found_here) {
            int check = prioidx;
            for (prioidx = 0; prioidx < check-1 && (bitarray_findlowest(runnablearrays[prioidx], 0) < 0 || (thiscpu->allowstarvation && (bitarray_findlowest(exhaustedarrays[prioidx], 0) < 0))); prioidx++) {
              found_here = 0; // Reset the outer loop if more important processes may be runnable
              //break; // End the inner loop
            }
          }
        }
      } while (found_here);
      if (thiscpu->allowstarvation) {
       // sched_reviveexhausted(prioidx);
      }
    }
    if (!thiscpu->allowstarvation) {
      for (int i = 0; i < NPRIO; i++) {
        if (sched_reviveexhausted(i) > 0) {
          found_any = 1;
        }
      }
    }
    if (!found_any) {
      //printf("Not found any processes\n");
      // This function will enable interrupts and put the CPU in wait mode
      sched_waitforinterrupts();
    }
  }
}

const char* sched_statename(sched_state_t s, int padded) {
  switch (s) {
    case SCHED_STATE_UNUSED:
      return padded ? "SCHED_STATE_UNUSED  " : "SCHED_STATE_UNUSED";
    case SCHED_STATE_USED:
      return padded ? "SCHED_STATE_USED    " : "SCHED_STATE_USED";
    case SCHED_STATE_SLEEPING:
      return padded ? "SCHED_STATE_SLEEPING" : "SCHED_STATE_SLEEPING";
    case SCHED_STATE_RUNNABLE:
      return padded ? "SCHED_STATE_RUNNABLE" : "SCHED_STATE_RUNNABLE";
    case SCHED_STATE_RUNNING:
      return padded ? "SCHED_STATE_RUNNING " : "SCHED_STATE_RUNNING";
    case SCHED_STATE_ZOMBIE:
      return padded ? "SCHED_STATE_ZOMBIE  " : "SCHED_STATE_ZOMBIE";
    default:
      return padded ? "BADSTATE" : "BADSTATE";
  }
}

// Called when CTRL+p is pressed
void sched_dumpstatus() {
  int i; // Iterator variable is reused
  
  printf("\n"); // Print a newline to improve alignment
  
  for (i = 0; i < NPROC; i++) {
    if (proc[i].state != SCHED_STATE_UNUSED) {
      printf("proc[%d] pid=%d state=%s name=\"%s\" prio=%d maxprio=%d\n", i, proc[i].pid, sched_statename(proc[i].state, 1), proc[i].name, proc[i].prio, proc[i].maxprio);
    }
  }
  
  printf("Sleeping: ");
  for (i = 0; i < NPROC; i++) {
	if (bitarray_getnolock(sleeping, i)) {
      printf("%d ", i);
    }
  }
  printf("\n");
  int prio;
  for (prio = 0; prio < NPRIO; prio++) {
    printf("Runnables priority #%d (main set): ", prio);
    for (i = 0; i < NPROC; i++) {
      if (bitarray_getnolock(runnablearrays[prio], i)) {
        printf("%d ", i);
      }
    }
    printf("\n");
    printf("Runnables priority #%d (exhausted): ", prio);
    for (i = 0; i < NPROC; i++) {
      if (bitarray_getnolock(exhaustedarrays[prio], i)) {
        printf("%d ", i);
      }
    }
    printf("\n");
  }
  printf("RAM: %d MB total %d MB free\n", (int) (physpg_totalram()/MB), (int)(physpg_freeram()/MB));
}

/* 
 * hypothetical thread-sync code I never finished
int
thrsync() {
  struct proc* p = myproc();
  struct proc* mainthread = p->mainthread;
  if (!mainthread) {
    return 0;
  }
  int unsynced;
  do {
	unsynced = 0;
    for (int i = 0; i < NPROC; i++) {
      struct proc* thr = &proc[i];
      if (thr != p) {
        acquire(&thr->lock);
        if (thr->mainthread == mainthread) {
          // TODO ...
		}
        release(&thr->lock);
      }
    }
  } while (unsynced);
  return unsynced;
}
*/