import * as reg from './regfile';


// Op-code of instructions.
//
//     |76543210|
//     |......ds|
//     |.......s <--- Operation size: 1 = word, 0 = byte.
//     |......d <--- Operation direction: 1 = register is destination, 0 = register is source.
//     |765432 <--- Op-code
export const enum OP {
    MOV     = 0x89,
    MOVL    = 0xB8,
    MOVQ    = 0xC7,
    MOVABS  = 0xB8, // Lower 3 bits encode destination register.
    INC     = 0xff,
    DEC     = 0xff,
    PUSH    = 0x50,
    POP     = 0x58,
}

export const enum OP_SIZE {
    BYTE = 0b0,
    WORD = 0b1,
}

export const enum OP_DIRECTION {
    REG_IS_SRC = 0b00,
    REG_IS_DST = 0b10,
}

// Two bits of `MOD` field in `Mod-R/M` byte.
export const enum MOD {
    INDIRECT = 0b00,
    DISP8 = 0b01,
    DISP32 = 0b10,
    REG_TO_REG = 0b11,
}

// Op-code extension into `REG` field of `Mod-R/M` byte.
export const enum OPREG {
    INC = 0,
    DEC = 1,
}

// Displacement.
export const enum DISP {
    BYTE,
    WORD,
    LONG,
}

export const RM_NEEDS_SIB = 0b100;

export interface IModRM {
    mod: number;
    reg: number;
    rm: number;
}

export interface ISib {
    scale: number;
    index: number;
    base: number;
}

export type number64 = [number, number];


// Transpiler returns a machine code of every instruction.
// It does not have any logic and does not do any optimizations or type
// checking, it just returns the machine code for every ASM instruction.
export class Encoder {

    // See Intel manual: 2-8 Vol. 2A
    // See: http://wiki.osdev.org/X86-64_Instruction_Encoding#REX_prefix
    // `W`, `R`, `X`, `B` are numbers `0` or `1`.
    protected rex(W: number, R: number, X: number, B: number) {
        return 0b01000000 | (W << 3) + (R << 2) + (X << 1) + B;
    }

    protected rexW(R: number = 0, X: number = 0, B: number = 0) {
        return this.rex(1, R, X, B);
    }

    protected rexOneOperand(r: reg.Register, needs_rex = true) {
        var W = 0, B = 0;
        if((r.size == 64) && needs_rex) W = 1;
        if(r.isExtended) B = 1;
        return W || B ? this.rex(W, 0, 0, B) : null;
    }

    protected rexRegToReg(src: reg.Register, dst: reg.Register) {
        var W = 0, R = 0, X = 0, B = 0;

        // Determine if we operate on 64-bit registers.
        if(src.size == 64) W = 1;
        else if(dst.size == 64) W = 1;

        // Extend bits in `Mod-R/M` byte.
        if(src.isExtended) B = 1;
        if(dst.isExtended) R = 1;

        return W || R || X || B ? this.rex(W, R, X, B) : null;
    }

    protected rexMemToReg(src: reg.MemoryReference, dst: reg.Register) {
        var W = 0, R = 0, X = 0, B = 0;
        
        if(src.base && src.base.size == 64) W = 1;
        else if(dst.size == 64) W = 1;
        
        if(src.base && src.base.isExtended) B = 1;
        if(src.index && src.index.isExtended) X = 1;
        if(dst.isExtended) R = 1;

        return W || R || X || B ? this.rex(W, R, X, B) : null;
    }

    protected rexRegToMem(src: reg.Register, dst: reg.MemoryReference) {
        var W = 0, R = 0, X = 0, B = 0;

        if(dst.base && dst.base.size == 64) W = 1;
        else if(src.size == 64) W = 1;

        if(src.isExtended) R = 1;
        if(dst.base && dst.base.isExtended) B = 1;
        if(dst.index && dst.index.isExtended) X = 1;

        return W || R || X || B ? this.rex(W, R, X, B) : null;
    }

