#include "memory.h"
#include <cstring>
#include <cstdio>
#include <elf.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>
#include <vector>

Memory::Memory() : base_addr(MEM_BASE), entry_point(MEM_BASE), tohost_addr(TOHOST_ADDR) {
    mem.resize(MEM_SIZE, 0);
}

uint64_t Memory::to_offset(uint64_t addr) {
    if (addr < base_addr || addr >= base_addr + MEM_SIZE) {
        fprintf(stderr, "Memory access out of bounds: 0x%lx\n", addr);
        return 0;
    }
    return addr - base_addr;
}

uint8_t Memory::read8(uint64_t addr) {
    return mem[to_offset(addr)];
}

uint16_t Memory::read16(uint64_t addr) {
    uint64_t off = to_offset(addr);
    return mem[off] | (mem[off + 1] << 8);
}

uint32_t Memory::read32(uint64_t addr) {
    uint64_t off = to_offset(addr);
    return mem[off] | (mem[off + 1] << 8) | (mem[off + 2] << 16) | (mem[off + 3] << 24);
}

uint64_t Memory::read64(uint64_t addr) {
    uint64_t off = to_offset(addr);
    uint64_t lo = read32(addr);
    uint64_t hi = read32(addr + 4);
    return lo | (hi << 32);
}

void Memory::write8(uint64_t addr, uint8_t data) {
    mem[to_offset(addr)] = data;
}

void Memory::write16(uint64_t addr, uint16_t data) {
    uint64_t off = to_offset(addr);
    mem[off] = data & 0xFF;
    mem[off + 1] = (data >> 8) & 0xFF;
}

void Memory::write32(uint64_t addr, uint32_t data) {
    uint64_t off = to_offset(addr);
    mem[off] = data & 0xFF;
    mem[off + 1] = (data >> 8) & 0xFF;
    mem[off + 2] = (data >> 16) & 0xFF;
    mem[off + 3] = (data >> 24) & 0xFF;
}

void Memory::write64(uint64_t addr, uint64_t data) {
    write32(addr, data & 0xFFFFFFFF);
    write32(addr + 4, (data >> 32) & 0xFFFFFFFF);
}

bool Memory::load_elf(const std::string& filename) {
    int fd = open(filename.c_str(), O_RDONLY);
    if (fd < 0) {
        fprintf(stderr, "Failed to open ELF file: %s\n", filename.c_str());
        return false;
    }

    Elf64_Ehdr ehdr;
    if (read(fd, &ehdr, sizeof(ehdr)) != sizeof(ehdr)) {
        fprintf(stderr, "Failed to read ELF header\n");
        close(fd);
        return false;
    }

    if (memcmp(ehdr.e_ident, ELFMAG, SELFMAG) != 0) {
        fprintf(stderr, "Not a valid ELF file\n");
        close(fd);
        return false;
    }

    entry_point = ehdr.e_entry;

    // Load program headers
    for (int i = 0; i < ehdr.e_phnum; i++) {
        Elf64_Phdr phdr;
        lseek(fd, ehdr.e_phoff + i * sizeof(phdr), SEEK_SET);
        if (read(fd, &phdr, sizeof(phdr)) != sizeof(phdr)) {
            fprintf(stderr, "Failed to read program header %d\n", i);
            continue;
        }

        if (phdr.p_type == PT_LOAD) {
            uint64_t offset = to_offset(phdr.p_paddr);
            lseek(fd, phdr.p_offset, SEEK_SET);
            if (read(fd, &mem[offset], phdr.p_filesz) != (ssize_t)phdr.p_filesz) {
                fprintf(stderr, "Failed to load segment %d\n", i);
                close(fd);
                return false;
            }
            printf("Loaded segment: paddr=0x%lx size=%lu\n", phdr.p_paddr, phdr.p_filesz);
        }
    }

    if (ehdr.e_shoff != 0 && ehdr.e_shnum != 0) {
        std::vector<Elf64_Shdr> shdrs(ehdr.e_shnum);
        lseek(fd, ehdr.e_shoff, SEEK_SET);
        if (read(fd, shdrs.data(), ehdr.e_shnum * sizeof(Elf64_Shdr)) == (ssize_t)(ehdr.e_shnum * sizeof(Elf64_Shdr))) {
            for (int i = 0; i < ehdr.e_shnum; i++) {
                const Elf64_Shdr& symtab = shdrs[i];
                if (symtab.sh_type != SHT_SYMTAB || symtab.sh_entsize == 0 ||
                    symtab.sh_link >= ehdr.e_shnum) {
                    continue;
                }

                const Elf64_Shdr& strtab = shdrs[symtab.sh_link];
                std::vector<char> strings(strtab.sh_size);
                lseek(fd, strtab.sh_offset, SEEK_SET);
                if (read(fd, strings.data(), strtab.sh_size) != (ssize_t)strtab.sh_size) {
                    continue;
                }

                size_t nsyms = symtab.sh_size / symtab.sh_entsize;
                std::vector<Elf64_Sym> symbols(nsyms);
                lseek(fd, symtab.sh_offset, SEEK_SET);
                if (read(fd, symbols.data(), symtab.sh_size) != (ssize_t)symtab.sh_size) {
                    continue;
                }

                for (const Elf64_Sym& sym : symbols) {
                    if (sym.st_name >= strings.size()) {
                        continue;
                    }
                    if (strcmp(&strings[sym.st_name], "tohost") == 0) {
                        tohost_addr = sym.st_value;
                        break;
                    }
                }
            }
        }
    }

    close(fd);
    printf("ELF loaded: entry=0x%lx\n", entry_point);
    return true;
}