#pragma once
#include <stddef.h>

#include <stdio.h>

namespace assembler {

typedef unsigned char byte; 

struct Register; 
struct MemAddress; 

enum RegCode : byte {
    EAX = 0,
    ECX = 1,
    EDX = 2,
    EBX = 3,
    ESP = 4, SIB = 4,
    EBP = 5, ADDR = 5,
    ESI = 6,
    EDI = 7,

    R8 = 8,
    R9 = 9,
    R10 = 10,
    R12 = 12,
    R13 = 13,
    R14 = 14,
    R15 = 15,

    //R11 is supervolatile. Can change in between
    //virtual asm ops (used as a temporary), so be careful
    //with using it.
    R11 = 11,

    //XMM registers have unique numbers so we can recognize them
    XMM0 = 16,
    XMM1 = 17,
    XMM2 = 18,
    XMM3 = 19,
    XMM4 = 20,
    XMM5 = 21,
    XMM6 = 22,
    XMM7 = 23,

    XMM8 = 24,
    XMM9 = 25,
    XMM10 = 26,
    XMM11 = 27,
    XMM12 = 28,
    XMM13 = 29,
    XMM14 = 30,
    XMM15 = 31,

    NONE = 0,
}; 

//Floating Point Register codes representing the stack registers on the FPU
// The top of the stack is always FPU_0
enum FloatReg : byte {
    FPU_0 = 0,
    FPU_1 = 1,
    FPU_2 = 2,
    FPU_3 = 3,
    FPU_4 = 4,
    FPU_5 = 5,
    FPU_6 = 6,
    FPU_7 = 7,
}; 

enum JumpType {
    Overflow,
    NotOverflow,
    Below, Carry = Below,
    NotBelow, NotCarry = NotBelow,
    Equal, Zero = Equal,
    NotEqual, NotZero = NotEqual,
    NotAbove,
    Above,
    Sign,
    NotSign,
    Parity,
    NotParity,
    Less,
    GreaterOrEqual,
    LessOrEqual,
    Greater,
    Jump,

    JumpTypeCount
}; 


//Handles thread safety for the JIT
struct CriticalSection {
    void* pLock; 

    void enter(); 
    void leave(); 

    CriticalSection(); 
    ~CriticalSection(); 
}; 

struct AddrPrefix {
    MemAddress& adr; 
    bool defLong; 
    unsigned char further; 

    AddrPrefix(MemAddress& Adr, bool DefLong, unsigned char Further)
        : adr(Adr), defLong(DefLong), further(Further) {
    }
}; 

struct RegPrefix {
    Register& reg; 
    unsigned short other; 
    bool defLong; 

    RegPrefix(Register& Reg, unsigned short Other, bool DefLong)
        : reg(Reg), other(Other), defLong(DefLong) {
    }
}; 

//Stores information about the code page
//    Generates an executable page in memory when created
//    Deletes the asssociated page when deleted
//Implementation in virtual_asm_<operating system>.cpp (e.g. virtual_asm_windows.cpp)
struct CodePage {
    void* page; 
    unsigned int size, used, references; 
    bool final; 

    CodePage(unsigned int Size, void* requestedStart = 0); 
    ~CodePage(); 

    void grab(); 
    void drop(); 

    //Call finalize when done writing to the code page to guarantee that it can be executed
    //No more writing may be done to the allocated pages
    void finalize(); 

    //Returns the pointer to the first currently unused chunk of the page
    template<class T>
    T getFunctionPointer() {
        return reinterpret_cast<T>((byte*)page+used); 
    }

    byte* getActivePage() const {
        return (byte*)page+used; 
    }

    //Marks bytes as used; 
    //future calls to getFunctionPointer() will not reference the location that is being marked as used
    void markBytesUsed(unsigned int count) {
        used += count; 
    }

    //Marks bytes up to <address> as used
    void markUsedAddress(void* address) {
        unsigned newUsed = (unsigned)((byte*)address - (byte*)page); 
        if(newUsed > used && newUsed <= size)
            used = newUsed; 
    }

    //Returns the number of bytes not yet allocated to a function
    unsigned int getFreeSize() const {
        return size-used; 
    }

    //Returns the smallest page (in bytes) that can be allocated by a code page (Sizes other than multiples of this size allocate an extra page)
    static unsigned int getMinimumPageSize(); 

private:
    CodePage() {}
}; 

//Stores the code pointer and provides access to various processor-level operations
// To work with the processor, create a set of 'Register' instances, each taking the RegCode of the associated register (e.g. Register eax(cpu, EAX))
//Implementation in virtual_asm_<processor instruction set>.cpp (e.g. virtual_asm_x86.cpp)
struct Processor {
    //Pointer to the location for the next opcode
    byte* op; 
    byte* pageStart; 
    //The current mode of operation, in bits
    // e.g. 32 bits for x86, indicating that operations should treat addresses as if they were unsigned integers
    unsigned bitMode, lastBitMode; 
    //The number of bytes currently on the stack that we are responsible for
    unsigned stackDepth; 
    //Reserved jump space
    unsigned jumpSpace; 
    byte* jumpPtr; 

