#include "../../Headers/Assembler/SagittariusAssembler.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* Lexer / scanner helpers */
bool Sag_NextWord(FILE* f, Sag_Str* out) {
    if (!f || !out) return false;
    
    int c;
    // Skip whitespace and comments
    while (1) {
        c = fgetc(f);
        if (c == EOF) return false;
        
        if (c == ' ' || c == '\t' || c == '\n' || c == '\r') {
            continue;
        }
        if (c == ';' || c == '#') {
            // Skip until end of line
            while (1) {
                c = fgetc(f);
                if (c == EOF || c == '\n' || c == '\r') break;
            }
            continue;
        }
        break;
    }
    
    // Read word/token
    uint64_t cap = 32;
    uint64_t len = 0;
    char *buf = malloc(cap);
    if (!buf) return false;
    
    if (c == '"') {
        // Read string literal
        buf[len++] = c;
        while (1) {
            c = fgetc(f);
            if (c == EOF) {
                free(buf);
                return false;
            }
            if (len + 2 >= cap) {
                cap *= 2;
                buf = realloc(buf, cap);
            }
            buf[len++] = c;
            if (c == '"') {
                break;
            }
        }
    } else {
        // Read normal word until whitespace, comment, or EOF
        buf[len++] = c;
        while (1) {
            c = fgetc(f);
            if (c == EOF) break;
            if (c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == ';' || c == '#') {
                ungetc(c, f);
                break;
            }
            if (len + 2 >= cap) {
                cap *= 2;
                buf = realloc(buf, cap);
            }
            buf[len++] = c;
        }
    }
    
    buf[len] = '\0';
    out->head = buf;
    out->start = 0;
    out->length = len;
    return true;
}

bool Sag_MatchStr(Sag_Str* str, char** match_list, uint64_t match_count, uint64_t* out_index) {
    if (!str || !str->head) return false;
    for (uint64_t i = 0; i < match_count; ++i) {
        if (strlen(match_list[i]) == str->length &&
            strncmp(str->head + str->start, match_list[i], str->length) == 0) {
            if (out_index) *out_index = i;
            return true;
        }
    }
    return false;
}

static uint64_t get_instruction_arg_count(const char *name) {
    if (strcmp(name, "add") == 0 || strcmp(name, "sub") == 0 || strcmp(name, "mul") == 0 ||
        strcmp(name, "div") == 0 || strcmp(name, "mod") == 0 || strcmp(name, "pow") == 0) {
        return 4;
    }
    if (strcmp(name, "sin") == 0 || strcmp(name, "cos") == 0 || strcmp(name, "tan") == 0 ||
        strcmp(name, "sinh") == 0 || strcmp(name, "cosh") == 0 || strcmp(name, "tanh") == 0 ||
        strcmp(name, "asin") == 0 || strcmp(name, "acos") == 0 || strcmp(name, "atan") == 0 ||
        strcmp(name, "abs") == 0 || strcmp(name, "exp") == 0) {
        return 3;
    }
    if (strcmp(name, "cvt") == 0) return 4;
    if (strcmp(name, "set") == 0) return 3;
    if (strncmp(name, "set.", 4) == 0) return 2; // alias!
    if (strcmp(name, "mv") == 0) return 3;
    if (strcmp(name, "cp") == 0) return 3;
    if (strcmp(name, "save") == 0) return 3;
    if (strcmp(name, "load") == 0) return 3;
    if (strcmp(name, "jmp") == 0) return 3;
    if (strcmp(name, "jmp_if") == 0) return 4;
    if (strcmp(name, "call") == 0) return 4;
    if (strcmp(name, "ret") == 0) return 1;
    if (strcmp(name, "cmp") == 0) return 5;
    if (strcmp(name, "mathv") == 0) return 6;
    if (strcmp(name, "halt") == 0) return 0;
    if (strcmp(name, "syscall") == 0) return 2;
    if (strcmp(name, "tsyscall") == 0) return 3;
    return 0xFFFFFFFFFFFFFFFFULL;
}

