/*
 * CS 261 PA3: Mini-ELF disassembler
 *
 * Name: Nicholas Tamassia
 *
 * This code was developed in compliance with the JMU Honor code.
 */

#include "p3-disas.h"

/**********************************************************************
 *                         REQUIRED FUNCTIONS
 *********************************************************************/

bool has_register_byte(y86_inst_t ins);

bool has_register_byte(y86_inst_t ins) {
    switch (ins.icode) {
        case CMOV:
        case IRMOVQ:
        case RMMOVQ:
        case MRMOVQ:
        case OPQ:
        case PUSHQ:
        case POPQ:
            return true;
        default:
            return false;
    }
}

bool has_valc_bytes(y86_inst_t ins);

bool has_valc_bytes(y86_inst_t ins) {
    switch (ins.icode) {
        case IRMOVQ:
        case RMMOVQ:
        case MRMOVQ:
        case JUMP:
        case CALL:
            return true;
        default:
            return false;
    }
}

void validate_opcode(y86_inst_t *ins, byte_t *current_byte);

void validate_opcode(y86_inst_t *ins, byte_t *current_byte) {
    if (ins == NULL || current_byte == NULL) {
        return;
    }

    y86_icode_t opcode_high = *current_byte >> 4;
    byte_t opcode_low = *current_byte & 0xF;
    ins->icode = opcode_high;

    switch (ins->icode) {
        case CMOV:
            if (opcode_low >= BADCMOV) {
                ins->icode = INVALID;
            }
            ins->ifun.cmov = opcode_low;
            break;

        case OPQ:
            if (opcode_low >= BADOP) {
                ins->icode = INVALID;
            }
            ins->ifun.op = opcode_low;
            break;

        case JUMP:
            if (opcode_low >= BADJUMP) {
                ins->icode = INVALID;
            }
            ins->ifun.jump = opcode_low;
            break;

        case IOTRAP:
            if (opcode_low >= BADTRAP) {
                ins->icode = INVALID;
            }
            ins->ifun.trap = opcode_low;
            break;

        case HALT:
        case NOP:
        case IRMOVQ:
        case RMMOVQ:
        case MRMOVQ:
        case CALL:
        case RET:
        case PUSHQ:
        case POPQ:
            if (opcode_low != 0x0) {
                ins->icode = INVALID;
            }
            ins->ifun.b = opcode_low;
            break;

        default:
            ins->icode = INVALID;
            break;
    }
}

void validate_registers(y86_inst_t *ins, byte_t *current_byte);

void validate_registers(y86_inst_t *ins, byte_t *current_byte) {
    if (ins == NULL || current_byte == NULL) {
        return;
    }

    ins->ra = *current_byte >> 4;
    ins->rb = *current_byte & 0xF;

    switch (ins->icode) {
        case IRMOVQ:
            if (ins->ra != NOREG || ins->rb == NOREG) {
                ins->icode = INVALID;
            }
            break;

        case PUSHQ:
        case POPQ:
            if (ins->ra == NOREG || ins->rb != NOREG) {
                ins->icode = INVALID;
            }
            break;

        case RMMOVQ:
        case MRMOVQ:
            if (ins->ra == NOREG) {
                ins->icode = INVALID;
            }
            break;

        case CMOV:
        case OPQ:
            if (ins->ra == NOREG || ins->rb == NOREG) {
                ins->icode = INVALID;
            }
            break;
        default:
            break;
    }
}

void assign_valc(y86_inst_t *ins, byte_t *current_byte);

void assign_valc(y86_inst_t *ins, byte_t *current_byte) {
    if (ins == NULL || current_byte == NULL) {
        return;
    }

    uint64_t valc = 0;
    for (int i = 0; i < 8; i++) {
        valc |= (uint64_t) * (current_byte + i) << (8 * i);
    }

    switch (ins->icode) {
        case IRMOVQ:
            ins->valC.v = valc;
            break;
        case RMMOVQ:
        case MRMOVQ:
            ins->valC.d = valc;
            break;
        case JUMP:
        case CALL:
            ins->valC.dest = valc;
            break;
        default:
            break;
    }
}

