#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;
}
};