bool Sag_Scan(FILE* f, Sag_IntermediateProgram* out) {
    if (!f || !out) return false;
    
    out->insts = NULL;
    out->inst_count = 0;
    out->data = NULL;
    out->data_count = 0;
    out->Consts = NULL;
    out->Consts_count = 0;
    
    typedef enum Section { SEC_CODE, SEC_DATA, SEC_CONST } Section;
    Section active_sec = SEC_CODE;
    Sag_Str* pending_label = NULL;
    
    Sag_Str word;
    while (Sag_NextWord(f, &word)) {
        // Check for sections
        if (strcmp(word.head, ".code:") == 0) {
            active_sec = SEC_CODE;
            free(word.head);
            continue;
        }
        if (strcmp(word.head, ".data:") == 0) {
            active_sec = SEC_DATA;
            free(word.head);
            continue;
        }
        if (strcmp(word.head, ".const:") == 0) {
            active_sec = SEC_CONST;
            free(word.head);
            continue;
        }
        
        // Check if label definition
        if (word.length > 1 && word.head[word.length - 1] == ':') {
            word.head[word.length - 1] = '\0';
            word.length--;
            pending_label = malloc(sizeof(Sag_Str));
            *pending_label = word;
            continue;
        }
        
        if (active_sec == SEC_CODE) {
            uint64_t arg_count = get_instruction_arg_count(word.head);
            if (arg_count == 0xFFFFFFFFFFFFFFFFULL) {
                // Unknown instruction name
                printf("Error: Unknown instruction '%s'\n", word.head);
                free(word.head);
                return false;
            }
            
            Sag_IntermediateInst ii;
            ii.inst.data = 0;
            ii.label = pending_label;
            pending_label = NULL; // Consumed
            ii.arg_count = arg_count;
            
            ii.args = malloc((arg_count + 1) * sizeof(Sag_Str));
            ii.args[0] = word; // First is instruction name
            
            for (uint64_t a = 1; a <= arg_count; ++a) {
                if (!Sag_NextWord(f, &ii.args[a])) {
                    printf("Error: Missing argument for '%s'\n", word.head);
                    free(ii.args);
                    return false;
                }
            }
            
            out->insts = realloc(out->insts, (out->inst_count + 1) * sizeof(Sag_IntermediateInst));
            out->insts[out->inst_count] = ii;
            out->inst_count++;
            
        } else if (active_sec == SEC_DATA) {
            // Read next 2 words (type and value)
            Sag_Str type_str, val_str;
            if (!Sag_NextWord(f, &type_str) || !Sag_NextWord(f, &val_str)) {
                printf("Error: Incomplete data section entry\n");
                return false;
            }
            
            out->data = realloc(out->data, (out->data_count + 3) * sizeof(Sag_Str));
            out->data[out->data_count] = word;       // name
            out->data[out->data_count + 1] = type_str; // type
            out->data[out->data_count + 2] = val_str;  // value
            out->data_count += 3;
            
        } else if (active_sec == SEC_CONST) {
            // Read next 1 word (value)
            Sag_Str val_str;
            if (!Sag_NextWord(f, &val_str)) {
                printf("Error: Incomplete const section entry\n");
                return false;
            }
            
            out->Consts = realloc(out->Consts, (out->Consts_count + 2) * sizeof(Sag_Str));
            out->Consts[out->Consts_count] = word;      // name
            out->Consts[out->Consts_count + 1] = val_str; // value
            out->Consts_count += 2;
        }
    }
    return true;
}