    //Initializes the processor to point to the active page of the code page
    //Optionally takes a bitMode override (defaults to the same bitMode as the exe)
    Processor(CodePage& codePage, unsigned defaultBitMode = sizeof(void*)*8 ); 

    //Creates a jump to the new code page, and marks the current address as used on the old code page
    //Updates output pointer to the new code page's active page
    void migrate(CodePage& prevPage, CodePage& newPage); 

    //Changes the current bitMode, and stores the previous bitMode
    void setBitMode(unsigned bits) {
        lastBitMode = bitMode; 
        bitMode = bits; 
    }

    //Restores the previous bitMode
    void resetBitMode() {
        bitMode = lastBitMode; 
    }

    //Returns the alignment of the stack (number of bytes a push increments esp)
    static unsigned pushSize(); 

    //Pushes data to the opcode output
    template<class T>
    Processor& operator<<(T b) {
        *(T*)op = b;  op += sizeof(T); 
        return *this; 
    }

    //Pushes bytes representing a memory address to the opcode output
    template<class T>
    Processor& operator<<(MemAddress addr); 

    //Pushes bytes representing a prefix
    template<class T>
    Processor& operator<<(AddrPrefix pr); 

    template<class T>
    Processor& operator<<(RegPrefix pr); 

    //Calls the function, passing the arguments specified by 'args'
    //args is a string like "rrcmrm" which specifies arguments as sourced by a Register*, MemAddres*, or a constant
    //EBP is invalid during the call
    void call_cdecl(void* func, const char* args, va_list ap); 
    void call_cdecl(void* func, const char* args, ...); 

    //Use call() in between these to set up a call with an arbitrary function
    unsigned call_cdecl_args(const char* args, ...); 
    unsigned call_cdecl_args(const char* args, va_list ap); 
    unsigned call_thiscall_args(Register* obj, const char* args, ...); 
    unsigned call_thiscall_args(Register* obj, const char* args, va_list ap); 

    //Prepares for a call to manual call to a cdecl function (Do not use with call_cdecl)
    // Use before pushing arguments
    // Invalidates EBP until call_cdecl_end()
    void call_cdecl_prep(unsigned argBytes); 
    //Ends a manual call to a cdecl function (Do not use with call_cdecl)
    // Use after returning from the function
    void call_cdecl_end(unsigned argBytes, bool returnPointer = false); 

    //Note: stdcall is like cdecl, but does not use cdecl_end

    //Calls the function, passing the arguments specified by 'args'
    //args is a string like "rrcmrm" which specifies arguments as sourced by a Register*, MemAddres*, or a constant
    //EBP is invalid during the call
    void call_stdcall(void* func, const char* args, ...); 

    //To call a thiscall:
    // cpu.call_thiscall_prep(total argument size)
    // cpu.push(arguments)
    // cpu.call_thiscall_this(source of 'this' pointer)
    // cpu.call(function)
    // cpu.call_thiscall_end(total argument size)
    void call_thiscall_prep(unsigned argBytes); 
    void call_thiscall_this(MemAddress address); 
    void call_thiscall_this(Register& reg); 
    void call_thiscall_this_mem(MemAddress address, Register& memreg); 
    void call_thiscall_this_mem(Register& reg, Register& memreg); 
    void call_thiscall_end(unsigned argBytes, bool returnPointer = false); 

    //Calls a function (push code pointer, jump to function)
    void call(Register& reg); 
    void call(void* func); 

    //Pushes a constant value onto the stack (Pushes are always pushSize() large, values beyond this size are an error)
    void push(size_t value); 
    //Pops <count> times (Pops are always pushSize() large)
    void pop(unsigned int count); 

    //Pushes the value of <reg> onto the stack
    void push(Register& reg); 
    //Pops the alue of <reg> from the stack
    void pop(Register& reg); 

    //Get a register corresponding to an argument on 64-bit calling convention
    unsigned maxIntArgs64(); 
    unsigned maxFloatArgs64(); 
    bool isIntArg64Register(unsigned char number, unsigned char arg); 
    bool isFloatArg64Register(unsigned char number, unsigned char arg); 
    Register intArg64(unsigned char number, unsigned char arg); 
    Register floatArg64(unsigned char number, unsigned char arg); 
    Register intArg64(unsigned char number, unsigned char arg, Register defaultReg); 
    Register floatArg64(unsigned char number, unsigned char arg, Register defaultReg); 
    Register floatReturn64(); 
    Register intReturn64(); 