    protected rexMem(r: reg.Register, ref: reg.MemoryReference) {
        var W = 0, R = 0, X = 0, B = 0;

        if(ref.base && ref.base.size == 64) W = 1;
        else if(r.size == 64) W = 1;

        if(r.isExtended) R = 1;
        if(ref.base && ref.base.isExtended) B = 1;
        if(ref.index && ref.index.isExtended) X = 1;

        return W || R || X || B ? this.rex(W, R, X, B) : null;
    }

    protected getDisplacementSize(disp) {
        if((disp <= 0x7f) && (disp >= -0x80)) return DISP.BYTE;
        if((disp <= 0x7fffffff) && (disp >= -0x80000000)) return DISP.LONG;
        throw Error(`Displacement ${disp} too big.`);
    }

    // Creates a Mod-REG-R/M byte, `mod` is mode, `register` is ID of the first register,
    // `rm` is ID of the 2nd register or memory value.
    protected modrm(mod, register, rm) {
        // return ((mod & 0b11) << 6) | ((register & 0b111) << 3) | (rm & 0b111);
        // return (mod << 6) + (register << 3) + rm;
        return (mod << 6) | (register << 3) | rm;
    }

    protected modrmPack(modrm: IModRM) {
        return this.modrm(modrm.mod, modrm.reg, modrm.rm);
    }

    protected modrmOneOperand(dst: reg.Register, opreg: OPREG = 0): IModRM {
        return {
            mod: MOD.REG_TO_REG,
            reg: opreg,
            rm: dst.id,
        };
    }

    protected modrmRegToReg(src: reg.Register, dst: reg.Register): IModRM {
        return {
            mod: MOD.REG_TO_REG,
            reg: src.id,
            rm: dst.id,
        };
    }

    protected modrmMemToReg(src: reg.MemoryReference, dst: reg.Register, mod = MOD.INDIRECT): IModRM {
        var rm = src.base.id;

        // There will be a `SIB` byte, we have to set `R/M` to `0b100` = `RM_NEEDS_SIB`.
        var need_sib = !!src.index;
        if(need_sib) rm = RM_NEEDS_SIB;

        return {
            mod: mod,
            reg: dst.id,
            rm: rm,
        };
    }

    protected modrmRegToMem(src: reg.Register, dst: reg.MemoryReference, mod = MOD.INDIRECT): IModRM {
        var rm = dst.base.id;

        // There will be a `SIB` byte, we have to set `R/M` to `0b100`.
        var need_sib = !!dst.index;
        if(need_sib) rm = RM_NEEDS_SIB;

        return {
            mod: mod,
            reg: src.id,
            rm: rm,
        };
    }

    protected modrmMem(r: reg.Register, ref: reg.MemoryReference, mod = MOD.INDIRECT): IModRM {
        var rm = ref.base ? ref.base.id : 0;

        // There will be a `SIB` byte, we have to set `R/M` to `0b100`.
        var need_sib = !!ref.index;
        if(need_sib) rm = RM_NEEDS_SIB;

        return {
            mod: mod,
            reg: r.id,
            rm: rm,
        };
    }

    protected isSibNeeded(modrm: IModRM, ref: reg.MemoryReference): boolean {
         return !!ref.index || (modrm.rm == RM_NEEDS_SIB) ? true : false;
    }

    protected sib(scale, index, base): number {
        return (scale << 6) | (index << 3) | base;
    }

    protected sibPack(sib: ISib): number {
        return this.sib(sib.scale, sib.index, sib.base);
    }

    protected sibFromRef(ref: reg.MemoryReference): ISib {
        return {
            scale: ref.scale,
            index: ref.index ? ref.index.id : 0,
            base: ref.base ? ref.base.id : 0,
        };
    }

