#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <exec/mem.h>
#include <exec/debcon.h>
#include <exec/library.h>
#include <exec/elf.h>

/* From https://chromium.googlesource.com/native_client/src/native_client/+/branch_heads/2311/src/include */

#define EI_NIDENT       16   /* fwd, see rest of EI_* below */
#define ET_NONE         0   /* no file type */
#define ET_REL          1   /* relocatable file */
#define ET_EXEC         2   /* executable file */
#define ET_DYN          3   /* shared object file */
#define ET_CORE         4   /* core file */
/* TODO(karl) figure out effect of adding ET_LOOS through ET_HIOS */
#define ET_LOOS    0xfe00   /* Environment-specific */
#define ET_HIOS    0xfeff   /* Environment-specific */
#define ET_LOPROC  0xff00   /* processor-specific */
#define ET_HIPROC  0xffff   /* processor-specific */
#define EM_NONE         0   /* no machine */
#define EM_M32          1   /* at&t we 32100 */
#define EM_SPARC        2   /* sparc */
#define EM_386          3   /* intel architecture */
#define EM_68K          4   /* motorola 68000 */
#define EM_88K          5   /* motorola 88000 */
#define EM_860          7   /* intel 80860 */
#define EM_MIPS         8   /* mips rs3000 */
#define EM_MIPS_RS4_BE  10  /* mips rs4000 big-endian */
#define EM_LORESERVED   11
#define EM_HIRESERVED   16
#define EM_ARM          40  /* arm */
#define EM_X86_64       62  /* x86-64 */
#define EV_NONE         0   /* invalid version */
#define EV_CURRENT      1   /* current version */
#define EI_MAG0         0   /* file identification */
#define EI_MAG1         1   /* file identification */
#define EI_MAG2         2   /* file identification */
#define EI_MAG3         3   /* file identification */
#define EI_CLASS        4   /* file class */
#define EI_DATA         5   /* data encoding */
#define EI_VERSION      6   /* file version */

/*
 * EI_PAD deviates from the pdf specification, where its value is 7, since
 * EI_OSABI and EI_ABIVERSION have been introduced.  EI_OSABI and
 * EI_OSABIVERSION are from linux elf.h for code usage compatibility.
 * Also, for Elf 64, the value for EI_PAD is also 9.
 */
#define EI_PAD          9   /* start of padding bytes */
#define EI_OSABI        7
#define EI_ABIVERSION   8
/*
 * ELFMAG and SELFMAG are names/values from linux elf.h, for code usage
 * compatibility.
 */
#define ELFMAG          "\177ELF"
#define SELFMAG         4
/* EI_CLASS values */
#define ELFCLASSNONE    0
#define ELFCLASS32      1
#define ELFCLASS64      2
/* EI_DATA values */
#define ELFDATANONE     0
#define ELFDATA2LSB     1
#define ELFDATA2MSB     2
#define PT_NULL       0           /* Unused entry */
#define PT_LOAD       1           /* Loadable segment */
#define PT_DYNAMIC    2           /* Dynamic linking tables */
#define PT_INTERP     3           /* Program interpreter path name */
#define PT_NOTE       4           /* Note section */
#define PT_SHLIB      5           /* Reserved */
#define PT_PHDR       6           /* Program header table */
#define PT_LOOS       0x60000000  /* Environment-specific low */
#define PT_HIOS       0x6fffffff  /* Environment-specific high */
#define PT_LOPROC     0x70000000  /* Processor-specific low */
#define PT_HIPROC     0x7fffffff  /* Processor-specific high */
/*
 * These are from linux elf.h, for code usage
 * compatibility.
 */
#define PT_TLS        7
#define PT_GNU_STACK  0x6474e551
#define PT_GNU_EH_FRAME 0x6474e550      /* GCC .eh_frame_hdr segment */
#define PT_GNU_RELRO    0x6474e552      /* Read-only after relocation */
#define PF_X          1
#define PF_W          2
#define PF_R          4
/*
 * PF_MASKOS is from linux elf.h, for code usage compatibility
 */
#define PF_MASKOS     0x0ff00000  /* os specific */
#define SHF_WRITE       0x1         /* Has writable data */
#define SHF_ALLOC       0x2         /* Allocated in memory image of program */
#define SHF_EXECINSTR   0x4         /* Contains executable instructions */
#define SHF_MASKOS      0x0f000000  /* Environment-specific use */
#define SHF_MASKPROC    0xf0000000  /* Processor-specific use */
#define DT_NULL         0
#define DT_REL          17
#define DT_RELSZ        18
#define ELF_NOTE_GNU    "GNU"
/* n_type value for build ID notes generated by "ld --build-id". */
#define NT_GNU_BUILD_ID 3
#define R_386_GLOB_DAT  6
#define R_386_RELATIVE  8
#define R_ARM_GLOB_DAT  21
#define R_ARM_RELATIVE  23