    //Pushes the memory at <address> onto the stack (Pushes are always pushSize() large, pushing larger values invokes multiple pushes)
    void push(MemAddress address); 
    //Pops the value on the stack to the memory at <address> (Pops are always pushSize() large, popping larger values invokes multiple pops)
    void pop(MemAddress address); 

    //Prepares a short jump (fewer than approx. 120 bytes in either direction)
    // Pass the return to a matching end_short_jump
    void* prep_short_jump(JumpType type); 
    //Ends a short jump
    void end_short_jump(void* p); 

    //Prepares a large jump (can jump to any location)
    // Pass the return to a matching end_long_jump
    void* prep_long_jump(JumpType type); 
    //Ends a large jump
    void end_long_jump(void* p); 

    //Jumps to <dest>
    void jump(JumpType type, volatile byte* dest); 
    //Jumps to the address in <reg>
    void jump(Register& reg); 
    //Decrements ecx and jumps if it becomes 0;  Optionally conditionally jumps based on a Zero/NotZero test
    void loop(volatile byte* dest, JumpType type = Jump); 

    //Copies from *esi to *edi, and adjusts them both by the data size according to the direction flag
    void string_copy(unsigned size); 
    //Sets direction flag for string copy
    void setDirFlag(bool forward); 

    //Returns from a function (pop code pointer, jump there)
    void ret(); 

    //Triggers a debug break
    void debug_interrupt(); 

private:
    Processor() {}
}; 

//Provides access to the floating point unit's state
//Implementation in virtual_asm_<processor instruction set>.cpp (e.g. virtual_asm_x86.cpp)
struct FloatingPointUnit {
    Processor& cpu; 

    FloatingPointUnit(Processor& CPU); 

    //Clears the FPU's state and registers
    void init(); 

    //Negates FPU_0
    void negate(); 

    //Pushes
    void load_const_0(); 
    void load_const_1(); 

    //FPU_1 becomes FPU_0 (Pops the fpu stack)
    void pop(); 

    //Exchanges contents of FPU_n and FPU_0
    void exchange(FloatReg floatReg); 

    //Compares FPU_0 to floatReg, setting the CPU's flags according to the values' relation
    // Optionally pops the fpu stack
    void compare_toCPU(FloatReg floatReg, bool pop = true); 

    //Pushes the specified data type stored at <address> onto the FPU stack (becomes FPU_0)
    void load_float(MemAddress address); 
    void load_dword(MemAddress address); 
    void load_qword(MemAddress address); 
    void load_double(MemAddress address); 

    //Stores the value on FPU_0 to <address> according to the data type
    // Optionally pops the fpu stack
    void store_float(MemAddress address, bool pop = true); 
    void store_dword(MemAddress address, bool pop = true); 
    void store_double(MemAddress address, bool pop = true); 

    //Control words
    void store_control_word(MemAddress address); 
    void load_control_word(MemAddress address); 

    //Effect: FPU_0 -= <reg>
    void operator-=(FloatReg reg); 

    //Effect: FPU_0 += *(float*)address
    void add_float(MemAddress address); 
    //Effect: FPU_0 -= *(float*)address
    void sub_float(MemAddress address); 
    //Effect: FPU_0 *= *(float*)address
    void mult_float(MemAddress address); 
    //Effect: FPU_0 /= *(float*)address
    void div_float(MemAddress address); 

    //Effect: FPU_0 += *(double*)address
    void add_double(MemAddress address); 
    void add_double(FloatReg reg, bool pop = true); 
    //Effect: FPU_0 -= *(double*)address
    // If Reversed: FPU_0 = *(double*)address - FPU_0
    void sub_double(MemAddress address, bool reversed = false); 
    void sub_double(FloatReg reg, bool reversed = false, bool pop = true); 
    //Effect: FPU_0 *= *(double*)address
    void mult_double(MemAddress address); 
    void mult_double(FloatReg reg, bool pop = true); 
    //Effect: FPU_0 /= *(double*)address
    // If Reversed: FPU_0 = *(double*)address / FPU_0
    void div_double(MemAddress address, bool reversed = false); 
    void div_double(FloatReg reg, bool reversed = false, bool pop = true); 
}; 

//Temporary struct that represents an addition to a memory address, with optional scaling
struct ScaledIndex {
    RegCode reg; 
    unsigned char scaleFactor; 