y86_inst_t fetch(y86_t *cpu, byte_t *memory) {
    y86_inst_t ins = {0};

    if (cpu == NULL || memory == NULL) {
        ins.icode = INVALID;
        return ins;
    }

    if (cpu->pc > MEMSIZE - 10) {
        ins.icode = INVALID;
        cpu->stat = ADR;
        return ins;
    }

    byte_t *current_byte = &memory[cpu->pc];

    validate_opcode(&ins, current_byte);
    if (ins.icode == INVALID) {
        cpu->stat = INS;
        return ins;
    }

    current_byte++;

    if (has_register_byte(ins)) {
        validate_registers(&ins, current_byte);

        if (ins.icode == INVALID) {
            cpu->stat = INS;
            return ins;
        }

        current_byte++;
    }

    if (has_valc_bytes(ins)) {
        assign_valc(&ins, current_byte);
        current_byte += 8;
    }

    ins.valP = current_byte - memory;

    if (ins.icode == HALT) {
        cpu->stat = HLT;
    } else {
        cpu->stat = AOK;
    }

    return ins;
}

/**********************************************************************
 *                         OPTIONAL FUNCTIONS
 *********************************************************************/

void print_register(y86_regnum_t y86_register);

void print_register(y86_regnum_t y86_register) {
    switch (y86_register) {
        case RAX:
            printf("%%rax");
            break;
        case RCX:
            printf("%%rcx");
            break;
        case RDX:
            printf("%%rdx");
            break;
        case RBX:
            printf("%%rbx");
            break;
        case RSP:
            printf("%%rsp");
            break;
        case RBP:
            printf("%%rbp");
            break;
        case RSI:
            printf("%%rsi");
            break;
        case RDI:
            printf("%%rdi");
            break;
        case R8:
        case R9:
        case R10:
        case R11:
        case R12:
        case R13:
        case R14:
            printf("%%r%d", y86_register);
            break;
        case NOREG:
            break;
    }
}

void print_cmov(y86_cmov_t cmov);

void print_cmov(y86_cmov_t cmov) {
    switch (cmov) {
        case RRMOVQ:
            printf("rrmovq");
            break;
        case CMOVLE:
            printf("cmovle");
            break;
        case CMOVL:
            printf("cmovl");
            break;
        case CMOVE:
            printf("cmove");
            break;
        case CMOVNE:
            printf("cmovne");
            break;
        case CMOVGE:
            printf("cmovge");
            break;
        case CMOVG:
            printf("cmovg");
            break;
        case BADCMOV:
            break;
    }
}

void print_opq(y86_op_t opq);

void print_opq(y86_op_t opq) {
    switch (opq) {
        case ADD:
            printf("addq");
            break;
        case SUB:
            printf("subq");
            break;
        case AND:
            printf("andq");
            break;
        case XOR:
            printf("xorq");
            break;
        case BADOP:
            break;
    }
}

void print_jump(y86_jump_t jump);

void print_jump(y86_jump_t jump) {
    switch (jump) {
        case JMP:
            printf("jmp");
            break;
        case JLE:
            printf("jle");
            break;
        case JL:
            printf("jl");
            break;
        case JE:
            printf("je");
            break;
        case JNE:
            printf("jne");
            break;
        case JGE:
            printf("jge");
            break;
        case JG:
            printf("jg");
            break;
        case BADJUMP:
            break;
    }
}

void disassemble(y86_inst_t *inst) {
    if (inst == NULL) {
        return;
    }

    switch (inst->icode) {
        case HALT:
            printf("halt");
            break;
        case NOP:
            printf("nop");
            break;
        case CMOV:
            print_cmov(inst->ifun.cmov);
            printf(" ");
            print_register(inst->ra);
            printf(", ");
            print_register(inst->rb);
            break;
        case IRMOVQ:
            printf("irmovq 0x%lx, ", inst->valC.v);
            print_register(inst->rb);
            break;
        case RMMOVQ:
            printf("rmmovq ");
            print_register(inst->ra);
            printf(", ");
            printf("0x%lx", inst->valC.d);
            if (inst->rb != NOREG) {
                printf("(");
                print_register(inst->rb);
                printf(")");
            }
            break;
        case MRMOVQ:
            printf("mrmovq 0x%lx", inst->valC.d);
            if (inst->rb != NOREG) {
                printf("(");
                print_register(inst->rb);
                printf(")");
            }
            printf(", ");
            print_register(inst->ra);
            break;
        case OPQ:
            print_opq(inst->ifun.op);
            printf(" ");
            print_register(inst->ra);
            printf(", ");
            print_register(inst->rb);
            break;
        case JUMP:
            print_jump(inst->ifun.jump);
            printf(" 0x%lx", inst->valC.dest);
            break;
        case CALL:
            printf("call 0x%lx", inst->valC.dest);
            break;
        case RET:
            printf("ret");
            break;
        case PUSHQ:
            printf("pushq ");
            print_register(inst->ra);
            break;
        case POPQ:
            printf("popq ");
            print_register(inst->ra);
            break;
        case IOTRAP:
            printf("iotrap %d", inst->ifun.trap);
            break;
        default:
            break;
    }
}