    protected insOneOperand(r: reg.Register, op: OP, opreg: OPREG = 0, hasRex = true, reg_in_op = false, imm: number[] = []) {
        var ins = [];

        // REX.W | REX.B prefix.
        if(hasRex || r.isExtended) {
            var rex = this.rexOneOperand(r, hasRex);
            if(rex) ins.push(rex);
        }

        // Op-code
        if(reg_in_op) { // If 3 lower bits in Op-code are used to encode register.
            op |= r.id;
            ins.push(op);
        } else {
            ins.push(op);

            // Mod-R/M
            var modrm = this.modrmOneOperand(r, opreg);
            ins.push(this.modrmPack(modrm));
        }

        // Immediate
        if(imm.length) for(var long of imm) this.pushConstant(ins, long);

        return ins;
    }

    protected insRegToReg(src: reg.Register, dst: reg.Register, op: OP, hasRex = true) {
        var ins = [];

        // REX.W | REX.B prefix.
        if(hasRex) {
            var rex = this.rexRegToReg(src, dst);
            if(rex) ins.push(rex);
        }

        // Op-code
        op = (op & 0b11111101) | OP_DIRECTION.REG_IS_SRC;
        ins.push(op);

        // Mod-R/M
        var modrm = this.modrmRegToReg(src, dst);
        ins.push(this.modrmPack(modrm));

        return ins;
    }

    protected insMemToReg(src: reg.MemoryReference, dst: reg.Register, op: OP, hasRex = true, op_size: OP_SIZE = OP_SIZE.WORD) {
        var ins = [];

        if(hasRex) {
            var rex = this.rexMemToReg(src, dst);
            if(rex) ins.push(rex);
        }

        // Set op-code's direction and size bits.
        // op |= OP_DIRECTION.REG_IS_DST;
        op = (op & 0b11111101) | OP_DIRECTION.REG_IS_DST;
        op = (op & 0b11111110) | op_size;
        ins.push(op);

        var mod = MOD.INDIRECT;
        if(src.displacement) {
            var disp_size = this.getDisplacementSize(src.displacement);
            if(disp_size == DISP.BYTE) mod = MOD.DISP8;
            else mod = MOD.DISP32;
        }

        // Mod-R/M
        var modrm = this.modrmMemToReg(src, dst, mod);
        ins.push(this.modrmPack(modrm));

        // SIB
        if(this.isSibNeeded(modrm, src)) {
            var sib = this.sibFromRef(src);
            ins.push(this.sibPack(sib));
        }

        // Displacement
        if(src.displacement) {
            if(mod == MOD.DISP8) ins.push(src.displacement); // Only one byte.
            else this.pushConstant(ins, src.displacement);
        }

        return ins;
    }

    protected insRegToMem(src: reg.Register, dst: reg.MemoryReference, op: OP, hasRex = true, op_size: OP_SIZE = OP_SIZE.WORD) {
        var ins = [];

        if(hasRex) {
            var rex = this.rexRegToMem(src, dst);
            if(rex) ins.push(rex);
        }

        op = (op & 0b11111101) | OP_DIRECTION.REG_IS_SRC;
        op = (op & 0b11111110) | op_size;
        ins.push(op);

        var mod = MOD.INDIRECT;
        if(dst.displacement) {
            var disp_size = this.getDisplacementSize(dst.displacement);
            if(disp_size == DISP.BYTE) mod = MOD.DISP8;
            else mod = MOD.DISP32;
        }

        // Mod-R/M
        var modrm = this.modrmRegToMem(src, dst, mod);
        ins.push(this.modrmPack(modrm));

        // SIB
        if(this.isSibNeeded(modrm, dst)) {
            var sib = this.sibFromRef(dst);
            ins.push(this.sibPack(sib));
        }

        // Displacement
        if(dst.displacement) {
            if(mod == MOD.DISP8) ins.push(dst.displacement); // Only one byte.
            else this.pushConstant(ins, dst.displacement);
        }

        return ins;
    }

