/* This file is part of asmc, a bootstrapping OS with minimal seed
   Copyright (C) 2019 Giovanni Mascellani <gio@debian.org>
   https://gitlab.com/giomasce/asmc

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>. */

#include <stdint.h>
#include <assert.h>
#include <stdio.h>
#include <string.h>
#include <asmc_api.h>

struct gdt_entry {
    unsigned int limit_low : 16;
    unsigned int base_low : 24;
    // Access byte
    unsigned int accessed : 1;
    unsigned int read_write : 1; // Readable for code, writable for data
    unsigned int conforming : 1; // Conforming for code, expand down for data
    unsigned int code : 1; // 1 for code, 0 for data
    unsigned int type : 1; // 1 for code/data, 0 for system segments (like TSS)
    unsigned int DPL : 2; //priviledge level
    unsigned int present : 1;
    // Flags
    unsigned int limit_high : 4;
    unsigned int always_0 : 2;
    unsigned int size : 1; // 1 for 32 bits, 0 for 16 bits
    unsigned int gran : 1; // 1 for 4k page addressing, 0 for byte addressing
    unsigned int base_high : 8;
} __attribute__((packed));

struct gdt_desc {
    uint16_t size;
    struct gdt_entry *offset;
} __attribute__((packed));

void reload_gdt(struct gdt_entry *gdt, int entry_num) {
    struct gdt_desc desc;
    desc.size = entry_num * 8 - 1;
    desc.offset = gdt;
    asm volatile("lgdt %0 \n\t"
                 "push $0x10 \n\t"
                 "push $reload_cs \n\t"
                 // TCC does not know retf
                 ".byte 0xcb \n\t"
                 "reload_cs: \n\t"
                 "mov $0x18, %%ax \n\t"
                 // TCC generates a 0x66 prefix that does not look like right
                 /*"mov %%ax, %%ds \n\t"
                 "mov %%ax, %%es \n\t"
                 "mov %%ax, %%fs \n\t"
                 "mov %%ax, %%gs \n\t"
                 "mov %%ax, %%ss \n\t"*/
                 ".byte 0x8e, 0xd8 \n\t"
                 ".byte 0x8e, 0xc0 \n\t"
                 ".byte 0x8e, 0xe0 \n\t"
                 ".byte 0x8e, 0xe8 \n\t"
                 ".byte 0x8e, 0xd0 \n\t"
                 : : "m" (desc));
    printf("Still alive after reloading code and data segments!\n");
}

// Load GDT, code and data segments according to Linux Boot Protocol
void load_linux_gdt() {
    assert(sizeof(struct gdt_desc) == 6);
    assert(sizeof(struct gdt_entry) == 8);
    // Random address, hoping that it does not bother anybody
    struct gdt_entry *gdt = (struct gdt_entry*) 0x800;
    memset(gdt, 0, sizeof(gdt[0]) * 4);

    // 0x10: code segment
    gdt[2].limit_low = 0xffff;
    gdt[2].base_low = 0x000000;
    gdt[2].accessed = 0;
    gdt[2].read_write = 1;
    gdt[2].conforming = 0;
    gdt[2].code = 1;
    gdt[2].type = 1;
    gdt[2].DPL = 0;
    gdt[2].present = 1;
    gdt[2].limit_high = 0xf;
    gdt[2].size = 1;
    gdt[2].gran = 1;
    gdt[2].base_high = 0x00;

    // 0x18: data segment
    gdt[3].limit_low = 0xffff;
    gdt[3].base_low = 0x000000;
    gdt[3].accessed = 0;
    gdt[3].read_write = 1;
    gdt[3].conforming = 0;
    gdt[3].code = 0;
    gdt[3].type = 1;
    gdt[3].DPL = 0;
    gdt[3].present = 1;
    gdt[3].limit_high = 0xf;
    gdt[3].size = 1;
    gdt[3].gran = 1;
    gdt[3].base_high = 0x00;

    reload_gdt(gdt, 4);
}

