#include "elf_header.h"
#include "elf_osabi.h"
#include "elf_machine_type.h"
#include <stdio.h>
#include <string.h>

#define ELF_IDENTITIFCATION_MAX_SIZE 16
#define ELF_CLASS_TYPE32 1
#define ELF_CLASS_TYPE64 2

#define ELF_BYTE_ORDER_LSB 1
#define ELF_BYTE_ORDER_MSB 2

#define ELF64_ADDR_SIZE sizeof(uint64_t)
#define ELF64_OFFSET_SIZE ELF64_ADDR_SIZE
#define ELF64_WORD_SIZE ELF64_ADDR_SIZE / 2
#define ELF64_HALF_SIZE ELF64_WORD_SIZE / 2

typedef enum EIdent
{
	EI_MAG0 = 0, //Magic number position zero - three.
	EI_MAG1 = 1,
	EI_MAG2 = 2,
	EI_MAG3 = 3,
	EI_CLASS = 4,
	EI_DATA = 5,
	EI_HEADER_VERSION = 6,
	EI_OSABI = ELF_OSABI_IDENT_INDEX,
	EI_ABIVERSION = ELF_OSABI_VERSION_IDENT_INDEX,
	EI_PAD = 9,
	EI_NIDENT = ELF_IDENTITIFCATION_MAX_SIZE - 1, //End of reserved 16 bytes for the header.
} ELFIdentOffsets;

struct elf_header {
    unsigned char Identitication[ELF_IDENTITIFCATION_MAX_SIZE];
    uint32_t FileType;
    uint32_t ISA;
    uint32_t ELFVersion;
    uint64_t EntryAddress;
	uint64_t ProgramHeaderOffset;
	uint64_t SectionHeaderOffset;
	uint32_t Flags;
	uint16_t HeaderSize;
	uint16_t ProgramHeaderEntrySize;
	uint16_t ProgramHeaderEntryCount;
	uint16_t SectionHeaderEntrySize;
	uint16_t SectionHeaderEntryCount;
	uint16_t SectionNamesIndex;
};

struct elf_header* CreateELFHeader(void);
int ELFMagicNumberValid(const unsigned char* ident);
void ELFPrintMagic(const struct elf_header* header);

struct elf_header* ParseObjectFile(const char* bytes, size_t count) {
    if (count < 0x40) return NULL;

    struct elf_header* header = CreateELFHeader();

    memcpy(header, bytes, ELF_IDENTITIFCATION_MAX_SIZE);

	if (!ELFMagicNumberValid(header->Identitication)) 
	{
		fprintf(stderr, "Failed to find ELF magic number in file header.\n");
        ELFPrintMagic(header);
        printf("\n");

		return NULL;
	}

	uint64_t offsetSize = header->Identitication[EI_CLASS] == ELF_CLASS_TYPE32 ? ELF64_OFFSET_SIZE / 2 : ELF64_OFFSET_SIZE;
	uint64_t wordSize = ELF64_WORD_SIZE;//header->Identitication[EI_CLASS] == ELF_CLASS_TYPE32 ? ELF64_WORD_SIZE / 2 : ELF64_WORD_SIZE;
	uint64_t halfSize = ELF64_HALF_SIZE;//header->Identitication[EI_CLASS] == ELF_CLASS_TYPE32 ? ELF64_HALF_SIZE / 2 : ELF64_HALF_SIZE;
	size_t cursor = ELF_IDENTITIFCATION_MAX_SIZE;
	
	memcpy(&header->FileType, &bytes[cursor], halfSize); //e_type

	cursor += halfSize;

	memcpy(&header->ISA, &bytes[cursor], halfSize); //e_machine

	cursor += halfSize;

	memcpy(&header->ELFVersion, &bytes[cursor], wordSize); //e_version

	cursor += wordSize;

	memcpy(&header->EntryAddress, &bytes[cursor], offsetSize); //e_entry

	cursor += offsetSize;

	memcpy(&header->ProgramHeaderOffset, &bytes[cursor], offsetSize); //e_phoff

	cursor += offsetSize;

	memcpy(&header->SectionHeaderOffset, &bytes[cursor], offsetSize); //e_shoff

	cursor += offsetSize;

	memcpy(&header->Flags, &bytes[cursor], wordSize); //e_flags

	cursor += wordSize;

	memcpy(&header->HeaderSize, &bytes[cursor], halfSize); //e_ehsize

	cursor += halfSize;

	memcpy(&header->ProgramHeaderEntrySize, &bytes[cursor], halfSize); //e_phentsize

	cursor += halfSize;

	memcpy(&header->ProgramHeaderEntryCount, &bytes[cursor], halfSize); //e_phnum

	cursor += halfSize;

	memcpy(&header->SectionHeaderEntrySize, &bytes[cursor], halfSize); //e_shentsize

	cursor += halfSize;

	memcpy(&header->SectionHeaderEntryCount, &bytes[cursor], halfSize); //e_shnum

	cursor += halfSize;

	memcpy(&header->SectionNamesIndex, &bytes[cursor], halfSize); //e_shstrndx

	switch(bytes[EI_DATA])
	{
		case ELF_BYTE_ORDER_LSB:
			printf("Data: 2's complement, little endian.\n");
			break;
		case ELF_BYTE_ORDER_MSB:
			printf("Data: 2's complement, big endian.\n");
			break;
		default:
			printf("Data: Unknown byte order\n");
	}

	if (bytes[EI_HEADER_VERSION] == 1) printf("ELF File Version: Current (1)\n");
	else printf("ELF File Version: Invalid (%d)\n", bytes[EI_HEADER_VERSION]);

	char buffer[256];

	GetELFOSABIString(bytes[EI_OSABI], buffer);

	printf("OS ABI: %s\n", buffer);

	GetELFMachineTypeText((ELFMachine)header->ISA, buffer);

	printf("Machine type: %s\n", buffer);

	printf("Object File Version: %#02x\n", header->ELFVersion);

	printf("Entry point address: %#06lx\n", header->EntryAddress);

	printf("Program Header Offset: %ld bytes\n", header->ProgramHeaderOffset);

	printf("Section Header Offset: %ld bytes\n", header->SectionHeaderOffset);

	printf("Flags: %#04x\n", header->Flags);

	printf("Size of Header: %d bytes\n", header->HeaderSize);

	printf("Size of Program Headers: %d bytes\n", header->ProgramHeaderEntrySize);

	printf("Number of Programer Headers: %d\n", header->ProgramHeaderEntryCount);

	printf("Size of Section Header: %d bytes\n", header->SectionHeaderEntrySize);

	printf("Number of Section Headers: %d\n", header->SectionHeaderEntryCount);

	printf("Section Header string table index: %d\n", header->SectionNamesIndex);
	
    return header;
}

int ELFMagicNumberValid(const unsigned char* ident) {
    char magic[5] = { ident[EI_MAG0], ident[EI_MAG1], ident[EI_MAG2], ident[EI_MAG3], '\0' };

    return strcmp(magic, "\x7f""ELF") == 0;
}

void ELFPrintMagic(const struct elf_header* header) {
    if (!header) return;

    printf("Magic: ");

    for(int i = 0; i < ELF_IDENTITIFCATION_MAX_SIZE; i++)
	{
		printf("%02x ", header->Identitication[i]);
	}

	printf("\n");
}

struct elf_header* CreateELFHeader(void) {
    struct elf_header* header = calloc(sizeof(struct elf_header), 1);

    if (!header) 
    {
        fprintf(stderr, "Failed to calloc\n");

        exit(1);
    }

    return header;
}