bool Sag_Combine(Sag_IntermediateProgram* L, Sag_IntermediateProgram* R, Sag_IntermediateProgram* out) {
    if (!L || !R || !out) return false;
    
    out->inst_count = L->inst_count + R->inst_count;
    if (out->inst_count > 0) {
        out->insts = malloc(out->inst_count * sizeof(Sag_IntermediateInst));
        if (L->inst_count > 0) memcpy(out->insts, L->insts, L->inst_count * sizeof(Sag_IntermediateInst));
        if (R->inst_count > 0) memcpy(out->insts + L->inst_count, R->insts, R->inst_count * sizeof(Sag_IntermediateInst));
    } else {
        out->insts = NULL;
    }
    
    out->data_count = L->data_count + R->data_count;
    if (out->data_count > 0) {
        out->data = malloc(out->data_count * sizeof(Sag_Str));
        if (L->data_count > 0) memcpy(out->data, L->data, L->data_count * sizeof(Sag_Str));
        if (R->data_count > 0) memcpy(out->data + L->data_count, R->data, R->data_count * sizeof(Sag_Str));
    } else {
        out->data = NULL;
    }
    
    out->Consts_count = L->Consts_count + R->Consts_count;
    if (out->Consts_count > 0) {
        out->Consts = malloc(out->Consts_count * sizeof(Sag_Str));
        if (L->Consts_count > 0) memcpy(out->Consts, L->Consts, L->Consts_count * sizeof(Sag_Str));
        if (R->Consts_count > 0) memcpy(out->Consts + L->Consts_count, R->Consts, R->Consts_count * sizeof(Sag_Str));
    } else {
        out->Consts = NULL;
    }
    
    return true;
}

/* Base64 & String parsing helpers */
static inline int b64_char_val(char c) {
    if (c >= 'A' && c <= 'Z') return c - 'A';
    if (c >= 'a' && c <= 'z') return c - 'a' + 26;
    if (c >= '0' && c <= '9') return c - '0' + 52;
    if (c == '+') return 62;
    if (c == '/') return 63;
    return -1;
}

static inline uint8_t* base64_decode(const char* in, size_t* out_len) {
    size_t len = strlen(in);
    size_t padding = 0;
    if (len > 0 && in[len - 1] == '=') padding++;
    if (len > 1 && in[len - 2] == '=') padding++;
    
    *out_len = (len * 3) / 4 - padding;
    uint8_t* out = malloc(*out_len);
    if (!out) return NULL;
    
    size_t j = 0;
    uint32_t val = 0;
    int valb = -8;
    for (size_t i = 0; i < len; ++i) {
        char c = in[i];
        if (c == '=') break;
        int v = b64_char_val(c);
        if (v == -1) continue;
        val = (val << 6) | v;
        valb += 6;
        if (valb >= 0) {
            out[j++] = (val >> valb) & 0xFF;
            valb -= 8;
        }
    }
    return out;
}

static inline uint8_t* parse_string_bytes(const char* in, size_t* out_len) {
    size_t len = strlen(in);
    size_t start = 0;
    size_t end = len;
    if (len >= 2 && in[0] == '"' && in[len - 1] == '"') {
        start = 1;
        end = len - 1;
    }
    
    uint8_t* out = malloc(len + 1);
    size_t j = 0;
    for (size_t i = start; i < end; ++i) {
        if (in[i] == '\\' && i + 1 < end) {
            i++;
            switch (in[i]) {
                case 'n': out[j++] = '\n'; break;
                case 't': out[j++] = '\t'; break;
                case 'r': out[j++] = '\r'; break;
                case '\"': out[j++] = '\"'; break;
                case '\\': out[j++] = '\\'; break;
                default: out[j++] = in[i]; break;
            }
        } else {
            out[j++] = in[i];
        }
    }
    out[j++] = '\0';
    *out_len = j;
    return out;
}

static inline uint8_t* read_file_bytes(const char* path, size_t* out_len) {
    // Strip quotes if path has them
    size_t len = strlen(path);
    char* clean_path = malloc(len + 1);
    size_t start = 0, end = len;
    if (len >= 2 && path[0] == '"' && path[len-1] == '"') {
        start = 1;
        end = len - 1;
    }
    size_t j = 0;
    for (size_t i = start; i < end; ++i) clean_path[j++] = path[i];
    clean_path[j] = '\0';
    
    FILE* f = fopen(clean_path, "rb");
    free(clean_path);
    if (!f) return NULL;
    
    fseek(f, 0, SEEK_END);
    long sz = ftell(f);
    if (sz < 0) { fclose(f); return NULL; }
    fseek(f, 0, SEEK_SET);
    
    uint8_t* out = malloc(sz);
    if (out) {
        size_t read_bytes = fread(out, 1, sz, f);
        *out_len = read_bytes;
    }
    fclose(f);
    return out;
}

