#include <gelf.h>
#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>

#include "cpu.h"

static inline int process_elf(Elf *elf, CPU *cpu) {
    AddressSpace *addressSpace = cpu->addressSpace;

    if (elf_kind(elf) != ELF_K_ELF) {
        fprintf(stderr, "File is not a valid ELF file\n");
        return 1;
    }

    GElf_Ehdr ehdr;
    if (gelf_getehdr(elf, &ehdr) == NULL) {
        fprintf(stderr, "gelf_getehdr() failed: %s\n", elf_errmsg(-1));
        return 1;
    }

    if (ehdr.e_type != ET_EXEC) {
        printf("ELF is not an executable, proceeding anyway.\n");
    }

    if (ehdr.e_machine != EM_RISCV) {
        fprintf(stderr, "ELF is not for RISC-V. Machine: 0x%X\n", ehdr.e_machine);
        return 1;
    }

    Elf_Scn *scn = NULL;
    GElf_Shdr shdr;

    int warnedArch = 0;
    
    size_t shdrstrndx; // Section HeaDeR STRing table iNDeX
    if (elf_getshdrstrndx(elf, &shdrstrndx) != 0) {
        fprintf(stderr, "elf_getshdrstrndx error: %s\n", elf_errmsg(-1));
        return 1;
    }

    while ((scn = elf_nextscn(elf, scn)) != NULL) {
        if (gelf_getshdr(scn, &shdr) != &shdr) {
            fprintf(stderr, "getshdr() failed: %s\n", elf_errmsg(-1));
            continue;
        }

        char *sectionName = elf_strptr(elf, shdrstrndx, shdr.sh_name);

        if (shdr.sh_type == SHT_RISCV_ATTRIBUTES) {
            Elf_Data *data = elf_getdata(scn, NULL);
            
            if (data == NULL || data->d_size == 0) {
                fprintf(stderr, "elf_getdata() failed: %s\n", elf_errmsg(-1));
                continue;
            }

            // TODO this is highly fragile
            // https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-elf.adoc
            char *content = (char *)data->d_buf;
            content += 19; // first 17 chars are subsection headers, and I have to +2 to make strcmp match
            if (strcmp(content, "rv32i2p1") == 0) {
                warnedArch = 1;
            }

            break;
        } else if (strcmp(sectionName, ".text") == 0) {
            Elf_Data *data = elf_getdata(scn, NULL);
            
            if (data == NULL || data->d_size == 0) {
                fprintf(stderr, "elf_getdata() failed: %s\n", elf_errmsg(-1));
                continue;
            }

            char *content = (char *)data->d_buf;
            uint32_t contentSize = data->d_size;

            if (contentSize > addressSpace->romSize) {
                fprintf(stderr, "Code has %d bytes, and doesn't fit in ROM of capacity %d bytes\n", contentSize, addressSpace->romSize);
                return 1;
            }

            memcpy(addressSpace->rom, content, contentSize);

            printf("Loaded %d bytes or %d instructions from an ELF file\n", contentSize, contentSize / 4);
        } else {
            printf("Unrecognized section: %s\n", sectionName);
        }
    }

    if (warnedArch == 0) {
        printf("I couldn't verify whether the ELF was compiled for the correct instruction set!\n");
    }

    cpu->programCounter = ehdr.e_entry;

    return 0;
}

int load_elf_to_cpu_and_rom(const char filename[], CPU *cpu) {
    if (elf_version(EV_CURRENT) == EV_NONE) {
        fprintf(stderr, "Invalid ELF version: %s\n", elf_errmsg(-1)); // TODO can perror be used?
        return 1;
    }

    int code = 0;

    int fd = open(filename, O_RDONLY, 0);
    if (fd < 0) {
        perror("Failed opening file");
        return 1;
    }

    Elf *elf = elf_begin(fd, ELF_C_READ, NULL);
    if (elf == NULL) {
        fprintf(stderr, "Error reading ELF file: %s\n", elf_errmsg(-1));
        if (close(fd) < 0) {
            perror("Also failed closing file");
        }
        return 1;
    }

    if ((code = process_elf(elf, cpu)) != 0) {
        fprintf(stderr, "Failed processing ELF\n");
    }

    if (elf_end(elf) < 0) {
        fprintf(stderr, "Failed closing ELF: %s\n", elf_errmsg(-1)); // TODO can perror be used?
    }

    if (close(fd) < 0) {
        perror("Failed closing file"); // TODO perhaps make one close for all the errors above
    }

    return code;
}