void disassemble_code(byte_t *memory, elf_phdr_t *phdr, elf_hdr_t *hdr) {
    if (memory == NULL || phdr == NULL || hdr == NULL) {
        printf("Failed to disassemble code\n");
        return;
    }

    y86_t cpu = {0};
    y86_inst_t ins = {0};
    uint32_t p_vaddr_end = phdr->p_vaddr + phdr->p_size;
    cpu.pc = phdr->p_vaddr;

    printf("  0x%03lx:%*s| .pos 0x%03lx code\n", cpu.pc, 31, "", cpu.pc);

    while (cpu.pc < p_vaddr_end) {
        if (cpu.pc == hdr->e_entry) {
            printf("  0x%03lx:%*s| _start:\n", cpu.pc, 31, "");
        }

        ins = fetch(&cpu, memory);

        if (ins.icode == INVALID) {
            printf("Invalid opcode: 0x%02x\n", memory[cpu.pc]);
            break;
        }

        printf("  0x%03lx: ", cpu.pc);

        int bytes_printed = 1;
        printf("%02x ", ins.icode << 4 | ins.ifun.b);

        if (has_register_byte(ins)) {
            printf("%02x ", (ins.ra << 4 | ins.rb));
            bytes_printed++;
        }

        if (has_valc_bytes(ins)) {
            for (int i = 0; i < 8; i++) {
                printf("%02lx ", (ins.valC.d >> (8 * i)) & 0xFF);
            }

            bytes_printed += 8;
        }

        /* Print #x spaces to align the "|" */
        printf("%*s|   ", 30 - 3 * bytes_printed, "");
        disassemble(&ins);
        printf("\n");

        cpu.pc = ins.valP;
    }

    printf("\n");
}

void disassemble_data(byte_t *memory, elf_phdr_t *phdr) {
    if (memory == NULL || phdr == NULL) {
        printf("Failed to disassemble data\n");
        return;
    }

    uint32_t pc = phdr->p_vaddr;
    uint32_t p_vaddr_end = phdr->p_vaddr + phdr->p_size;

    printf("  0x%03x:%*s| .pos 0x%03x data\n", pc, 31, "", pc);

    while (pc < p_vaddr_end) {
        uint64_t byte_string = 0;

        printf("  0x%03x: ", pc);
        for (int i = 0; i < 8; i++) {
            byte_string |= (uint64_t)(memory[pc + i]) << (8 * i);
            printf("%02x ", memory[pc + i]);
        }
        printf("      |   .quad 0x%lx\n", byte_string);

        pc += 8;
    }

    printf("\n");
}

void disassemble_rodata(byte_t *memory, elf_phdr_t *phdr) {
    if (memory == NULL || phdr == NULL) {
        printf("Failed to disassemble read-only data\n");
        return;
    }

    uint32_t pc = phdr->p_vaddr;
    uint32_t p_vaddr_end = phdr->p_vaddr + phdr->p_size;

    printf("  0x%03x:%*s| .pos 0x%03x rodata\n", pc, 31, "", pc);

    uint32_t string_address = pc; /* Address of current .string */
    uint32_t char_count = 0;

    while (pc < p_vaddr_end) {
        bool on_line_start = (char_count + 1) % 10 == 1;
        bool on_null_terminator = memory[pc] == 0x00;
        bool break_line = on_null_terminator || (char_count + 1) % 10 == 0;

        if (on_line_start) {
            printf("  0x%03x: ", pc);
        }

        printf("%02x ", memory[pc]);

        /* If reached end of line, print #x spaces to align the "|" */
        if (break_line) {
            int spaces = 27 - (char_count % 10) * 3;
            printf("%*s| ", spaces, "");
        }

        /* Prints out the full string if reached the end of the first line */
        if (break_line && char_count <= 9) {
            byte_t *current_byte = &memory[string_address];

            printf("  .string \"");
            while (*current_byte != 0x00) {
                printf("%c", *current_byte);
                current_byte++;
            }
            printf("\"");
        }

        if (break_line) {
            printf("\n");
        }

        /* Reset char count to indicate new word on null terminator */
        if (on_null_terminator) {
            string_address = pc + 1;
            char_count = 0;
        } else {
            char_count++;
        }

        pc++;
    }

    printf("\n");
}