typedef struct LabelMap {
    const char* name;
    uint64_t offset;
} LabelMap;

static inline const char* lookup_const(Sag_IntermediateProgram* prog, const char* name) {
    for (uint64_t i = 0; i < prog->Consts_count; i += 2) {
        if (strcmp(prog->Consts[i].head, name) == 0) {
            return prog->Consts[i+1].head;
        }
    }
    return NULL;
}

static inline bool resolve_token(Sag_IntermediateProgram* prog, LabelMap* labels, size_t label_count, const char* token, int64_t* out_int, double* out_double) {
    const char* resolved = lookup_const(prog, token);
    if (resolved) {
        token = resolved;
    }
    
    for (size_t i = 0; i < label_count; ++i) {
        if (strcmp(labels[i].name, token) == 0) {
            *out_int = labels[i].offset;
            *out_double = (double)labels[i].offset;
            return true;
        }
    }
    
    char* end;
    if (strncmp(token, "0x", 2) == 0 || strncmp(token, "0X", 2) == 0) {
        *out_int = strtoll(token, &end, 16);
        *out_double = (double)*out_int;
        if (*end == '\0') return true;
    }
    
    double d = strtod(token, &end);
    if (*end == '\0') {
        *out_double = d;
        *out_int = (int64_t)d;
        return true;
    }
    
    return false;
}

static inline sagittarius_type parse_type(const char* name) {
    if (strcmp(name, "uint8") == 0) return st_uint8;
    if (strcmp(name, "uint16") == 0) return st_uint16;
    if (strcmp(name, "uint32") == 0) return st_uint32;
    if (strcmp(name, "uint64") == 0) return st_uint64;
    if (strcmp(name, "int8") == 0) return st_int8;
    if (strcmp(name, "int16") == 0) return st_int16;
    if (strcmp(name, "int32") == 0) return st_int32;
    if (strcmp(name, "int64") == 0) return st_int64;
    if (strcmp(name, "single") == 0 || strcmp(name, "float") == 0) return st_single;
    if (strcmp(name, "double") == 0) return st_double;
    return st_uint8;
}

static inline uint8_t parse_math2_op(const char* name) {
    if (strcmp(name, "add") == 0) return sag_math2_add;
    if (strcmp(name, "sub") == 0) return sag_math2_sub;
    if (strcmp(name, "mul") == 0) return sag_math2_mul;
    if (strcmp(name, "div") == 0) return sag_math2_div;
    if (strcmp(name, "mod") == 0) return sag_math2_mod;
    if (strcmp(name, "pow") == 0) return sag_math2_pow;
    return 0;
}

static inline uint8_t parse_math1_op(const char* name) {
    if (strcmp(name, "sin") == 0) return sag_math1_sin;
    if (strcmp(name, "cos") == 0) return sag_math1_cos;
    if (strcmp(name, "tan") == 0) return sag_math1_tan;
    if (strcmp(name, "sinh") == 0) return sag_math1_sinh;
    if (strcmp(name, "cosh") == 0) return sag_math1_cosh;
    if (strcmp(name, "tanh") == 0) return sag_math1_tanh;
    if (strcmp(name, "asin") == 0) return sag_math1_asin;
    if (strcmp(name, "acos") == 0) return sag_math1_acos;
    if (strcmp(name, "atan") == 0) return sag_math1_atan;
    if (strcmp(name, "abs") == 0) return sag_math1_abs;
    if (strcmp(name, "exp") == 0) return sag_math1_exp;
    return 0;
}