    ScaledIndex(RegCode Reg, unsigned char Scale) : reg(Reg), scaleFactor(Scale) {}
}; 

//Temporary struct that stores data necessary for memory access
//  Provides operations that can be performed on a memory address
//Implementation in virtual_asm_<processor instruction set>.cpp (e.g. virtual_asm_x86.cpp)
struct MemAddress {
    Processor& cpu; 
    RegCode code; 
    void* absolute_address; 
    unsigned char other; 
    int offset; 
    unsigned bitMode; 
    bool Float; 
    bool Signed; 
    unsigned char scaleFactor; 
    RegCode scaleReg; 

    MemAddress(Processor& CPU, void* address); 
    MemAddress(Processor& CPU, RegCode Code); 
    MemAddress(Processor& CPU, RegCode Code, int Offset); 
    MemAddress operator+(ScaledIndex scale); 
    MemAddress operator+(int Offset); 
    MemAddress operator-(int Offset); 

    void operator++(); 
    void operator--(); 

    void operator-(); 
    void operator~(); 

    void operator+=(unsigned int amount); 
    void operator-=(unsigned int amount); 

    void operator=(unsigned int value); 
    void operator=(void* pointer); 
    void operator=(Register fromReg); 

    void operator&=(unsigned int value); 
    void operator|=(unsigned int value); 

    //Copies memory using an intermediate register
    void direct_copy(MemAddress address, Register& intermediate); 
    AddrPrefix prefix(unsigned char further = 0, bool defLong = false); 
}; 

//Converts a MemAddress from the default unsigned <cpu bit mode> to match the passed type
template<class T>
MemAddress as(MemAddress addr) {
    addr.bitMode = sizeof(T) * 8; 
    addr.Signed = (T)-1 < (T)0; 
    return addr; 
}

template<>
MemAddress as<float>(MemAddress addr); 

template<>
MemAddress as<double>(MemAddress addr); 

//Structure that provides operations that can be performed on a register
//  Also provides the means to generate MemAddresses relative to a register via dereference (e.g. *eax+8)
//Implementation in virtual_asm_<processor instruction set>.cpp (e.g. virtual_asm_x86.cpp)
struct Register {
    Processor& cpu; 
    RegCode code; 
    unsigned bitMode; 

    Register(Processor& CPU, RegCode Code); 
    Register(Processor& CPU, RegCode Code, unsigned BitModeOverride); 

    void set_regCode(Register& other) {
        code = other.code; 
        bitMode = other.bitMode; 
    }

    unsigned getBitMode() const; 
    unsigned getBitMode(const MemAddress& addr) const; 

    MemAddress operator*() const; 
    ScaledIndex operator*(unsigned char scale) const; 

    //Loads the address pointed to by <address> into this register
    void copy_address(MemAddress address); 

    void swap(MemAddress address); 
    void swap(Register& other); 

    void operator<<=(Register& other); 
    void operator>>=(Register& other); 
    void rightshift_logical(Register& other); 

    void operator+=(unsigned int amount); 
    void operator+=(MemAddress address); 
    void operator+=(Register& other); 

    void operator-=(unsigned int amount); 
    void operator-=(Register& other); 
    void operator-=(MemAddress address); 

    void operator*=(MemAddress address); 

    void operator-(); 
    void operator~(); 

    void operator--(); 
    void operator++(); 

    void operator&=(unsigned long long mask); 
    void operator&=(MemAddress address); 
    void operator&=(Register other); 

    void operator^=(MemAddress address); 
    void operator^=(Register& other); 

    void operator|=(MemAddress address); 
    void operator|=(unsigned long long mask); 

    //Copies a smaller data type, retaining the sign
    void copy_expanding(MemAddress address); 
    //Copies an 8 bit register, leaving 0s in higher bytes
    void copy_zeroing(Register& other); 

    void operator=(unsigned long long value); 
    void operator=(void* pointer); 
    void operator=(Register other); 
    void operator=(MemAddress addr); 

    void operator==(Register other); 
    void operator==(MemAddress addr); 
    void operator==(unsigned int test); 

    void setIf(JumpType condition); 
    void* setDeferred(unsigned long long def = 0); 

    bool xmm(); 
    bool extended(); 
    RegCode index(); 

    RegPrefix prefix(unsigned short other = 0, bool defaultLong = false); 
    RegPrefix prefix(Register& other, bool defaultLong = false); 
    unsigned char modrm(unsigned short other); 
    unsigned char modrm(Register& other); 

    //Multiplies *address with value, stores the result in this register
    void multiply_signed(MemAddress address, int value); 

    //Divides {eax,edx} by this register;  result in eax, remainder in edx
    void divide(); 
    void divide_signed(); 
}; 

//Converts a MemAddress from the default unsigned <cpu bit mode> to match the passed type
template<class T>
Register as(Register reg) {
    reg.bitMode = sizeof(T) * 8; 
    return reg; 
}
}; 
