#include "stack.h"
#include "elf.h"
#include "klibc.h"
#include "paging.h"
#include "types.h"

/** Defined in linker.ld script  **/
extern char __multiboot_data_start;
extern char __multiboot_data_end;
static size_t multibootDataSize;
static size_t multibootUsage;

/** Initialized at 'StackSymbolSetup'. */
static Elf32_Sym_t *elfSymtab;
static size_t elfSymtabSize;

static const char *elfStrtab;
static size_t elfStrtabSize;

static void *stackSaveSection(const void *addr, size_t size, const char *sectionName)
{
    if (size > (multibootDataSize - multibootUsage)) {
        printf("No enough room to store %s (%d/%d). Increase multiboot_data section size",
               sectionName, size, (multibootDataSize - multibootUsage));
        return NULL;
    }
    void *destAddr = (void *)&__multiboot_data_start + multibootUsage;

    memcpy((void *)destAddr, addr, size);
    multibootUsage += size;

    return destAddr;
}

/**
 * Copy the symbols table and string table from a multiboot information.
 * We need to copy them as they are loaded by the bootloader and could be override as soon
 * as we play with memory around our loaded kernel.
 * As special ELF section is allocated for this purpose.
 */
void stackSymbolSetup(multiboot_info_t *mbi)
{
    /** Get the section header table as an array of section headers. */
    Elf32_Shdr_t *sht = (Elf32_Shdr_t *)mbi->u.elf_sec.addr;
    size_t sec_len    = (size_t)mbi->u.elf_sec.num;
    assert(mbi->u.elf_sec.size == sizeof(Elf32_Shdr_t));
    multibootDataSize = &__multiboot_data_end - &__multiboot_data_start;
    multibootUsage    = 0;

    /**
     * The section header at index 'shndx' in the table is a meta section
     * header. It contains a list of section header names in the table. We
     * should look for section header names ".symtab" & ".strtab".
     */
    const char *sh_names = (const char *)sht[mbi->u.elf_sec.shndx].sh_addr;

    /** Loop through the table and look for ".symtab" & ".strtab". */
    for (size_t i = 0; i < sec_len; ++i) {
        const char *name = sh_names + sht[i].sh_name;

        if (strcmp(name, ".symtab") == 0) {
            assert(sht[i].sh_type == SHT_SYMTAB);
            elfSymtab = stackSaveSection((void *)sht[i].sh_addr, sht[i].sh_size, name);
            if (elfSymtab)
                elfSymtabSize = sht[i].sh_size;

        } else if (strcmp(name, ".strtab") == 0) {
            assert(sht[i].sh_type == SHT_STRTAB);
            elfStrtab = stackSaveSection((void *)sht[i].sh_addr, sht[i].sh_size, name);
            if (elfStrtab)
                elfStrtabSize = sht[i].sh_size;
        }
    }
}

/** Look up an address in symbols map and return its function name or NULL. */
const char *lookupSymbolName(uint32_t addr)
{
    size_t symtab_len = elfSymtabSize / sizeof(Elf32_Sym_t);

    for (size_t i = 0; i < symtab_len; ++i) {
        if ((ELF_SYM_TYPE(elfSymtab[i].st_info) == ELF_SYM_TYPE_FUNC) &&
            (addr >= elfSymtab[i].st_value) &&
            (addr <= elfSymtab[i].st_value + elfSymtab[i].st_size)) {
            return elfStrtab + elfSymtab[i].st_name;
        }
    }

    return NULL;
}

void printStackTrace(unsigned int maxFrames)
{
#ifdef DEBUG
    // Now on Stack:
    //	( potential second function argument )
    //	first function argument (maxFrames)
    //	return address from caller
    //	EBP (Extended Base Pointer) of calling function
    //
    // retrace function from address could done by optaining function address with gdb or
    // objdump -S kernel
    unsigned int *ebp = __builtin_frame_address(0);
    for (unsigned int frame = 0; frame < maxFrames; frame++) {
        unsigned int eip = *(ebp + 1);
        if (eip == 0) {
            // No caller on stack
            break;
        }
        unsigned int *arguments = ebp + 2;
        const char *functionName = lookupSymbolName(eip);
        if (functionName != NULL) {
            printf("[%d] %s (", frame, functionName);
        } else {
            printf("[%d] 0x%x (", frame, eip);
        }
        int nbArg = 0;
        do {
            if ((_stack_bottom <= (vaddr_t)arguments) && ((vaddr_t)arguments <= _stack_top)) {
                printf(" 0x%x", *arguments);
                arguments += 1;
            } else {
                break;
            }
            nbArg++;
            if (nbArg >= 4) {
                break;
            }
        } while (1);
        printf(")\n");
        ebp = (unsigned int *)(ebp[0]);
    }
#else
    printf("Must be compiled with -fno-omit-frame-pointer for full stack\n");
    unsigned int *ebp = &maxFrames - 2;
    unsigned int *eip  = ebp + sizeof(unsigned int);
    printf("[0] 0x%x\n", (unsigned int)eip);
#endif
}