bool Sag_Finalize(Sag_IntermediateProgram* intermediate, SagittariusProgram* out) {
    if (!intermediate || !out) return false;
    
    // Pass 1: Calculate Code Labels & instruction sizes
    size_t label_capacity = 32;
    size_t label_count = 0;
    LabelMap* labels = malloc(label_capacity * sizeof(LabelMap));
    
    uint64_t byte_offset = 0;
    for (uint64_t i = 0; i < intermediate->inst_count; ++i) {
        Sag_IntermediateInst* ii = &intermediate->insts[i];
        if (ii->label) {
            if (label_count >= label_capacity) {
                label_capacity *= 2;
                labels = realloc(labels, label_capacity * sizeof(LabelMap));
            }
            labels[label_count].name = ii->label->head;
            labels[label_count].offset = byte_offset;
            label_count++;
        }
        
        // Calculate instruction size
        const char* name = ii->args[0].head;
        uint64_t sz = 8;
        if (strcmp(name, "set") == 0 || strncmp(name, "set.", 4) == 0) {
            sz = 16;
        } else if (strcmp(name, "call") == 0) {
            // Immediate mode if third argument is "1"
            if (ii->arg_count >= 2 && strcmp(ii->args[2].head, "1") == 0) {
                sz = 16;
            }
        }
        byte_offset += sz;
    }
    
    uint64_t total_inst_bytes = byte_offset;
    out->instCount = total_inst_bytes / 8;
    out->instructions = malloc(out->instCount * sizeof(SagittariusInst));
    memset(out->instructions, 0, out->instCount * sizeof(SagittariusInst));
    
    // Pass 2: Resolve Data Section & Data Labels
    uint64_t data_capacity = 1024;
    uint64_t data_size = 0;
    uint8_t* data_buf = malloc(data_capacity);
    
    for (uint64_t i = 0; i < intermediate->data_count; i += 3) {
        const char* d_name = intermediate->data[i].head;
        const char* d_type = intermediate->data[i+1].head;
        const char* d_val = intermediate->data[i+2].head;
        
        // Record label offset relative to program start (which is loaded at 0 in final VM space)
        if (label_count >= label_capacity) {
            label_capacity *= 2;
            labels = realloc(labels, label_capacity * sizeof(LabelMap));
        }
        labels[label_count].name = d_name;
        labels[label_count].offset = total_inst_bytes + data_size;
        label_count++;
        
        size_t entry_len = 0;
        uint8_t* entry_bytes = NULL;
        if (strcmp(d_type, "string") == 0) {
            entry_bytes = parse_string_bytes(d_val, &entry_len);
        } else if (strcmp(d_type, "base64") == 0) {
            entry_bytes = base64_decode(d_val, &entry_len);
        } else if (strcmp(d_type, "file") == 0) {
            entry_bytes = read_file_bytes(d_val, &entry_len);
        }
        
        if (entry_bytes && entry_len > 0) {
            if (data_size + entry_len > data_capacity) {
                while (data_size + entry_len > data_capacity) data_capacity *= 2;
                data_buf = realloc(data_buf, data_capacity);
            }
            memcpy(data_buf + data_size, entry_bytes, entry_len);
            data_size += entry_len;
            free(entry_bytes);
        }
    }
    
    out->data_size = data_size;
    if (data_size > 0) {
        out->data = malloc(data_size);
        memcpy(out->data, data_buf, data_size);
    } else {
        out->data = NULL;
    }
    free(data_buf);
    
    // Pass 3: Assemble & Encode instructions
    uint64_t final_inst_idx = 0;
    for (uint64_t i = 0; i < intermediate->inst_count; ++i) {
        Sag_IntermediateInst* ii = &intermediate->insts[i];
        const char* name = ii->args[0].head;
        
        uint64_t current_inst_offset = final_inst_idx * 8;
        SagittariusInst main_inst;
        main_inst.data = 0;
        
        if (strcmp(name, "add") == 0 || strcmp(name, "sub") == 0 || strcmp(name, "mul") == 0 ||
            strcmp(name, "div") == 0 || strcmp(name, "mod") == 0 || strcmp(name, "pow") == 0) {
            
            uint8_t opcode = math2;
            sagittarius_type type = parse_type(ii->args[1].head);
            uint8_t op = parse_math2_op(name);
            
            int64_t L = 0, R = 0, Dest = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &L, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &R, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[4].head, &Dest, &dummy);
            
            main_inst.data = opcode | ((uint64_t)type << 8) | ((uint64_t)op << 16) | 
                             ((uint64_t)(L & 0xFF) << 24) | ((uint64_t)(R & 0xFF) << 32) | 
                             ((uint64_t)(Dest & 0xFF) << 40);
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "sin") == 0 || strcmp(name, "cos") == 0 || strcmp(name, "tan") == 0 ||
                   strcmp(name, "sinh") == 0 || strcmp(name, "cosh") == 0 || strcmp(name, "tanh") == 0 ||
                   strcmp(name, "asin") == 0 || strcmp(name, "acos") == 0 || strcmp(name, "atan") == 0 ||
                   strcmp(name, "abs") == 0 || strcmp(name, "exp") == 0) {
            
            uint8_t opcode = math1;
            sagittarius_type type = parse_type(ii->args[1].head);
            uint8_t op = parse_math1_op(name);
            
            int64_t L = 0, Dest = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &L, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &Dest, &dummy);
            
            main_inst.data = opcode | ((uint64_t)type << 8) | ((uint64_t)op << 16) | 
                             ((uint64_t)(L & 0xFF) << 24) | ((uint64_t)(Dest & 0xFF) << 32);
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "cvt") == 0) {
            uint8_t opcode = cvt;
            int64_t src_reg = 0, dst_reg = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &src_reg, &dummy);
            sagittarius_type src_type = parse_type(ii->args[2].head);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &dst_reg, &dummy);
            sagittarius_type dst_type = parse_type(ii->args[4].head);
            
            main_inst.data = opcode | ((uint64_t)(src_reg & 0xFF) << 8) | ((uint64_t)src_type << 16) | 
                             ((uint64_t)(dst_reg & 0xFF) << 24) | ((uint64_t)dst_type << 32);
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "set") == 0) {
            uint8_t opcode = set;
            int64_t length = 0, reg = 0, val_int = 0;
            double val_double = 0;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &length, &val_double);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &reg, &val_double);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &val_int, &val_double);
            
            main_inst.data = opcode | ((uint64_t)(length & 0xFF) << 8) | ((uint64_t)(reg & 0xFF) << 16);
            out->instructions[final_inst_idx++] = main_inst;
            
            SagittariusInst payload_inst;
            payload_inst.data = (uint64_t)val_int;
            out->instructions[final_inst_idx++] = payload_inst;
            
        } else if (strncmp(name, "set.", 4) == 0) {
            uint8_t opcode = set;
            sagittarius_type type = parse_type(name + 4);
            uint8_t length = 8;
            switch (type) {
                case st_uint8:  case st_int8:   length = 1; break;
                case st_uint16: case st_int16:  length = 2; break;
                case st_uint32: case st_int32:  case st_single: length = 4; break;
                default: length = 8; break;
            }
            
            int64_t reg = 0, val_int = 0;
            double val_double = 0;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &reg, &val_double);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &val_int, &val_double);
            
            main_inst.data = opcode | ((uint64_t)length << 8) | ((uint64_t)(reg & 0xFF) << 16);
            out->instructions[final_inst_idx++] = main_inst;
            
            SagittariusInst payload_inst;
            payload_inst.data = 0;
            if (type == st_single) {
                float f = (float)val_double;
                uint32_t u32;
                memcpy(&u32, &f, 4);
                payload_inst.data = u32;
            } else if (type == st_double) {
                double d = val_double;
                uint64_t u64;
                memcpy(&u64, &d, 8);
                payload_inst.data = u64;
            } else {
                payload_inst.data = (uint64_t)val_int;
            }
            out->instructions[final_inst_idx++] = payload_inst;
            
        } else if (strcmp(name, "mv") == 0 || strcmp(name, "cp") == 0) {
            uint8_t opcode = (strcmp(name, "mv") == 0) ? mv : cp;
            int64_t length = 0, src_reg = 0, dst_reg = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &length, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &src_reg, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &dst_reg, &dummy);
            
            main_inst.data = opcode | ((uint64_t)(length & 0xFF) << 8) | 
                             ((uint64_t)(src_reg & 0xFF) << 16) | ((uint64_t)(dst_reg & 0xFF) << 24);
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "save") == 0) {
            uint8_t opcode = save;
            int64_t length = 0, src_reg = 0, dest_mem_ptr_reg = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &length, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &src_reg, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &dest_mem_ptr_reg, &dummy);
            
            main_inst.data = opcode | ((uint64_t)(length & 0xFF) << 8) | 
                             ((uint64_t)(src_reg & 0xFF) << 16) | ((uint64_t)(dest_mem_ptr_reg & 0xFF) << 24);
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "load") == 0) {
            uint8_t opcode = load;
            int64_t length = 0, dest_reg = 0, src_mem_ptr_reg = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &length, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &dest_reg, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &src_mem_ptr_reg, &dummy);
            
            main_inst.data = opcode | ((uint64_t)(length & 0xFF) << 8) | 
                             ((uint64_t)(dest_reg & 0xFF) << 16) | ((uint64_t)(src_mem_ptr_reg & 0xFF) << 24);
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "jmp") == 0) {
            uint8_t opcode = jmp;
            int64_t mode = 0, val_mode = 0, target_val = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &mode, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &val_mode, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &target_val, &dummy);
            
            if (val_mode == 1 && mode == 1) {
                // Relative immediate: adjust target label to be relative offset
                target_val = target_val - current_inst_offset;
            }
            
            main_inst.data = opcode | ((uint64_t)(mode & 0xFF) << 8) | ((uint64_t)(val_mode & 0xFF) << 16);
            if (val_mode == 0) {
                main_inst.data |= ((uint64_t)(target_val & 0xFF) << 24);
            } else {
                main_inst.data |= ((uint64_t)(uint32_t)target_val << 24);
            }
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "jmp_if") == 0) {
            uint8_t opcode = jmp_if;
            int64_t mode = 0, val_mode = 0, flag_reg = 0, target_val = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &mode, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &val_mode, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &flag_reg, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[4].head, &target_val, &dummy);
            
            if (val_mode == 1 && mode == 1) {
                target_val = target_val - current_inst_offset;
            }
            
            main_inst.data = opcode | ((uint64_t)(mode & 0xFF) << 8) | ((uint64_t)(val_mode & 0xFF) << 16) | 
                             ((uint64_t)(flag_reg & 0xFF) << 24);
            if (val_mode == 0) {
                main_inst.data |= ((uint64_t)(target_val & 0xFF) << 32);
            } else {
                main_inst.data |= ((uint64_t)(uint32_t)target_val << 32);
            }
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "call") == 0) {
            uint8_t opcode = call;
            int64_t mode = 0, val_mode = 0, reg_to_rem = 0, target_val = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &mode, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &val_mode, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &reg_to_rem, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[4].head, &target_val, &dummy);
            
            if (val_mode == 0) {
                main_inst.data = opcode | ((uint64_t)(mode & 0xFF) << 8) | ((uint64_t)(val_mode & 0xFF) << 16) | 
                                 ((uint64_t)(reg_to_rem & 0xFF) << 24) | ((uint64_t)(target_val & 0xFF) << 32);
                out->instructions[final_inst_idx++] = main_inst;
            } else {
                if (mode == 1) {
                    target_val = target_val - current_inst_offset;
                }
                main_inst.data = opcode | ((uint64_t)(mode & 0xFF) << 8) | ((uint64_t)(val_mode & 0xFF) << 16) | 
                                 ((uint64_t)(reg_to_rem & 0xFF) << 24);
                out->instructions[final_inst_idx++] = main_inst;
                
                SagittariusInst payload_inst;
                payload_inst.data = (uint64_t)target_val;
                out->instructions[final_inst_idx++] = payload_inst;
            }
            
        } else if (strcmp(name, "ret") == 0) {
            uint8_t opcode = ret;
            int64_t reg = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &reg, &dummy);
            
            main_inst.data = opcode | ((uint64_t)(reg & 0xFF) << 8);
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "cmp") == 0) {
            uint8_t opcode = cmp;
            sagittarius_type type = parse_type(ii->args[1].head);
            
            uint8_t op_val = 0;
            const char* op_name = ii->args[2].head;
            if (strcmp(op_name, "eq") == 0) op_val = 0;
            else if (strcmp(op_name, "ne") == 0) op_val = 1;
            else if (strcmp(op_name, "lt") == 0) op_val = 2;
            else if (strcmp(op_name, "le") == 0) op_val = 3;
            else if (strcmp(op_name, "gt") == 0) op_val = 4;
            else if (strcmp(op_name, "ge") == 0) op_val = 5;
            
            int64_t L = 0, R = 0, Dest = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &L, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[4].head, &R, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[5].head, &Dest, &dummy);
            
            main_inst.data = opcode | ((uint64_t)type << 8) | ((uint64_t)op_val << 16) | 
                             ((uint64_t)(L & 0xFF) << 24) | ((uint64_t)(R & 0xFF) << 32) | 
                             ((uint64_t)(Dest & 0xFF) << 40);
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "mathv") == 0) {
            uint8_t opcode = mathv;
            int64_t w = 0, h = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &w, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &h, &dummy);
            
            uint8_t op_val = 0;
            const char* op_name = ii->args[3].head;
            if (strcmp(op_name, "add") == 0) op_val = 0;
            else if (strcmp(op_name, "sub") == 0) op_val = 1;
            else if (strcmp(op_name, "mul") == 0) op_val = 2;
            else if (strcmp(op_name, "div") == 0) op_val = 3;
            
            int64_t L_ptr = 0, R_ptr = 0, Dest_ptr = 0;
            resolve_token(intermediate, labels, label_count, ii->args[4].head, &L_ptr, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[5].head, &R_ptr, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[6].head, &Dest_ptr, &dummy);
            
            main_inst.data = opcode | ((uint64_t)(w & 0xFF) << 8) | ((uint64_t)(h & 0xFF) << 16) | 
                             ((uint64_t)op_val << 24) | ((uint64_t)(L_ptr & 0xFF) << 32) | 
                             ((uint64_t)(R_ptr & 0xFF) << 40) | ((uint64_t)(Dest_ptr & 0xFF) << 48);
            
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "halt") == 0) {
            main_inst.data = halt;
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "syscall") == 0) {
            uint8_t opcode = syscall;
            int64_t ns_reg = 0, func_reg = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &ns_reg, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &func_reg, &dummy);
            
            main_inst.data = opcode | ((uint64_t)(ns_reg & 0xFF) << 8) | ((uint64_t)(func_reg & 0xFF) << 16);
            out->instructions[final_inst_idx++] = main_inst;
            
        } else if (strcmp(name, "tsyscall") == 0) {
            uint8_t opcode = tsyscall;
            int64_t ns_reg = 0, func_reg = 0, result_reg = 0;
            double dummy;
            resolve_token(intermediate, labels, label_count, ii->args[1].head, &ns_reg, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[2].head, &func_reg, &dummy);
            resolve_token(intermediate, labels, label_count, ii->args[3].head, &result_reg, &dummy);
            
            main_inst.data = opcode | ((uint64_t)(ns_reg & 0xFF) << 8) | ((uint64_t)(func_reg & 0xFF) << 16) | 
                             ((uint64_t)(result_reg & 0xFF) << 24);
            out->instructions[final_inst_idx++] = main_inst;
        }
    }
    
    free(labels);
    return true;
}

bool Sag_WriteProgram(FILE* f, SagittariusProgram* program) {
    if (!f || !program) return false;
    
    if (fwrite(SAGITTARUIS_PROGRAM_HEAD, 1, 8, f) != 8) return false;
    
    uint32_t version = SAGITTARIUS_PROGRAM_FORMAT_VERSION;
    if (fwrite(&version, 4, 1, f) != 1) return false;
    
    if (fwrite(&program->instCount, 8, 1, f) != 1) return false;
    if (fwrite(&program->data_size, 8, 1, f) != 1) return false;
    
    if (program->instCount > 0) {
        if (fwrite(program->instructions, sizeof(SagittariusInst), program->instCount, f) != program->instCount) {
            return false;
        }
    }
    
    if (program->data_size > 0) {
        if (fwrite(program->data, 1, program->data_size, f) != program->data_size) {
            return false;
        }
    }
    
    return true;
}