/*
 * vm86.c
 *
 * Copyright (C) 2016 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero 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 Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <vm86.h>
#include <segments.h>
#include <exception.h>
#include <process.h>
#include <thread.h>
#include <memory.h>

#define PREFIX_LOCK     (1 << 0)
#define PREFIX_OPSIZE   (1 << 1)
#define PREFIX_ADDRSIZE (1 << 2)
#define PREFIX_REP      (1 << 3)
#define PREFIX_REPNZ    (1 << 4)
#define PREFIX_ES       (1 << 5)
#define PREFIX_SS       (1 << 6)
#define PREFIX_FS       (1 << 7)
#define PREFIX_GS       (1 << 8)

#define VM86_MEM_START      0x10000
#define VM86_MEM_END        0x90000
#define VM86_MEM_PARAGRAPHS ((VM86_MEM_END - VM86_MEM_START) >> 4)

#define VM86_TRAMPOLINE_CS 0x0000
#define VM86_TRAMPOLINE_IP 0xE000
#define VM86_TRAMPOLINE_SS 0x0000
#define VM86_TRAMPOLINE_SP 0xDFFA

extern void vm86_start(vm86_registers_t input_regs, vm86_registers_t *output_regs);

static bool_t vm86_interrupts = TRUE;
static dword_t vm86_mem_bitmap[VM86_MEM_PARAGRAPHS / 32] = { 0 };
static dword_t vm86_mem_last_alloc_bitmap[VM86_MEM_PARAGRAPHS / 32] = { 0 };

static inline dword_t get_bios_interrupt(byte_t number)
{
    return *((dword_t*)(number * sizeof(dword_t)));
}

static inline byte_t peekb(word_t segment, word_t offset)
{
    return *((byte_t*)((segment << 4) + offset));
}

static inline void pokeb(word_t segment, word_t offset, byte_t value)
{
    *((byte_t*)((segment << 4) + offset)) = value;
}

static inline word_t peekw(word_t segment, word_t offset)
{
    return *((word_t*) ((segment << 4) + offset));
}

static inline void pokew(word_t segment, word_t offset, word_t value)
{
    *((word_t*)((segment << 4) + offset)) = value;
}

static inline dword_t peekl(word_t segment, word_t offset)
{
    return *((dword_t*)((segment << 4) + offset));
}

static inline void pokel(word_t segment, word_t offset, dword_t value)
{
    *((dword_t*)((segment << 4) + offset)) = value;
}

word_t vm86_alloc(word_t size)
{
    dword_t i;
    bool_t found = FALSE;
    dword_t first_free = VM86_MEM_PARAGRAPHS;

    for (i = 0; i < VM86_MEM_PARAGRAPHS; i++)
    {
        if (test_bit(vm86_mem_bitmap, i)) first_free = VM86_MEM_PARAGRAPHS;
        else if (first_free == VM86_MEM_PARAGRAPHS) first_free = i;

        if ((first_free != VM86_MEM_PARAGRAPHS) && ((i - first_free + 1) == size))
        {
            found = TRUE;
            break;
        }
    }

    if (!found) return 0;

    for (i = 0; i < size; i++)
    {
        set_bit(vm86_mem_bitmap, first_free + i);

        if (i != (size - 1)) clear_bit(vm86_mem_last_alloc_bitmap, first_free + i);
        else set_bit(vm86_mem_last_alloc_bitmap, first_free + i);
    }

    return (VM86_MEM_START >> 4) + first_free;
}

void vm86_free(word_t paragraph)
{
    dword_t i;

    for (i = paragraph - (VM86_MEM_START >> 4); !test_bit(vm86_mem_last_alloc_bitmap, i); i++)
    {
        clear_bit(vm86_mem_bitmap, i);
    }
}

void vm86_handler(registers_ext_vm86_t *regs)
{
    dword_t i, prefix = 0, count = 0, segment = 0, operand = 0;
    byte_t instruction;

    while (TRUE)
    {
        instruction = peekb(regs->cs, regs->eip);
        regs->eip++;

        if (instruction == 0x26) prefix |= PREFIX_ES;
        else if (instruction == 0x36) prefix |= PREFIX_SS;
        else if (instruction == 0x64) prefix |= PREFIX_FS;
        else if (instruction == 0x65) prefix |= PREFIX_GS;
        else if (instruction == 0x66) prefix |= PREFIX_OPSIZE;
        else if (instruction == 0x67) prefix |= PREFIX_ADDRSIZE;
        else if (instruction == 0xF0) prefix |= PREFIX_LOCK;
        else if (instruction == 0xF2) prefix |= PREFIX_REPNZ;
        else if (instruction == 0xF3) prefix |= PREFIX_REP;
        else break;
    }

    if (prefix & PREFIX_REP)
    {
        if (prefix & PREFIX_ADDRSIZE) count = regs->ecx;
        else count = regs->ecx & 0xFFFF;
    }
    else
    {
        count = 1;
    }

    if (prefix & PREFIX_ES) segment = regs->es;
    else if (prefix & PREFIX_SS) segment = regs->ss;
    else if (prefix & PREFIX_FS) segment = regs->fs;
    else if (prefix & PREFIX_GS) segment = regs->gs;
    else segment = regs->ds;

    switch (instruction)
    {
        case 0x6C: // insb byte ptr [es:di], dx
        {
            for (i = 0; i < count; i++)
            {
                pokeb(regs->es, regs->edi, cpu_read_port_byte(regs->edx));
                if (!(regs->eflags & CPU_STATUS_FLAG_DF)) regs->edi++;
                else regs->edi--;
            }

            break;
        }

        case 0x6D: // insw/insd (d)word ptr [es:di], dx
        {
            for (i = 0; i < count; i++)
            {
                if (!(prefix & PREFIX_OPSIZE))
                {
                    pokew(regs->es, regs->edi, cpu_read_port_word(regs->edx));

                    if (!(regs->eflags & CPU_STATUS_FLAG_DF)) regs->edi += 2;
                    else regs->edi -= 2;
                }
                else
                {
                    pokel(regs->es, regs->edi, cpu_read_port_dword(regs->edx));

                    if (!(regs->eflags & CPU_STATUS_FLAG_DF)) regs->edi += 4;
                    else regs->edi -= 4;
                }
            }

            break;
        }

        case 0x6E: // outsb dx, byte ptr [ds:si]
        {
            for (i = 0; i < count; i++)
            {
                cpu_write_port_byte(regs->edx, peekb(segment, regs->esi));

                if (!(regs->eflags & CPU_STATUS_FLAG_DF)) regs->esi++;
                else regs->esi--;
            }

            break;
        }

        case 0x6F: // outsw/outsd dx, (d)word ptr [ds:si]
        {
            for (i = 0; i < count; i++)
            {
                if (!(prefix & PREFIX_OPSIZE))
                {
                    cpu_write_port_word(regs->edx, peekw(segment, regs->esi));

                    if (!(regs->eflags & CPU_STATUS_FLAG_DF)) regs->esi += 2;
                    else regs->esi -= 2;
                }
                else
                {
                    cpu_write_port_dword(regs->edx, peekl(segment, regs->esi));

                    if (!(regs->eflags & CPU_STATUS_FLAG_DF)) regs->esi += 4;
                    else regs->esi -= 4;
                }
            }

            break;
        }

        case 0x9C: // pushf
        {
            operand = regs->eflags;
            operand &= ~CPU_STATUS_FLAG_VM;

            if (vm86_interrupts) operand |= CPU_STATUS_FLAG_IF;
            else operand &= ~CPU_STATUS_FLAG_IF;

            if (!(prefix & PREFIX_OPSIZE))
            {
                regs->esp3 -= 2;
                pokew(regs->ss, regs->esp3, operand & 0xFFFF);
            }
            else
            {
                regs->esp3 -= 4;
                pokel(regs->ss, regs->esp3, operand);
            }

            break;
        }

        case 0x9D: // popf
        {
            if (!(prefix & PREFIX_OPSIZE))
            {
                regs->eflags &= 0xFFFF0000;
                regs->eflags |= peekw(regs->ss, regs->esp3) | CPU_STATUS_FLAG_VM;
                regs->esp3 += 2;
            }
            else
            {
                regs->eflags = peekl(regs->ss, regs->esp3) | CPU_STATUS_FLAG_VM;
                regs->esp3 += 4;
            }

            vm86_interrupts = (regs->eflags & CPU_STATUS_FLAG_IF) ? TRUE : FALSE;
            regs->eflags |= CPU_STATUS_FLAG_IF;

            break;
        }

        case 0xCD: // int
        {
            operand = peekb(regs->cs, regs->eip);
            regs->eip++;

            regs->esp3 -= 2;
            pokew(regs->ss, regs->esp3, (regs->eflags & (~CPU_STATUS_FLAG_VM)) | (vm86_interrupts ? CPU_STATUS_FLAG_IF : 0));

            regs->esp3 -= 2;
            pokew(regs->ss, regs->esp3, regs->cs);

            regs->esp3 -= 2;
            pokew(regs->ss, regs->esp3, regs->eip);

            operand = get_bios_interrupt(operand);

            regs->cs = (operand >> 16) & 0xFFFF;
            regs->eip = operand & 0xFFFF;

            break;
        }

        case 0xCF: // iret
        {
            if (regs->cs == VM86_TRAMPOLINE_CS && regs->eip == VM86_TRAMPOLINE_IP + 1)
            {
                registers_t *stack_regs = (registers_t*)get_kernel_esp();
                stack_regs->esp += 20;

                set_kernel_esp(stack_regs->esp);
                vm86_registers_t *results = *(vm86_registers_t**)(stack_regs->esp + sizeof(vm86_registers_t));

                results->gs = regs->gs;
                results->fs = regs->fs;
                results->ds = regs->ds;
                results->es = regs->es;
                results->ss = regs->ss;
                results->esp = regs->esp;
                results->eflags = regs->eflags;
                results->cs = regs->cs;
                results->eip = regs->eip;
                results->eax = regs->eax;
                results->ecx = regs->ecx;
                results->edx = regs->edx;
                results->ebx = regs->ebx;
                results->ebp = regs->ebp;
                results->esi = regs->esi;
                results->edi = regs->edi;

                memcpy(regs, stack_regs, sizeof(registers_t));
                return;
            }

            regs->eip = peekw(regs->ss, regs->esp3);
            regs->esp3 += 2;

            regs->cs = peekw(regs->ss, regs->esp3);
            regs->esp3 += 2;

            regs->eflags = peekw(regs->ss, regs->esp3) | CPU_STATUS_FLAG_VM;
            regs->esp3 += 2;

            vm86_interrupts = (regs->eflags & CPU_STATUS_FLAG_IF) ? TRUE : FALSE;
            regs->eflags |= CPU_STATUS_FLAG_IF;

            break;
        }

        case 0xE4: // in al, <byte>
        {
            operand = peekb(regs->cs, regs->eip);
            regs->eip++;
            regs->eax &= 0xFFFFFF00;
            regs->eax |= cpu_read_port_byte(operand & 0xFF) & 0xFF;

            break;
        }

        case 0xE5: // in (e)ax, <byte>
        {
            operand = peekb(regs->cs, regs->eip);
            regs->eip++;

            if (!(prefix & PREFIX_OPSIZE))
            {
                regs->eax &= 0xFFFF0000;
                regs->eax |= cpu_read_port_word(operand & 0xFF) & 0xFFFF;
            }
            else regs->eax = cpu_read_port_dword(operand & 0xFF);

            break;
        }

        case 0xE6: // out <byte>, al
        {
            operand = peekb(regs->cs, regs->eip);
            regs->eip++;
            cpu_write_port_byte(operand & 0xFF, regs->eax & 0xFF);

            break;
        }

        case 0xE7: // out <byte>, (e)ax
        {
            operand = peekb(regs->cs, regs->eip);
            regs->eip++;

            if (!(prefix & PREFIX_OPSIZE)) cpu_write_port_word(operand & 0xFF, regs->eax & 0xFFFF);
            else cpu_write_port_dword(operand & 0xFF, regs->eax);

            break;
        }

        case 0xEC: // in al, dx
        {
            regs->eax &= 0xFFFFFF00;
            regs->eax |= cpu_read_port_byte(regs->edx & 0xFFFF) & 0xFF;

            break;
        }

        case 0xED: // in (e)ax, dx
        {
            if (!(prefix & PREFIX_OPSIZE))
            {
                regs->eax &= 0xFFFF0000;
                regs->eax |= cpu_read_port_word(regs->edx & 0xFFFF) & 0xFFFF;
            }
            else regs->eax = cpu_read_port_dword(regs->edx & 0xFFFF);

            break;
        }

        case 0xEE: // out dx, al
        {
            cpu_write_port_byte(regs->edx & 0xFFFF, regs->eax & 0xFF);
            break;
        }

        case 0xEF: // out dx, (e)ax
        {
            if (!(prefix & PREFIX_OPSIZE)) cpu_write_port_word(regs->edx & 0xFFFF, regs->eax & 0xFFFF);
            else cpu_write_port_dword(regs->edx & 0xFFFF, regs->eax);

            break;
        }

        case 0xFA: // cli
        {
            vm86_interrupts = FALSE;
            break;
        }

        case 0xFB: // sti
        {
            vm86_interrupts = TRUE;
            break;
        }

        default:
        {
            KERNEL_CRASH_WITH_REGS("General Protection Fault (VM86)", (registers_t*)regs);
        }
    }
}

dword_t vm86_interrupt(byte_t number, vm86_registers_t *regs)
{
    critical_t critical;
    enter_critical(&critical);

    process_t *old_process = switch_process(kernel_process);

    dword_t ret = map_memory_internal(NULL, NULL, 0x100000, PAGE_PRESENT | PAGE_WRITABLE | PAGE_USERMODE);
    if (ret != ERR_SUCCESS) return ret;

    dword_t far_ptr = get_bios_interrupt(number);
    regs->cs = (far_ptr >> 16) & 0xFFFF;
    regs->eip = far_ptr & 0xFFFF;
    regs->ss = VM86_TRAMPOLINE_SS;
    regs->esp = VM86_TRAMPOLINE_SP;
    regs->eflags = 0x00020202;

    pokew(VM86_TRAMPOLINE_SS, VM86_TRAMPOLINE_SP, VM86_TRAMPOLINE_IP);
    pokew(VM86_TRAMPOLINE_SS, VM86_TRAMPOLINE_SP + 2, VM86_TRAMPOLINE_CS);
    pokew(VM86_TRAMPOLINE_SS, VM86_TRAMPOLINE_SP + 4, 0x0002);

    pokeb(VM86_TRAMPOLINE_CS, VM86_TRAMPOLINE_IP, 0xCF);

    vm86_start(*regs, regs);

    unmap_memory_internal(NULL, 0x100000);
    switch_process(old_process);
    leave_critical(&critical);
    return ERR_SUCCESS;
}