    // Operation where one operand is a memory reference.
    protected insMem(src: reg.Register|reg.MemoryReference, dst: reg.Register|reg.MemoryReference, op: OP, hasRex = true, op_size: OP_SIZE = OP_SIZE.WORD) {
        var ins = [];

        var r: reg.Register, ref: reg.MemoryReference;
        if(src instanceof reg.Register) {
            r = src as reg.Register;
            ref = dst as reg.MemoryReference;
        } else {
            r = dst as reg.Register;
            ref = src as reg.MemoryReference;
        }

        // REX prefix
        if(hasRex) {
            var rex = this.rexMem(r, ref);
            if(rex) ins.push(rex);
        }

        // Set direction of the reg-to-mem or mem-to-reg.
        op = (op & 0b11111101) | (r === src ? OP_DIRECTION.REG_IS_SRC : OP_DIRECTION.REG_IS_DST);

        // TODO: Size of the operands, make this actually useful.
        op = (op & 0b11111110) | op_size;

        // Op-code
        ins.push(op);

        var mod = MOD.INDIRECT;
        if(ref.displacement) {
            var disp_size = this.getDisplacementSize(ref.displacement);
            if(disp_size == DISP.BYTE) mod = MOD.DISP8;
            else mod = MOD.DISP32;
        }

        // Mod-R/M
        var modrm = this.modrmMem(r, ref, mod);
        ins.push(this.modrmPack(modrm));

        // SIB
        if(this.isSibNeeded(modrm, ref)) {
            var sib = this.sibFromRef(ref);
            ins.push(this.sibPack(sib));
        }

        // Displacement
        if(ref.displacement) {
            if(mod == MOD.DISP8) ins.push(ref.displacement); // Only one byte.
            else this.pushConstant(ins, ref.displacement); // Push bytes in reverse order.
        }

        return ins;
    }

    protected pushConstant(arr: number[], constant: number) {
        arr.push(constant & 0xff);
        arr.push((constant >> 8) & 0xff);
        arr.push((constant >> 16) & 0xff);
        arr.push((constant >> 24) & 0xff);
        return arr;
    }

    inc(register: reg.Register) {
        return this.insOneOperand(register, OP.INC, OPREG.INC);
    }

    dec(register: reg.Register) {
        return this.insOneOperand(register, OP.DEC, OPREG.DEC);
    }

    // mov_r_r(src: reg.Register, dst: reg.Register) {
    //     return this.insRegToReg(src, dst, OP.MOV);
    // }

    movq_r_r(src: reg.Register64, dst: reg.Register64) {
        if(!(src instanceof reg.Register64) || !(dst instanceof reg.Register64))
            throw Error('`movq` is defined only on 64-bit registers.');
        return this.insRegToReg(src, dst, OP.MOV);
    }

    mov_r_r(src: reg.Register, dst: reg.Register) {
        return this.insRegToReg(src, dst, OP.MOV);
    }

    movq_m_r(src: reg.MemoryReference, dst: reg.Register64) {
        return this.insMem(src, dst, OP.MOV);
        // return this.insMemToReg(src, dst, OP.MOV);
    }

    movq_imm_r(imm: number, dst: reg.Register64) {
        return this.insOneOperand(dst, OP.MOVQ, 0, true, false, [imm]);
    }

    movq_r_m(src: reg.Register64, dst: reg.MemoryReference) {
        return this.insMem(src, dst, OP.MOV);
        // return this.insRegToMem(src, dst, OP.MOV);
    }

    movq_rm(src: reg.Register64|reg.MemoryReference, dst: reg.Register64|reg.MemoryReference) {
        return this.insRegToReg(src, dst, OP.MOV);
    }

    mov_rm(src: reg.Register|reg.MemoryReference, dst: reg.Register|reg.MemoryReference) {
        return this.insRegToReg(src, dst, OP.MOV);
    }

    movabs(imm: [number, number], dst: reg.Register) {
        if(!(dst instanceof reg.Register64))
            throw Error('`movabs` operates only on 64-bit registers.');

        return this.insOneOperand(dst, OP.MOVABS, 0, true, true, imm);
    }

