#include "asm.h"
#include "common/op_str.h"
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

static inline void add_dst_reg(uint32_t* ins, uint16_t reg);
static inline void add_op1_reg(uint32_t* ins, uint16_t reg);
static inline void add_op2_reg(uint32_t* ins, uint16_t reg);
static inline void set_is_imm(uint32_t* ins);
static inline void set_mov_is_memory(uint32_t* ins);
static inline void set_mov_addr_is_reg(uint32_t* ins);
static inline void set_mov_is_store(uint32_t* ins);
static inline uint16_t linety_arithm_ins(LineTy ty);

uint16_t assemble_to_binary(uint16_t* out, const Line* lines, size_t lines_size)
{
    const bool debug = false;

    uint16_t ip = 0;

    typedef struct {
        int label;
        uint16_t ptr;
    } UnresolvedLabel;
    typedef struct {
        int label;
        uint16_t ip;
    } ResolvedLabel;

    UnresolvedLabel* unres_labels = malloc(sizeof(UnresolvedLabel) * 64);
    size_t unres_labels_size = 0;

    ResolvedLabel* res_labels = malloc(sizeof(ResolvedLabel) * 64);
    size_t res_labels_size = 0;

#define ADD_LABEL(LABEL)                                                       \
    {                                                                          \
        unres_labels[unres_labels_size++] = (UnresolvedLabel) { LABEL, ip };   \
        out[ip++] = 0;                                                         \
    }

    if (debug) {
        printf("assembling...\n");
        printf("ip op    n data...\n");
    }

    for (size_t i = 0; i < lines_size; ++i) {
        bool is_label = false;
        bool is_data = false;

        const Line* line = &lines[i];
        uint16_t ins_ip = ip;
        switch (line->ty) {
            case LineTy_Label: {
                res_labels[res_labels_size++]
                    = (ResolvedLabel) { line->op1.label, ip * 2 };

                is_label = true;
                break;
            }
            case LineTy_DataImm: {
                out[ip++] = line->op1.imm;

                is_data = true;
                break;
            }
            case LineTy_DataLabel: {
                ADD_LABEL(line->op1.label);

                is_data = true;
                break;
            }
            case LineTy_Nop: {
                out[ip++] = Op_Nop;
                break;
            }
            case LineTy_Hlt: {
                out[ip++] = Op_Hlt;
                break;
            }
            case LineTy_Jmp_Reg: {
                uint16_t op1 = line->op1.reg;

                uint32_t ins = Op_Jmp;
                add_op1_reg(&ins, op1);
                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Jmp_Imm: {
                uint16_t op1 = line->op1.imm;

                uint32_t ins = Op_Jmp;
                set_is_imm(&ins);
                out[ip++] = (uint16_t)ins;
                out[ip++] = op1;
                break;
            }
            case LineTy_Jmp_Label: {
                int op1 = line->op1.label;

                uint32_t ins = Op_Jmp;
                set_is_imm(&ins);
                out[ip++] = (uint16_t)ins;
                ADD_LABEL(op1);
                break;
            }
            case LineTy_Jnz_Reg: {
                uint16_t op1 = line->op1.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Jnz;
                add_op1_reg(&ins, op1);
                add_op2_reg(&ins, op2);

                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Jnz_Imm: {
                uint16_t op1 = line->op1.reg;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Jnz;
                set_is_imm(&ins);
                add_op1_reg(&ins, op1);

                out[ip++] = (uint16_t)ins;
                out[ip++] = op2;
                break;
            }
            case LineTy_Jnz_Label: {
                uint16_t op1 = line->op1.reg;
                int op2 = line->op2.label;

                uint32_t ins = Op_Jnz;
                set_is_imm(&ins);
                add_op1_reg(&ins, op1);

                out[ip++] = (uint16_t)ins;
                ADD_LABEL(op2);
                break;
            }
            case LineTy_Cmp_Reg: {
                uint16_t op1 = line->op1.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Cmp;
                add_op1_reg(&ins, op1);
                add_op2_reg(&ins, op2);

                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Cmp_Imm: {
                uint16_t op1 = line->op1.reg;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Cmp;
                set_is_imm(&ins);
                add_op1_reg(&ins, op1);

                out[ip++] = (uint16_t)ins;
                out[ip++] = op2;
                break;
            }
            case LineTy_Cmp_Label: {
                uint16_t op1 = line->op1.reg;
                int op2 = line->op2.label;

                uint32_t ins = Op_Cmp;
                set_is_imm(&ins);
                add_op1_reg(&ins, op1);

                out[ip++] = (uint16_t)ins;
                ADD_LABEL(op2);
                break;
            }
            case LineTy_Mov8_Reg_Reg: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov8;
                ins |= (op2 & 0xfu) << 7;
                ins |= (dst & 0xfu) << 12;

                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Mov8_Reg_Imm: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Mov8;
                set_is_imm(&ins);
                ins |= (dst & 0xfu) << 12;

                out[ip++] = (uint16_t)ins;
                out[ip++] = op2;
                break;
            }
            case LineTy_Mov8_Reg_MemReg: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov8;
                add_op2_reg(&ins, op2);
                set_mov_is_memory(&ins);
                set_mov_addr_is_reg(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = line->offset;
                break;
            }
            case LineTy_Mov8_Reg_MemImm: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Mov8;
                set_is_imm(&ins);
                set_mov_is_memory(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = op2;
                break;
            }
            case LineTy_Mov8_MemReg_Reg: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov8;
                add_op2_reg(&ins, op2);
                set_mov_is_memory(&ins);
                set_mov_addr_is_reg(&ins);
                set_mov_is_store(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = line->offset;
                break;
            }
            case LineTy_Mov8_MemReg_Imm: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov8;
                set_is_imm(&ins);
                set_mov_is_memory(&ins);
                set_mov_addr_is_reg(&ins);
                set_mov_is_store(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = line->offset;
                out[ip++] = op2;
                break;
            }
            case LineTy_Mov8_MemImm_Reg: {
                uint16_t dst = line->dst.imm;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov8;
                add_op2_reg(&ins, op2);
                set_mov_is_memory(&ins);
                set_mov_is_store(&ins);

                out[ip++] = (uint16_t)ins;
                out[ip++] = dst;
                break;
            }
            case LineTy_Mov8_MemImm_Imm: {
                uint16_t dst = line->dst.imm;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Mov8;
                set_is_imm(&ins);
                set_mov_is_memory(&ins);
                set_mov_is_store(&ins);

                out[ip++] = (uint16_t)ins;
                out[ip++] = dst;
                out[ip++] = op2;
                break;
            }
            case LineTy_Mov16_Reg_Reg: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov16;
                ins |= (op2 & 0xfu) << 7;
                ins |= (dst & 0xfu) << 12;

                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Mov16_Reg_Imm: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Mov16;
                set_is_imm(&ins);
                ins |= (dst & 0xfu) << 12;

                out[ip++] = (uint16_t)ins;
                out[ip++] = op2;
                break;
            }
            case LineTy_Mov16_Reg_Label: {
                uint16_t dst = line->dst.reg;
                int op2 = line->op2.label;

                uint32_t ins = Op_Mov16;
                set_is_imm(&ins);
                ins |= (dst & 0xfu) << 12;

                out[ip++] = (uint16_t)ins;
                ADD_LABEL(op2);
                break;
            }
            case LineTy_Mov16_Reg_MemReg: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov16;
                add_op2_reg(&ins, op2);
                set_mov_is_memory(&ins);
                set_mov_addr_is_reg(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = line->offset;
                break;
            }
            case LineTy_Mov16_Reg_MemImm: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Mov16;
                set_is_imm(&ins);
                set_mov_is_memory(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = op2;
                break;
            }
            case LineTy_Mov16_Reg_MemLabel: {
                uint16_t dst = line->dst.reg;

                uint32_t ins = Op_Mov16;
                set_is_imm(&ins);
                set_mov_is_memory(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                ADD_LABEL(line->op2.label);
                break;
            }
            case LineTy_Mov16_MemReg_Reg: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov16;
                add_op2_reg(&ins, op2);
                set_mov_is_memory(&ins);
                set_mov_addr_is_reg(&ins);
                set_mov_is_store(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = line->offset;
                break;
            }
            case LineTy_Mov16_MemReg_Imm: {
                uint16_t dst = line->dst.reg;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Mov16;
                set_is_imm(&ins);
                set_mov_is_memory(&ins);
                set_mov_addr_is_reg(&ins);
                set_mov_is_store(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = line->offset;
                out[ip++] = op2;
                break;
            }
            case LineTy_Mov16_MemImm_Reg: {
                uint16_t dst = line->dst.imm;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov16;
                add_op2_reg(&ins, op2);
                set_mov_is_memory(&ins);
                set_mov_is_store(&ins);

                out[ip++] = (uint16_t)ins;
                out[ip++] = dst;
                break;
            }
            case LineTy_Mov16_MemImm_Imm: {
                uint16_t dst = line->dst.imm;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = Op_Mov16;
                set_is_imm(&ins);
                set_mov_is_memory(&ins);
                set_mov_is_store(&ins);

                out[ip++] = (uint16_t)ins;
                out[ip++] = dst;
                out[ip++] = op2;
                break;
            }
            case LineTy_Mov16_MemLabel_Reg: {
                int dst = line->dst.label;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = Op_Mov16;
                add_op2_reg(&ins, op2);
                set_mov_is_memory(&ins);
                set_mov_is_store(&ins);

                out[ip++] = (uint16_t)ins;
                ADD_LABEL(dst);
                break;
            }
            case LineTy_In_Reg: {
                uint16_t dst = line->dst.reg;
                uint16_t op1 = line->op1.reg;

                uint32_t ins = Op_In;
                add_op1_reg(&ins, op1);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_In_Imm: {
                uint16_t dst = line->dst.reg;
                uint16_t op1 = line->op1.imm;

                uint32_t ins = Op_In;
                set_is_imm(&ins);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = op1;
                break;
            }
            case LineTy_Call_Reg: {
                uint16_t op1 = line->op1.reg;

                uint32_t ins = Op_Call;
                add_op1_reg(&ins, op1);

                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Call_Imm: {
                uint16_t op1 = line->op1.imm;

                uint32_t ins = Op_Call;
                set_is_imm(&ins);

                out[ip++] = (uint16_t)ins;
                out[ip++] = op1;
                break;
            }
            case LineTy_Call_Label: {
                int op1 = line->op1.label;

                uint32_t ins = Op_Call;
                set_is_imm(&ins);

                out[ip++] = (uint16_t)ins;
                ADD_LABEL(op1);
                break;
            }
            case LineTy_Ret: {
                uint32_t ins = Op_Ret;
                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Lit_Reg: {
                uint16_t op1 = line->op1.reg;

                uint32_t ins = Op_Lit;
                add_op1_reg(&ins, op1);

                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Lit_Imm: {
                uint16_t op1 = line->op1.imm;

                uint32_t ins = Op_Lit;
                set_is_imm(&ins);

                out[ip++] = (uint16_t)ins;
                out[ip++] = op1;
                break;
            }
            case LineTy_Lit_Label: {
                int op1 = line->op1.label;

                uint32_t ins = Op_Lit;
                set_is_imm(&ins);

                out[ip++] = (uint16_t)ins;
                ADD_LABEL(op1);
                break;
            }
            case LineTy_Int: {
                uint16_t int_id = line->op1.imm & 0xff;

                uint32_t ins = Op_Int;
                ins |= (uint16_t)(int_id << 8);
                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_IRet: {
                uint32_t ins = Op_Ret;
                ins |= 1 << 6;
                out[ip++] = (uint16_t)ins;
                break;
            }
            case LineTy_Or_Imm:
            case LineTy_Xor_Imm:
            case LineTy_And_Imm:
            case LineTy_Shl_Imm:
            case LineTy_RShl_Imm:
            case LineTy_Shr_Imm:
            case LineTy_RShr_Imm:
            case LineTy_Add_Imm:
            case LineTy_Sub_Imm:
            case LineTy_RSub_Imm:
            case LineTy_Mul_Imm:
            case LineTy_IMul_Imm:
            case LineTy_Div_Imm:
            case LineTy_IDiv_Imm:
            case LineTy_RDiv_Imm:
            case LineTy_RIDiv_Imm:
            case LineTy_Mod_Imm:
            case LineTy_RMod_Imm: {
                uint16_t dst = line->dst.reg;
                uint16_t op1 = line->op1.reg;
                uint16_t op2 = line->op2.imm;

                uint32_t ins = linety_arithm_ins(line->ty);
                set_is_imm(&ins);
                add_op1_reg(&ins, op1);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                out[ip++] = op2;
                break;
            }
            case LineTy_Or_Reg:
            case LineTy_Xor_Reg:
            case LineTy_And_Reg:
            case LineTy_Shl_Reg:
            case LineTy_RShl_Reg:
            case LineTy_Shr_Reg:
            case LineTy_RShr_Reg:
            case LineTy_Add_Reg:
            case LineTy_Sub_Reg:
            case LineTy_RSub_Reg:
            case LineTy_Mul_Reg:
            case LineTy_IMul_Reg:
            case LineTy_Div_Reg:
            case LineTy_IDiv_Reg:
            case LineTy_RDiv_Reg:
            case LineTy_RIDiv_Reg:
            case LineTy_Mod_Reg:
            case LineTy_RMod_Reg: {
                uint16_t dst = line->dst.reg;
                uint16_t op1 = line->op1.reg;
                uint16_t op2 = line->op2.reg;

                uint32_t ins = linety_arithm_ins(line->ty);
                add_op2_reg(&ins, op2);
                add_op1_reg(&ins, op1);
                add_dst_reg(&ins, dst);

                out[ip++] = (uint16_t)ins;
                break;
            }
        }

        if (debug) {
            if (!is_label) {
                printf("%02x %-5s %d",
                    ins_ip * 2,
                    is_data ? "data" : op_str(out[ins_ip] & 0x3f),
                    ip - ins_ip);
                for (uint16_t i = 0; i < ip - ins_ip; ++i) {
                    printf(" %02x %c%c%c%c %c%c%c%c  %02x %c%c%c%c %c%c%c%c ",
                        out[ins_ip + i] & 0xff,
                        fmt_binary(out[ins_ip + i] & 0xff),
                        out[ins_ip + i] >> 8,
                        fmt_binary(out[ins_ip + i] >> 8));
                }
                printf("\n");
            }
        }
    }

    if (debug) {
        printf("resolving...\n");
        printf(" l ip  v  data\n");
    }
    for (size_t i = 0; i < unres_labels_size; ++i) {
        bool found = false;
        for (size_t j = 0; j < res_labels_size; ++j) {
            if (res_labels[j].label == unres_labels[i].label) {
                out[unres_labels[i].ptr] = res_labels[j].ip;
                found = true;

                if (debug) {
                    printf(
                        "%2d %02x %02x  %02x %c%c%c%c %c%c%c%c  %02x %c%c%c%c "
                        "%c%c%c%c\n",
                        res_labels[j].label,
                        unres_labels[i].ptr * 2,
                        res_labels[j].ip,
                        out[unres_labels[i].ptr] & 0xff,
                        fmt_binary(out[unres_labels[i].ptr] & 0xff),
                        out[unres_labels[i].ptr] >> 8,
                        fmt_binary(out[unres_labels[i].ptr] >> 8));
                }
                break;
            }
        }
        if (!found) {
            fprintf(stderr,
                "warning: label '%d' could not be resolved\n",
                unres_labels[i].label);
        }
    }
    if (debug) {
        printf("done!\n");
    }

    free(unres_labels);
    free(res_labels);

    return ip * 2;
}

static inline void add_dst_reg(uint32_t* ins, uint16_t reg)
{
    *ins |= (reg & 0x7u) << 13;
}
static inline void add_op1_reg(uint32_t* ins, uint16_t reg)
{
    *ins |= (reg & 0x7u) << 10;
}
static inline void add_op2_reg(uint32_t* ins, uint16_t reg)
{
    *ins |= (reg & 0x7u) << 7;
}
static inline void set_is_imm(uint32_t* ins)
{
    *ins |= 1 << 6;
}
static inline void set_mov_is_memory(uint32_t* ins)
{
    *ins |= 1 << 10;
}
static inline void set_mov_addr_is_reg(uint32_t* ins)
{
    *ins |= 1 << 11;
}
static inline void set_mov_is_store(uint32_t* ins)
{
    *ins |= 1 << 12;
}

static inline uint16_t linety_arithm_ins(LineTy ty)
{
    switch (ty) {
        case LineTy_Or_Imm:
        case LineTy_Or_Reg:
            return Op_Or;
        case LineTy_Xor_Imm:
        case LineTy_Xor_Reg:
            return Op_Xor;
        case LineTy_And_Imm:
        case LineTy_And_Reg:
            return Op_And;
        case LineTy_Shl_Imm:
        case LineTy_Shl_Reg:
            return Op_Shl;
        case LineTy_RShl_Imm:
        case LineTy_RShl_Reg:
            return Op_RShl;
        case LineTy_Shr_Imm:
        case LineTy_Shr_Reg:
            return Op_Shr;
        case LineTy_RShr_Imm:
        case LineTy_RShr_Reg:
            return Op_RShr;
        case LineTy_Add_Imm:
        case LineTy_Add_Reg:
            return Op_Add;
        case LineTy_Sub_Imm:
        case LineTy_Sub_Reg:
            return Op_Sub;
        case LineTy_RSub_Imm:
        case LineTy_RSub_Reg:
            return Op_RSub;
        case LineTy_Mul_Imm:
        case LineTy_Mul_Reg:
            return Op_Mul;
        case LineTy_IMul_Imm:
        case LineTy_IMul_Reg:
            return Op_IMul;
        case LineTy_Div_Imm:
        case LineTy_Div_Reg:
            return Op_Div;
        case LineTy_IDiv_Imm:
        case LineTy_IDiv_Reg:
            return Op_IDiv;
        case LineTy_RDiv_Imm:
        case LineTy_RDiv_Reg:
            return Op_RDiv;
        case LineTy_RIDiv_Imm:
        case LineTy_RIDiv_Reg:
            return Op_RIDiv;
        case LineTy_Mod_Imm:
        case LineTy_Mod_Reg:
            return Op_Mod;
        case LineTy_RMod_Imm:
        case LineTy_RMod_Reg:
            return Op_RMod;
        default:
            fprintf(stderr, "error: line type '%d' not handled\n", ty);
            exit(1);
    }
}