// Prepare GDT to re-enter real mode
void prepare_real_mode_gdt() {
    assert(sizeof(struct gdt_desc) == 6);
    assert(sizeof(struct gdt_entry) == 8);
    // Random address, hoping that it does not bother anybody
    struct gdt_entry *gdt = (struct gdt_entry*) 0x800;
    memset(gdt, 0, sizeof(gdt[0]) * 4);

    // 0x10: code segment
    gdt[2].limit_low = 0xffff;
    gdt[2].base_low = 0x000000;
    gdt[2].accessed = 0;
    gdt[2].read_write = 1;
    gdt[2].conforming = 0;
    gdt[2].code = 1;
    gdt[2].type = 1;
    gdt[2].DPL = 0;
    gdt[2].present = 1;
    gdt[2].limit_high = 0x0;
    gdt[2].size = 0;
    gdt[2].gran = 0;
    gdt[2].base_high = 0x00;

    // 0x18: data segment
    gdt[3].limit_low = 0xffff;
    gdt[3].base_low = 0x000000;
    gdt[3].accessed = 0;
    gdt[3].read_write = 1;
    gdt[3].conforming = 0;
    gdt[3].code = 0;
    gdt[3].type = 1;
    gdt[3].DPL = 0;
    gdt[3].present = 1;
    gdt[3].limit_high = 0x0;
    gdt[3].size = 0;
    gdt[3].gran = 0;
    gdt[3].base_high = 0x00;
}

void prepare_setup_header(void *zero_page, void *initrd, size_t initrd_size) {
    char *czp = (char*) zero_page;
    *((uint16_t*)(czp+0x1fa)) = 0xffff;
    *((uint8_t*)(czp+0x210)) = 0xff;
    *((uint8_t*)(czp+0x211)) = 0x81;
    *((void**)(czp+0x218)) = initrd;
    *((uint32_t*)(czp+0x21c)) = initrd_size;
    *((uint16_t*)(czp+0x224)) = 0xde00;
    *((void**)(czp+0x228)) = (void*) 0x1e000;
}

void prepare_boot_params(void *zero_page, void *initrd, size_t initrd_size) {
    prepare_setup_header(zero_page, initrd, initrd_size);
}

#define SECT_SIZE 512

void load_linux32() {
    download_file("http://10.0.2.2:8080/bzImage", "/ram/cont/bzImage");
    download_file("http://10.0.2.2:8080/initrd.img", "/ram/cont/initrd.img");
    printf("Loading Linux kernel\n");
    FILE *fin = fopen("/ram/cont/bzImage", "rb");
    void *real_mode = malloc(SECT_SIZE);
    size_t res = fread(real_mode, SECT_SIZE, 1, fin);
    assert(res == 1);
    uint8_t setup_sects = *((uint8_t*)(((char*)real_mode)+0x1f1));
    setup_sects = setup_sects == 0 ? 4 : setup_sects;
    size_t real_mode_size = SECT_SIZE * (1 + setup_sects);
    printf("Real mode code/data is 0x%x bytes long\n", real_mode_size);
    real_mode = realloc(real_mode, real_mode_size);
    res = fread(((char*)real_mode) + SECT_SIZE, SECT_SIZE, setup_sects, fin);
    assert(res == setup_sects);
    assert(*((uint16_t*)(((char*)real_mode)+0x1fe)) == 0xaa55);
    assert(*((uint32_t*)(((char*)real_mode)+0x202)) == 0x53726448);
    size_t prot_mode_size;
    void *prot_mode = read_whole_file(fin, &prot_mode_size);
    printf("Protected mode code/data is 0x%x bytes long\n", prot_mode_size);
    fclose(fin);
    fin = fopen("/ram/cont/initrd.img", "rb");
    size_t initrd_size;
    void *initrd = read_whole_file(fin, &initrd_size);
    printf("Initrd is 0x%x bytes long\n", initrd_size);
    fclose(fin);
    // Align initrd
    initrd = realloc(initrd, initrd_size + 4096);
    void *initrd_aligned = (void*)(((((unsigned long)initrd)-1) | (4096-1)) + 1);
    memmove(initrd_aligned, initrd, initrd_size);
    printf("Initrd initially allocated at 0x%x, then moved at 0x%x\n", initrd, initrd_aligned);
    prepare_boot_params(real_mode, initrd_aligned, initrd_size);
    load_linux_gdt();
    char *cmdline = "verbose earlyprintk nokaslr console=uart8250,io,0x3f8,115200n8";

    // Here we really smash our operating system!
    memset((void*) 0x10000, 0, 0x10000);
    memcpy((void*) 0x10000, real_mode, real_mode_size);
    memcpy((void*) 0x1e000, cmdline, strlen(cmdline));
    //memcpy((void*) 0x100000, prot_mode, prot_mode_size);
    asm volatile("mov $0x100000, %%edi \n\t"
                 "cld \n\t"
                 "rep movsb \n\t"
                 "mov $0x10000, %%esi \n\t"
                 "xor %%eax, %%eax \n\t"
                 "mov %%eax, %%ebp \n\t"
                 "mov %%eax, %%edi \n\t"
                 "mov %%eax, %%ebx \n\t"
                 "mov %%eax, %%ecx \n\t"
                 "mov %%eax, %%edx \n\t"
                 "mov $0xe000, %%esp \n\t"
                 "push $0x100000 \n\t"
                 "ret \n\t"
                 : : "S" (prot_mode), "c" (prot_mode_size));

    free(initrd);
    free(real_mode);
    free(prot_mode);
}