    movq(src: reg.MemoryReference|reg.Register64|number|number64, dst: reg.MemoryReference|reg.Register64) {
        if(src instanceof reg.MemoryReference) {
            if(dst instanceof reg.MemoryReference)
                throw Error(`Cannot do memory-to-memory operation: movq ${src.toString()}, ${dst.toString()}`);
            else if(dst instanceof reg.Register) return this.insMem(src, dst, OP.MOV);
            else throw Error(`Invalid operand type: movq ${src.toString()}, ${dst}`);
        } else if(src instanceof reg.Register) {
            if(dst instanceof reg.MemoryReference) return this.insMem(src, dst, OP.MOV);
            else if(dst instanceof reg.Register) return this.insRegToReg(src, dst, OP.MOV);
            else throw Error(`Invalid operand type: movq ${src.toString()}, ${dst}`);
        } else if(typeof src == 'number') {
            var imm = src as number;
            if(dst instanceof reg.Register) {
                this.insOneOperand(dst, OP.MOVQ, 0, true, true, [imm]);
            } else throw Error(`Invalid operand type: movq $${src}, ${dst}`);
        } else if((src instanceof Array) && (src.length == 2)) {
            var imm64 = src as number64;
            if(dst instanceof reg.Register) return this.movabs(imm64, dst);
            else throw Error(`Invalid operand type: movq $, ${dst}`);
        }
        throw Error('Invalid operand types: movq');
    }

    mov(src: reg.MemoryReference|reg.Register|number|number64, dst: reg.MemoryReference|reg.Register) {
        if(src instanceof reg.Register64) {
            if(dst instanceof reg.Register64) return this.movq_r_r(src, dst);
            else if(dst instanceof reg.MemoryReference) return this.movq_r_m(src, dst);
            else throw Error(`Destination operand [${dst.toString()}] invalid.`);
        } else if(src instanceof reg.Register32) {

        }
    }



    movl_imm_r32(imm: number, dst: reg.Register32) {
        return this.insOneOperand(dst, OP.MOVL, 0, true, true, [imm]);
    }



    movl_r_r(src: reg.Register, dst: reg.Register) {
        return this.insRegToReg(src, dst, OP.MOV, false);
    }

    movb_r_r(src: reg.Register, dst: reg.Register) {
        return this.insRegToReg(src, dst, OP.MOV, false);
    }

    nop(size = 1) {

    }

    nopw() {
        return this.nop(2);
    }

    nopl() {
        return this.nop(4);
    }

    push_r(r: reg.Register64) {
        return this.insOneOperand(r, OP.PUSH, 0, false, true);
    }

    pop_r(r: reg.Register64) {
        return this.insOneOperand(r, OP.POP, 0, true, true);
    }

}


// Generates logically equivalent code to `Encoder` but the actual
// bytes of the machine code will likely differ, because `FuzzyEncoder`
// picks at random one of the possible instructions when multiple instructions
// can perform the same operation. Here are some examples:
//
//  - Bits in `REX` prefix are ignored if they don't have an effect on the instruction.
//  - Register-to-register `MOV` instruction can be encoded in two different ways.
//  - Up to four prefixes may be added to instruction, if they are not used, they are ignored.
//  - There can be many different *no-op* instruction that are used to fill in padding for byte alignment, for example:
//
//     mov %rax, %rax
//     add $0, %rax
//     ...
export class FuzzyEncoder extends Encoder {
    oneOrZero(): number {
        return Math.random() > 0.5 ? 1 : 0;
    }

    protected rexRegToReg(src: reg.Register, dst: reg.Register) {
        var rex = super.rexRegToReg(src, dst);
        if(rex) {
            // `REX.X` bit is not used in *register-to-register* instructions.
            rex = (rex & 0b01) | (this.oneOrZero() << 1);
        }
        return rex;
    }
}