/* Define 32-bit specific types */
typedef uint32_t    Elf32_Addr;   /* alignment 4 */
typedef uint16_t    Elf32_Half;   /* alignment 2 */
typedef uint32_t    Elf32_Off;    /* alignment 4 */
typedef int32_t     Elf32_Sword;  /* alignment 4 */
typedef uint32_t    Elf32_Word;   /* alignment 4 */
/* unsigned char, size 1, alignment 1 */

#pragma pack(push,1)
/* Define the structure of the file header for 32 bits. */
typedef struct {
  unsigned char e_ident[EI_NIDENT];
  Elf32_Half    e_type;
  Elf32_Half    e_machine;
  Elf32_Word    e_version;
  Elf32_Addr    e_entry;
  Elf32_Off     e_phoff;
  Elf32_Off     e_shoff;
  Elf32_Word    e_flags;
  Elf32_Half    e_ehsize;
  Elf32_Half    e_phentsize;
  Elf32_Half    e_phnum;
  Elf32_Half    e_shentsize;
  Elf32_Half    e_shnum;
  Elf32_Half    e_shstrndx;
} Elf32_Ehdr;

/* Define the structure of a program header table for 32-bits. */

typedef struct {
  Elf32_Word    p_type;
  Elf32_Off     p_offset;
  Elf32_Addr    p_vaddr;
  Elf32_Addr    p_paddr;
  Elf32_Word    p_filesz;
  Elf32_Word    p_memsz;
  Elf32_Word    p_flags;
  Elf32_Word    p_align;
} Elf32_Phdr;

/*
 * Define 32-bit section headers.
 * ncfileutil wants section headers, even though service runtime does
 * not.
 */
typedef struct {
  Elf32_Word  sh_name;
  Elf32_Word  sh_type;
  Elf32_Word  sh_flags;
  Elf32_Addr  sh_addr;
  Elf32_Off   sh_offset;
  Elf32_Word  sh_size;
  Elf32_Word  sh_link;
  Elf32_Word  sh_info;
  Elf32_Word  sh_addralign;
  Elf32_Word  sh_entsize;
} Elf32_Shdr;

typedef struct {
  Elf32_Addr r_offset;
  Elf32_Word r_info;
} Elf32_Rel;

#define ELF32_R_TYPE(val) ((val) & 0xff)

typedef struct {
  Elf32_Sword d_tag;
  union {
    Elf32_Word d_val;
    Elf32_Addr d_ptr;
  } d_un;
} Elf32_Dyn;

typedef struct {
  Elf32_Word n_namesz;
  Elf32_Word n_descsz;
  Elf32_Word n_type;
} Elf32_Nhdr;
#pragma pack(pop)


/* Simple ELF loader for 32-bit i386 images in memory.
 * - loads PT_LOAD segments into freshly allocated memory
 * - supports ET_DYN (shared object) by relocating the image to a new base
 * - applies relocations found via the PT_DYNAMIC table (DT_REL, DT_RELSZ)
 * - handles R_386_RELATIVE relocations
 *
 * This loader is intentionally small and only implements what is needed
 * for in-memory device/shared-object loading in this project.
 */

#define ALIGN_DOWN(x, a) ((x) & ~((a)-1))
#define ALIGN_UP(x, a) (((x) + (a)-1) & ~((a)-1))

