#include "param.h"
#include "types.h"
#include "defs.h"
#include "x86.h"
#include "memlayout.h"
#include "mmu.h"
#include "proc.h"

int mappages(pml4e_t *pgdir, void *va, addr_t size, addr_t pa, int perm);
static pte_t * walkpgdir(pml4e_t *pml4, const void *va, int alloc);


// `cpu` and `proc` are no longer __thread variables; they live at fixed offsets
// in the per-CPU block and are read through %gs at each use. See the long
// comment in include/proc.h for why -- the short version is that __thread means
// per *thread*, the compiler caches the thread pointer across calls on that
// basis, and a kernel thread changes CPU at every sched().

// Kernel FS_BASE set by seginit(); read by trapasm.S to restore after user TLS
uint64 kernel_fs_base;
// Single-CPU scratch for syscall_entry FS swap (safe while IF cleared by SYSCALL/SFMASK)
uint64 _sc_rax, _sc_rcx, _sc_rdx;

static pml4e_t *kpml4;
static pdpe_t *kpdpt;

struct tassdesc {
  uint lim_15_0 : 16;  
  uint base_15_0 : 16; 
  uint base_23_16 : 8; 
  uint type : 4;       
  uint s : 1;          
  uint dpl : 2;        
  uint p : 1;          
  uint lim_19_16 : 4;  
  uint avl : 1;        
  uint l : 1;          
  uint d : 1;          
  uint g : 1;          
  uint base_31_24 : 8; 
  uint32 base_63_32;   
  uint32 res;          
};

void
syscallinit(void)
{
  // SYSRETQ does NOT use USER_CS/USER_DS directly. Per the AMD64/Intel
  // convention, it computes CS = STAR[63:48]+16 and SS = STAR[63:48]+8
  // (RPL forced to 3), so STAR[63:48] must hold the *compat* 32-bit user
  // code selector (SEG_UCODE32), relying on the GDT's consecutive
  // [UCODE32, UDATA, UCODE] layout — NOT USER_CS itself. Using USER_CS
  // here made every sysretq load CS=USER_CS+16 (an invalid selector,
  // pointing past UCODE into the TSS descriptor) and SS=USER_CS+8. QEMU's
  // TCG apparently doesn't validate segment descriptors on SYSRET, so this
  // went unnoticed until an interrupt forced a strictly-validated iretq
  // with that same bogus CS still loaded, which correctly #GP'd.
  wrmsr(MSR_STAR,
    ((((uint64)((SEG_UCODE32 << 3) | DPL_USER) << 16) | (uint64)KERNEL_CS) << 32));
  wrmsr(MSR_LSTAR, (addr_t)syscall_entry);
  wrmsr(MSR_CSTAR, (addr_t)ignore_sysret);
  wrmsr(MSR_SFMASK, FL_TF|FL_DF|FL_IF|FL_IOPL_MASK|FL_AC|FL_NT);
}