void load_linux() {
    download_file("http://10.0.2.2:8080/bzImage", "/ram/cont/bzImage");
    download_file("http://10.0.2.2:8080/realmode.bin", "/ram/cont/realmode.bin");
    printf("Loading Linux kernel\n");
    FILE *fin = fopen("/ram/cont/bzImage", "rb");
    void *real_mode = malloc(SECT_SIZE);
    size_t res = fread(real_mode, SECT_SIZE, 1, fin);
    assert(res == 1);
    uint8_t setup_sects = *((uint8_t*)(((char*)real_mode)+0x1f1));
    setup_sects = setup_sects == 0 ? 4 : setup_sects;
    size_t real_mode_size = SECT_SIZE * (1 + setup_sects);
    printf("Real mode code/data is 0x%x bytes long\n", real_mode_size);
    real_mode = realloc(real_mode, real_mode_size);
    res = fread(((char*)real_mode) + SECT_SIZE, SECT_SIZE, setup_sects, fin);
    assert(res == setup_sects);
    assert(*((uint16_t*)(((char*)real_mode)+0x1fe)) == 0xaa55);
    assert(*((uint32_t*)(((char*)real_mode)+0x202)) == 0x53726448);
    size_t prot_mode_size;
    void *prot_mode = read_whole_file(fin, &prot_mode_size);
    printf("Protected mode code/data is 0x%x bytes long\n", prot_mode_size);
    fclose(fin);
    prepare_setup_header(real_mode, 0, 0);
    //load_linux_gdt();
    fin = fopen("/ram/cont/realmode.bin", "rb");
    size_t trampoline_size;
    void *trampoline = read_whole_file(fin, &trampoline_size);
    printf("Trampoline is 0x%x bytes long\n", trampoline_size);
    fclose(fin);
    prepare_real_mode_gdt();
    char *cmdline = "verbose";

    // Here we really smash our operating system!
    memset((void*) 0x10000, 0, 0x10000);
    memcpy((void*) 0x500, trampoline, trampoline_size);
    memcpy((void*) 0x10000, real_mode, real_mode_size);
    memcpy((void*) 0x1e000, cmdline, strlen(cmdline));
    memcpy((void*) 0x100000, prot_mode, prot_mode_size);
    ((void (*)()) 0x500)();

    free(trampoline);
    free(real_mode);
    free(prot_mode);
}

void load_mlb() {
    download_file("http://10.0.2.2:8080/realmode2", "/ram/cont/realmode.bin");
    download_file("http://10.0.2.2:8080/example", "/ram/cont/bootloader");
    FILE *fin = fopen("/ram/cont/realmode.bin", "rb");
    size_t trampoline_size;
    void *trampoline = read_whole_file(fin, &trampoline_size);
    printf("Trampoline is 0x%x bytes long\n", trampoline_size);
    fclose(fin);
    prepare_real_mode_gdt();
    fin = fopen("/ram/cont/bootloader", "rb");
    size_t bootloader_size;
    void *bootloader = read_whole_file(fin, &bootloader_size);
    printf("Bootloader is 0x%x bytes long\n", bootloader_size);
    fclose(fin);
    prepare_real_mode_gdt();

    memset((void*) 0x10000, 0, 0x10000);
    memcpy((void*) 0x500, trampoline, trampoline_size);
    memcpy((void*) 0x7c00, bootloader, bootloader_size);
    ((void (*)()) 0x500)();
}

int main(int argc, char *argv[]) {
    printf("Hello world from continue2.c!\n");
    //load_linux32();
    //load_linux();
    //load_mlb();
    return 0;
}