// NEW CODE for setting up the floating point unit
#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "riscv.h"
#include "defs.h"
#include "sched.h"
#include "proc.h"
#include "fpu.h"
#include "kprintf.h"

int _fpu_exists;

int fpu_status_read() {
  return (int) ((r_sstatus() >> 13) & 3);
}

void fpu_status_write(int status) {
  status = status & 3;
  uint64 mask = ~(3ULL << 13);
  w_sstatus((r_sstatus() & mask) | (status << 13));
}

void fpu_init() {
  int s1 = fpu_status_read();
  printf("FPU status is %d\n", s1);
  if (s1 != 0) {
    printf("ERROR: FPU should be in state 0 on startup. Will not initialise FPU.\n");
    return;
  }
  fpu_status_write(1);
  int s2 = fpu_status_read();
  printf("FPU status is %d\n", s2);
  if (s2 != 1) {
    printf("FPU is not present or not initialisable.\n");
    _fpu_exists = 0;
    return;
  } else {
    _fpu_exists = 1;
    printf("FPU appears to be present and now initialised!\n");
  }
  
  fpu_status_write(0);
  int s3 = fpu_status_read();
  printf("FPU status is %d\n", s3);
}

void fpu_setupinitial() {
  if (_fpu_exists) {
    fpu_status_write(3);
  }
}

int fpu_instrsizeattrap(struct proc* p) {
  uint32 instr;
  if (copyin(p->pagetable, (char*)&instr, p->trapframe->epc, 4) == -1) {
    printf("fpu_instrtrap(): invalid instruction address at %p, killing process\n", (void*)p->trapframe->epc);
    setkilled(p);
    return 0;
  }
  if ((instr & 3) == 3) {
    return 4;
  } else {
    return 2;
  }
}

void fpu_instrtrap(struct proc* p) {
  uint32 instr;
  if (copyin(p->pagetable, (char*)&instr, p->trapframe->epc, 4) == -1) {
    printf("fpu_instrtrap(): invalid instruction address at %p, killing process\n", (void*)p->trapframe->epc);
    setkilled(p);
    return;
  }
  
  uint32 opcode = (instr & 0x7F);
  
  if (fpu_status_read() == 0 && (opcode == 0x53 || opcode == 0x07 || opcode == 0x27 || opcode == 0x43 || opcode == 0x47 || opcode == 0x4B || opcode == 0x4F)) {
    printf("fpu_instrtrap(): floating point instruction at %p, TODO\n", (void*)p->trapframe->epc);
    if (p->fpu_saved) {
      fpu_status_write(1); // Clean state but enabled
      fpu_restore(&p->fpu_context);
      fpu_status_write(1); // Clean state but enabled
    } else {
      fpu_setupinitial();
    }
    p->fpu_active = 1;
  } else {
    printf("fpu_instrtrap(): invalid instruction %p at address %p, killing process\n", (void*)(uint64)instr, (void*)p->trapframe->epc);
    setkilled(p);
  }
}