void
seginit(void)
{
  struct segdesc *gdt;
  uint *tss;
  uint64 addr;
  void *local;
  struct cpu *c;

  // The per-CPU page, preallocated by the boot processor for every AP.
  //
  // It must not be kalloc()'d here on an application processor. kalloc takes a
  // spinlock; acquire() calls holding(); holding() reads `cpu`, which is a
  // __thread variable reached through %fs -- and %fs is set up four lines
  // below. So an AP allocating its own page reads the per-CPU pointer before
  // the per-CPU pointer exists, and the kernel dies in panic("acquire") with
  // every lock appearing to be held.
  //
  // startothers() fills c->local in before starting each AP, on the boot
  // processor, where locking works.
  c = &cpus[cpunum()];
  local = c->local;
  if (local == 0) {
    local = kalloc();          // the boot processor's own, taken normally
    memset(local, 0, PGSIZE);
  }

  gdt = (struct segdesc*) local;
  tss = (uint*) (((char*) local) + 1024);
  tss[16] = 0x00680000;

  kernel_fs_base = ((uint64) local) + 2048;
  wrmsr(0xC0000100, kernel_fs_base);

  // The TCB self-pointer: the word at %fs:0 must hold the thread pointer.
  //
  // This is part of the x86-64 TLS ABI and it was missing, which worked only
  // by accident. GCC with -mtls-direct-seg-refs compiles `cpu = c` into a
  // single `movq %rcx, %fs:-0x10` and never needs the thread pointer as a
  // value, so nothing ever read this word. clang materialises it instead --
  // `movq %fs:0x0, %rax; addq $-0x10, %rax; movq %rcx, (%rax)` -- which read
  // the zero left by the memset above and stored through address
  // 0xfffffffffffffff0, three instructions later.
  //
  // Writing it costs one store per CPU and removes the dependency on a
  // particular compiler's TLS optimisation.
  *(uint64 *) kernel_fs_base = kernel_fs_base;

  // Clear the two per-CPU slots before anything can read them. Their offsets
  // are PERCPU_CPU and PERCPU_PROC in include/proc.h and are shared with
  // trapasm.S, which reads %gs:16 for proc on the way out to user mode.
  ((uint64 *) kernel_fs_base)[2] = 0;      // struct cpu *   (PERCPU_CPU)
  ((uint64 *) kernel_fs_base)[3] = 0;      // struct proc *  (PERCPU_PROC)

  // The same address in GS base, per CPU.
  //
  // trapasm.S needs to find *this* CPU's kernel FS base on every entry from
  // user mode, and it has no register to work from at that instant -- which is
  // why it used to read the kernel_fs_base global. That global is written by
  // every CPU during its own seginit(), so with more than one processor the
  // last writer wins and every CPU restores one CPU's per-CPU block. GS base is
  // per-CPU in hardware, so %gs:0 answers the question with no global at all.
  //
  // Safe because GS belongs entirely to the kernel here: sys_arch_prctl()
  // implements ARCH_SET_FS and ARCH_GET_FS only, and musl uses FS for TLS on
  // x86-64. See the comment on RESTORE_KERNEL_FS.
  wrmsr(0xC0000101, kernel_fs_base);      // IA32_GS_BASE

  c->local = local;

  cpu_set(c);
  proc_set(0);

  // Two of the entries below are deliberately null descriptors, and neither is
  // dead code even though both look it:
  //
  //   SEG_UCODE32  is never loaded, but SYSRETQ computes CS from STAR[63:48]+16
  //                and SS from +8, so the *slot* must exist to keep the
  //                [UCODE32, UDATA, UCODE] triple consecutive. See syscallinit.
  //   SEG_KCPU     is a leftover slot number from 32-bit xv6, where %gs held the
  //                per-CPU pointer. SchOS uses FS_BASE/TLS instead, but the
  //                index has to stay occupied or SEG_TSS stops landing at 7.
  //
  // Deleting either moves every descriptor after it and breaks sysret.
  addr = (uint64) tss;
  gdt[0] = (struct segdesc) {0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
  gdt[SEG_KCODE] = SEG(STA_X|STA_R, 0, 0, 0);
  gdt[SEG_KDATA] = SEG(STA_W, 0, 0, 0);
  gdt[SEG_UCODE32] = (struct segdesc) {0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
  gdt[SEG_UDATA] = SEG(STA_W, 0, 0, DPL_USER);
  gdt[SEG_UCODE] = SEG(STA_X|STA_R, 0, 0, DPL_USER);
  gdt[SEG_KCPU]  = (struct segdesc) {0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
  
  struct tassdesc *td = (struct tassdesc*)(gdt + SEG_TSS);
  td->lim_15_0 = 0x67;
  td->base_15_0 = addr & 0xffff;
  td->base_23_16 = (addr >> 16) & 0xff;
  td->type = STS_T64A;
  td->s = 0;
  td->dpl = 0;
  td->p = 1;
  td->lim_19_16 = 0;
  td->avl = 0;
  td->l = 0;
  td->d = 0;
  td->g = 0;
  td->base_31_24 = (addr >> 24) & 0xff;
  td->base_63_32 = (addr >> 32) & 0xffffffff;
  td->res = 0;

  lgdt((void*) gdt, (NSEGS+1) * sizeof(struct segdesc));
  ltr(SEG_TSS << 3);
}

pml4e_t*
setupkvm(void)
{
  pml4e_t *pml4 = (pml4e_t*) kalloc();
  if(pml4 == 0) return 0;
  memset(pml4, 0, PGSIZE);

  pml4[PMX(KERNBASE)] = v2p(kpdpt) | PTE_P | PTE_W;

  // The framebuffer is deliberately NOT mapped here.
  //
  // This used to identity-map 0xfd000000 for 800*600*4 bytes, user-writable,
  // into every process that ever existed. Three things were wrong with that.
  // The address and the size were both guesses that only hold on QEMU's
  // stdvga at one resolution; setupkvm runs before vga_init, so on the first
  // call nothing had even read the VBE block yet; and every process could
  // scribble on the screen without asking anyone.
  //
  // Nothing used it -- userspace reaches the framebuffer through
  // mmap("/dev/fb0"), which maps the real base at the real size. So it is
  // gone, and a process gets the framebuffer exactly when it asks for it.

  return pml4;
}

void
kvmalloc(void)
{
  kpml4 = (pml4e_t*) kalloc();
  memset(kpml4, 0, PGSIZE);

  kpdpt = (pdpe_t*) kalloc();
  memset(kpdpt, 0, PGSIZE);

  kpml4[0] = v2p(kpdpt) | PTE_P | PTE_W;
  kpml4[PMX(KERNBASE)] = v2p(kpdpt) | PTE_P | PTE_W;
  
  for(uint64 i = 0; i < 4; i++) {
    uint64 flags = PTE_PS | PTE_P | PTE_W;
    if(i >= 3) {
      flags |= PTE_PCD | PTE_PWT;
    }
    kpdpt[i] = (i << 30) | flags;
  }

  switchkvm();
}

void
switchuvm(struct proc *p)
{
  pushcli();
  if(p->pgdir == 0)
    panic("switchuvm: no pgdir");
  uint *tss = (uint*) (((char*) mycpu()->local) + 1024);
  const addr_t stktop = (addr_t)p->kstack + KSTACKSIZE;
  tss[1] = (uint)stktop;
  tss[2] = (uint)(stktop >> 32);
  lcr3(v2p(p->pgdir));
  popcli();
}

static pte_t *
walkpgdir(pml4e_t *pml4, const void *va, int alloc)
{
  pml4e_t *pml4e;
  pdpe_t *pdp, *pdpe;
  pde_t *pd, *pde;
  pte_t *pgtab;

  pml4e = &pml4[PMX(va)];
  if(*pml4e & PTE_P)
    pdp = (pdpe_t*)P2V(PTE_ADDR(*pml4e));
  else {
    if(!alloc || (pdp = (pdpe_t*)kalloc()) == 0)
      return 0;
    memset(pdp, 0, PGSIZE);
    *pml4e = V2P(pdp) | PTE_P | PTE_W | PTE_U;
  }

  pdpe = &pdp[PDPX(va)];
  if(*pdpe & PTE_P)
    pd = (pde_t*)P2V(PTE_ADDR(*pdpe));
  else {
    if(!alloc || (pd = (pde_t*)kalloc()) == 0)
      return 0;
    memset(pd, 0, PGSIZE);
    *pdpe = V2P(pd) | PTE_P | PTE_W | PTE_U;
  }

  pde = &pd[PDX(va)];
  if(*pde & PTE_P)
    pgtab = (pte_t*)P2V(PTE_ADDR(*pde));
  else {
    if(!alloc || (pgtab = (pte_t*)kalloc()) == 0)
      return 0;
    memset(pgtab, 0, PGSIZE);
    *pde = V2P(pgtab) | PTE_P | PTE_W | PTE_U;
  }

  return (pte_t*)&pgtab[PTX(va)];
}

void
switchkvm(void)
{
  lcr3(v2p(kpml4));
}

int
mappages(pml4e_t *pgdir, void *va, addr_t size, addr_t pa, int perm)
{
  char *a, *last;
  pte_t *pte;

  a = (char*)PGROUNDDOWN((addr_t)va);
  last = (char*)PGROUNDDOWN(((addr_t)va) + size - 1);
  for(;;){
    if((pte = walkpgdir(pgdir, a, 1)) == 0)
      return -1;
    if(*pte & PTE_P)
      panic("remap");
    *pte = pa | perm | PTE_P | PTE_U;
    invlpg(a);
    if(a == last)
      break;
    a += PGSIZE;
    pa += PGSIZE;
  }
  return 0;
}

void
vm_unmap_range(pml4e_t *pgdir, uint64 va, uint64 length)
{
  uint64 a   = PGROUNDDOWN(va);
  uint64 end = PGROUNDUP(va + length);
  for (; a < end; a += PGSIZE) {
    pte_t *pte = walkpgdir(pgdir, (void*)a, 0);
    if (pte && (*pte & PTE_P)) {
      uint64 pa = PTE_ADDR(*pte);
      kfree((char*)P2V(pa));
      *pte = 0;
      invlpg((void*)a);
    }
  }
}

// Tear down a mapping without freeing what is underneath it.
//
// vm_unmap_range() kfree()s every page it unmaps, which is right for an
// anonymous mapping and catastrophic for anything shared: the framebuffer is
// device memory that was never allocated, and a shared memory segment's pages
// belong to the segment, not to whichever process happens to be detaching.
// Handing either of those to kfree puts them on the free list while another
// process is still reading and writing them.
void
vm_unmap_norelease(pml4e_t *pgdir, uint64 va, uint64 length)
{
  uint64 a   = PGROUNDDOWN(va);
  uint64 end = PGROUNDUP(va + length);
  for (; a < end; a += PGSIZE) {
    pte_t *pte = walkpgdir(pgdir, (void*)a, 0);
    if (pte && (*pte & PTE_P)) {
      *pte = 0;
      invlpg((void*)a);
    }
  }
}

// Duplicate an mmap()'d VMA region from src into dst for fork().
// Framebuffer mappings are re-pointed at the same physical MMIO range
// (shared, never copied/freed); anonymous mappings get fresh pages with
// the parent's contents copied in — mirrors what copyuvm() does for the
// [0, sz) heap range, which never covers mmap()'d addresses.
int
copy_vma_range(pml4e_t *dstpgdir, pml4e_t *srcpgdir, uint64 va, uint64 length, int share)
{
  uint64 a   = PGROUNDDOWN(va);
  uint64 end = PGROUNDUP(va + length);
  for (; a < end; a += PGSIZE) {
    if (share) {
      // Point the child at the very same physical pages, whatever they are.
      //
      // This used to special-case the framebuffer and map vga_fb_phys()
      // explicitly. Reading the parent's own PTE instead covers device memory
      // and shared memory segments with one rule, and cannot disagree with
      // what the parent actually has mapped -- which the explicit version
      // could, and did, when the address was hardcoded.
      pte_t *spte = walkpgdir(srcpgdir, (void*)a, 0);
      if (!spte || !(*spte & PTE_P))
        continue;
      if (mappages(dstpgdir, (void*)a, PGSIZE, PTE_ADDR(*spte), PTE_FLAGS(*spte)) < 0)
        return -1;
      continue;
    }
    pte_t *spte = walkpgdir(srcpgdir, (void*)a, 0);
    if (!spte || !(*spte & PTE_P))
      continue;
    char *mem = kalloc();
    if (!mem)
      return -1;
    memmove(mem, (char*)P2V(PTE_ADDR(*spte)), PGSIZE);
    if (mappages(dstpgdir, (void*)a, PGSIZE, V2P(mem), PTE_FLAGS(*spte)) < 0) {
      kfree(mem);
      return -1;
    }
  }
  return 0;
}

void
inituvm(pde_t *pgdir, char *sinit, uint64 sz)
{
  char *mem;

  if(sz >= PGSIZE)
    panic("inituvm: more than a page");
  mem = kalloc(); 
  memset(mem, 0, PGSIZE);
  mappages(pgdir, 0, PGSIZE, v2p(mem), PTE_W|PTE_U);
  memmove(mem, sinit, sz); 
}

int
loaduvm(pml4e_t *pgdir, char *addr, struct inode *ip, uint64 offset, uint64 sz)
{
  uint64 done, n, off_in_page;
  addr_t pa;
  pte_t *pte;
  char *cur_page;

  for(done = 0; done < sz; done += n){
    cur_page   = (char*)PGROUNDDOWN((uint64)(addr + done));
    off_in_page = (uint64)(addr + done) - (uint64)cur_page;
    if((pte = walkpgdir(pgdir, cur_page, 0)) == 0)
      panic("loaduvm: address should exist");
    pa = PTE_ADDR(*pte);
    n = PGSIZE - off_in_page;
    if(n > sz - done)
      n = sz - done;
    if(readi(ip, 0, (uint64)(P2V(pa) + off_in_page), offset + done, n) != (int)n)
      return -1;
  }
  return 0;
}

uint64
allocuvm(pml4e_t *pgdir, uint64 oldsz, uint64 newsz)
{
  char *mem;
  addr_t a;

  if(newsz >= KERNBASE)
    return 0;
  if(newsz < oldsz)
    return oldsz;

  a = PGROUNDUP(oldsz);
  for(; a < newsz; a += PGSIZE){
    mem = kalloc();
    if(mem == 0){
      deallocuvm(pgdir, newsz, oldsz);
      return 0;
    }
    memset(mem, 0, PGSIZE);
    if(mappages(pgdir, (char*)a, PGSIZE, V2P(mem), PTE_W|PTE_U) < 0){
      deallocuvm(pgdir, newsz, oldsz);
      kfree(mem);
      return 0;
    }
  }
  return newsz;
}

uint64
deallocuvm(pml4e_t *pgdir, uint64 oldsz, uint64 newsz)
{
  pte_t *pte;
  addr_t a, pa;

  if(newsz >= oldsz)
    return oldsz;

  a = PGROUNDUP(newsz);
  for(; a  < oldsz; a += PGSIZE){
    pte = walkpgdir(pgdir, (char*)a, 0);
    if(pte && (*pte & PTE_P) != 0){
      pa = PTE_ADDR(*pte);
      if(pa == 0)
        panic("kfree");
      if(pa < PHYSTOP) {
        char *v = P2V(pa);
        kfree(v);
      }
      *pte = 0;
    }
  }
  return newsz;
}

void
freevm(pml4e_t *pml4)
{
  uint i, j, k, l;
  pdpe_t *pdp;
  pde_t *pd;
  pte_t *pt;

  if(pml4 == 0)
    panic("freevm: no pgdir");

  for(i = 0; i < (512/2); i++){
    if(pml4[i] & PTE_P){
      pdp = (pdpe_t*)P2V(PTE_ADDR(pml4[i]));
      for(j = 0; j < 512; j++){
        if(pdp[j] & PTE_P){
          pd = (pde_t*)P2V(PTE_ADDR(pdp[j]));
          for(k = 0; k < 512; k++){
            if(pd[k] & PTE_P) {
              pt = (pte_t*)P2V(PTE_ADDR(pd[k]));
              for(l = 0; l < 512; l++){
                if(pt[l] & PTE_P) {
                  addr_t pa = PTE_ADDR(pt[l]);
                  if(pa < PHYSTOP && pa != 0) {
                    kfree(P2V(pa));
                  }
                }
              }
              kfree((char*)pt);
            }
          }
          kfree((char*)pd);
        }
      }
      kfree((char*)pdp);
    }
  }
  kfree((char*)pml4);
}

void
clearpteu(pml4e_t *pgdir, char *uva)
{
  pte_t *pte;
  pte = walkpgdir(pgdir, uva, 0);
  if(pte == 0)
    panic("clearpteu");
  *pte &= ~PTE_U;
}

pml4e_t*
copyuvm(pml4e_t *pgdir, uint64 sz)
{
  pml4e_t *d;
  pte_t *pte;
  addr_t pa, i, flags;
  char *mem;

  if((d = setupkvm()) == 0)
    return 0;
  for(i = 0; i < sz; i += PGSIZE){
    if((pte = walkpgdir(pgdir, (void*)i, 0)) == 0)
      continue;
    if(!(*pte & PTE_P))
      continue;
    pa = PTE_ADDR(*pte);
    flags = PTE_FLAGS(*pte);

    if(pa < PHYSTOP) {
      if((mem = kalloc()) == 0)
        goto bad;
      memmove(mem, (char*)P2V(pa), PGSIZE);
      if(mappages(d, (void*)i, PGSIZE, V2P(mem), flags) < 0) {
        kfree(mem);
        goto bad;
      }
    }
  }
  return d;

  bad:
  freevm(d);
  return 0;
}

char*
uva2ka(pml4e_t *pgdir, char *uva)
{
  if((addr_t)uva >= KERNBASE)
    return 0;

  pte_t *pte;
  pte = walkpgdir(pgdir, uva, 0);
  if(pte == 0 || (*pte & PTE_P) == 0)
    return 0;
  if((*pte & PTE_U) == 0)
    return 0;
  addr_t pa = (*pte) & 0x000FFFFFFFFFF000ull;
  // Include uva's in-page offset — callers dereference the returned pointer
  // directly at uva's address, not just somewhere within its page.
  return (char*)P2V(pa) + ((addr_t)uva & 0xFFF);
}

int
copyout(pml4e_t *pgdir, addr_t va, void *p, uint64 len)
{
  char *buf, *pa0;
  addr_t n, va0;

  buf = (char*)p;
  while(len > 0){
    va0 = PGROUNDDOWN(va);
    pa0 = uva2ka(pgdir, (char*)va0);

    if(pa0 == 0) {
      return -1;
    }

    n = PGSIZE - (va - va0);
    if(n > len)
      n = len;
    memmove(pa0 + (va - va0), buf, n);

    len -= n;
    buf += n;
    va += n;
  }
  return 0;
}

int
copyin(pml4e_t *pgdir, addr_t srcva, void *dst, uint64 len)
{
  uint64 n, va0;
  char *pa0;
  char *dst_buf = (char*)dst;

  while(len > 0){
    va0 = PGROUNDDOWN(srcva);
    pa0 = uva2ka(pgdir, (char*)va0);
    if(pa0 == 0)
      return -1;
    n = PGSIZE - (srcva - va0);
    if(n > len)
      n = len;
    memmove(dst_buf, pa0 + (srcva - va0), n);
    len -= n;
    dst_buf += n;
    srcva += n;
  }
  return 0;
}