uintptr_t ElfLoad(struct Library *lib, void *elf_image) {
  if (!elf_image) return 0;

  uint8_t *base = (uint8_t *)elf_image;
  Elf32_Ehdr *eh = (Elf32_Ehdr *)base;

  /* Basic validation: magic, class, data */
  if (memcmp(eh->e_ident, ELFMAG, SELFMAG) != 0) return 0;
  if (eh->e_ident[EI_CLASS] != ELFCLASS32) return 0;
  if (eh->e_ident[EI_DATA] != ELFDATA2LSB) return 0;

  /* Only handle i386 for now */
  if (eh->e_machine != EM_386) return 0;

  Elf32_Phdr *ph = (Elf32_Phdr *)(base + eh->e_phoff);

  /* Determine address range of loadable segments (virtual addresses).
   * For ET_DYN we'll treat p_vaddr as relative offsets and pick a
   * new base address. For ET_EXEC we'd ideally load at the linked
   * addresses; here we support ET_EXEC only if its p_vaddrs are
   * usable after allocation (best-effort).
   */
  uint32_t min_vaddr = UINT32_MAX;
  uint32_t max_vaddr = 0;
  for (int i = 0; i < eh->e_phnum; i++) {
    if (ph[i].p_type != PT_LOAD) continue;
    if (ph[i].p_memsz == 0) continue;
    if (ph[i].p_vaddr < min_vaddr) min_vaddr = ph[i].p_vaddr;
    uint32_t end = ph[i].p_vaddr + ph[i].p_memsz;
    if (end > max_vaddr) max_vaddr = end;
  }

  if (min_vaddr == UINT32_MAX) return 0; /* nothing to load */

  uint32_t span = max_vaddr - min_vaddr;
  /* align to page-size like boundary (4k) */
  uint32_t alloc_size = ALIGN_UP(span, 0x1000);

  void *load_base = AllocMem(alloc_size);
  if (!load_base) return 0;

  uint8_t *img_base = (uint8_t *)load_base;
  /* zero entire allocation to simplify bss handling */
  memset(img_base, 0, alloc_size);

  /* Copy PT_LOAD segments into the allocated image */
  for (int i = 0; i < eh->e_phnum; i++) {
    Elf32_Phdr *p = &ph[i];
    if (p->p_type != PT_LOAD) continue;
    uint32_t dest_off = p->p_vaddr - min_vaddr;
    uint8_t *dest = img_base + dest_off;
    if (p->p_filesz) {
      /* bounds check in source image */
      if ((uint32_t)p->p_offset + p->p_filesz > (uint32_t)eh->e_shoff && eh->e_shoff != 0) {
        /* best-effort: but don't fail here if file layout unknown */
      }
      memcpy(dest, base + p->p_offset, p->p_filesz);
    }
    if (p->p_memsz > p->p_filesz) {
      memset(dest + p->p_filesz, 0, p->p_memsz - p->p_filesz);
    }
  }

  /* Walk program headers to find PT_DYNAMIC (if any) and apply relocations */
  Elf32_Dyn *dynamic = NULL;
  uint32_t dynamic_size = 0;
  for (int i = 0; i < eh->e_phnum; i++) {
    Elf32_Phdr *p = &ph[i];
    if (p->p_type == PT_DYNAMIC) {
      /* dynamic table lives inside loaded memory at p_vaddr */
      uint32_t off = p->p_vaddr - min_vaddr;
      dynamic = (Elf32_Dyn *)(img_base + off);
      dynamic_size = p->p_memsz;
      break;
    }
  }

  uintptr_t entry = 0;
  if (eh->e_type == ET_DYN) {
    entry = (uintptr_t)img_base + (eh->e_entry - min_vaddr);
  } else if (eh->e_type == ET_EXEC) {
    /* For ET_EXEC assume entry is absolute in header */
    entry = (uintptr_t)img_base + (eh->e_entry - min_vaddr);
  } else {
    /* unsupported type */
    return 0;
  }

  /* Apply relocations from dynamic segment (DT_REL / DT_RELSZ) */
  if (dynamic) {
    uint32_t rel_addr = 0;
    uint32_t rel_sz = 0;

    for (Elf32_Dyn *d = dynamic; d->d_tag != DT_NULL; d++) {
      switch (d->d_tag) {
        case DT_REL:
          rel_addr = d->d_un.d_ptr;
          break;
        case DT_RELSZ:
          rel_sz = d->d_un.d_val;
          break;
        default:
          break;
      }
    }

    if (rel_addr && rel_sz) {
      /* rel_addr is a virtual address; convert into our load image */
      uint32_t rel_off = rel_addr - min_vaddr;
      Elf32_Rel *rels = (Elf32_Rel *)(img_base + rel_off);
      uint32_t rel_count = rel_sz / sizeof(Elf32_Rel);
      for (uint32_t i = 0; i < rel_count; i++) {
        Elf32_Rel *r = &rels[i];
        uint32_t type = ELF32_R_TYPE(r->r_info);
        uint32_t where_off = r->r_offset - min_vaddr;
        Elf32_Addr *where = (Elf32_Addr *)(img_base + where_off);
        switch (type) {
          case R_386_RELATIVE:
            /* *where = base + addend; For REL entries addend is the
             * original memory value at the location.
             */
            *where = (Elf32_Addr)((uintptr_t)img_base + (*where));
            break;
          case R_386_GLOB_DAT:
            /* Without a symbol resolver we cannot set external
             * addresses. Leave as-is (best-effort) or attempt to
             * treat it like RELATIVE if it appears to be relative.
             */
            /* no-op */
            break;
          default:
            /* unsupported relocation type */
            break;
        }
      }
    }
  }

  lib->entry = entry;
  lib->base = (uintptr_t)img_base;

  return entry;
}