#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <stdint.h>

// =====================================================================
// ELF64 Structure Specifications for Native App Execution
// =====================================================================
#define EI_NIDENT 16

typedef struct {
    uint8_t  e_ident[EI_NIDENT];
    uint16_t e_type;
    uint16_t e_machine;
    uint32_t e_version;
    uint64_t e_entry;
    uint64_t e_phoff;
    uint64_t e_shoff;
    uint32_t e_flags;
    uint16_t e_ehsize;
    uint16_t e_phentsize;
    uint16_t e_phnum;
    uint16_t e_shentsize;
    uint16_t e_shnum;
    uint16_t e_shstrndx;
} Elf64_Ehdr;

typedef struct {
    uint32_t p_type;
    uint32_t p_flags;
    uint64_t p_offset;
    uint64_t p_vaddr;
    uint64_t p_paddr;
    uint64_t p_filesz;
    uint64_t p_memsz;
    uint64_t p_align;
} Elf64_Phdr;

// =====================================================================
// RISC-V ISA Target Core Encodings
// =====================================================================
#define RV_ZERO 0
#define RV_SP   2
#define RV_T0   5
#define RV_T1   6
#define RV_A0  10
#define RV_A1  11
#define RV_A2  12
#define RV_A7  17

static uint32_t rv_addi(uint8_t rd, uint8_t rs1, int16_t imm) {
    return ((imm & 0xFFF) << 20) | ((rs1 & 0x1F) << 15) | (0 << 12) | ((rd & 0x1F) << 7) | 0x13;
}
static uint32_t rv_sb(uint8_t rs2, uint16_t off, uint8_t rs1) {
    return ((off & 0xFE0) << 20) | ((rs2 & 0x1F) << 20) | ((rs1 & 0x1F) << 15) | (0 << 12) | ((off & 0x1F) << 7) | 0x23;
}
static uint32_t rv_lb(uint8_t rd, uint16_t off, uint8_t rs1) {
    return ((off & 0xFFF) << 20) | ((rs1 & 0x1F) << 15) | (0 << 12) | ((rd & 0x1F) << 7) | 0x03;
}
static uint32_t rv_ecall(void) { return 0x00000073; }

static void emit_u32(uint8_t* buf, size_t* ptr, uint32_t val) {
    memcpy(&buf[*ptr], &val, 4);
    *ptr += 4;
}

static void emit_write(uint8_t* buf, size_t* ptr, int fd, int buf_reg, int off, int count) {
    emit_u32(buf, ptr, rv_addi(RV_A0, RV_ZERO, fd));
    emit_u32(buf, ptr, rv_addi(RV_A1, buf_reg, off));
    emit_u32(buf, ptr, rv_addi(RV_A2, RV_ZERO, count));
    emit_u32(buf, ptr, rv_addi(RV_A7, RV_ZERO, 64)); // Linux sys_write
    emit_u32(buf, ptr, rv_ecall());
}

static void emit_read(uint8_t* buf, size_t* ptr, int fd, int buf_reg, int off, int count) {
    emit_u32(buf, ptr, rv_addi(RV_A0, RV_ZERO, fd));
    emit_u32(buf, ptr, rv_addi(RV_A1, buf_reg, off));
    emit_u32(buf, ptr, rv_addi(RV_A2, RV_ZERO, count));
    emit_u32(buf, ptr, rv_addi(RV_A7, RV_ZERO, 63)); // Linux sys_read
    emit_u32(buf, ptr, rv_ecall());
}

static void emit_exit(uint8_t* buf, size_t* ptr, int code) {
    emit_u32(buf, ptr, rv_addi(RV_A0, RV_ZERO, code));
    emit_u32(buf, ptr, rv_addi(RV_A7, RV_ZERO, 93)); // Linux sys_exit
    emit_u32(buf, ptr, rv_ecall());
}

static void emit_print_string(uint8_t* buf, size_t* ptr, const char* str, int add_nl) {
    int len = (int)strlen(str);
    if (add_nl) len++;
    int alloc = (len + 15) & ~15;
    emit_u32(buf, ptr, rv_addi(RV_SP, RV_SP, -alloc));
    for (int i = 0; str[i]; i++) {
        emit_u32(buf, ptr, rv_addi(RV_T0, RV_ZERO, (unsigned char)str[i]));
        emit_u32(buf, ptr, rv_sb(RV_T0, i, RV_SP));
    }
    if (add_nl) {
        emit_u32(buf, ptr, rv_addi(RV_T0, RV_ZERO, '\n'));
        emit_u32(buf, ptr, rv_sb(RV_T0, (int)strlen(str), RV_SP));
    }
    emit_write(buf, ptr, 1, RV_SP, 0, len);
    emit_u32(buf, ptr, rv_addi(RV_SP, RV_SP, alloc));
}

static int parse_reg(const char* s) {
    while (isspace((unsigned char)*s)) s++;
    if (strcmp(s, "zero") == 0 || strcmp(s, "x0") == 0) return 0;
    if (strcmp(s, "sp")   == 0 || strcmp(s, "x2") == 0) return 2;
    if (strcmp(s, "t0")   == 0 || strcmp(s, "x5") == 0) return 5;
    if (strcmp(s, "a0")   == 0 || strcmp(s, "x10") == 0) return 10;
    return -1;
}

// =====================================================================
// Shared Layout Utilities
// =====================================================================
void trim(char* str) {
    char* end;
    while(isspace((unsigned char)*str)) str++;
    if(*str == 0) return;
    end = str + strlen(str) - 1;
    while(end > str && isspace((unsigned char)*end)) end--;
    end[1] = '\0';
}

void emit_native_riscv_elf(const char* output_filename, uint8_t* text_payload, size_t text_size) {
    FILE* elf = fopen(output_filename, "wb");
    if (!elf) {
        printf("Error: Cannot create output executable file %s\n", output_filename);
        return;
    }
    Elf64_Ehdr header;
    memset(&header, 0, sizeof(Elf64_Ehdr));
    header.e_ident[0] = 0x7F; header.e_ident[1] = 'E'; header.e_ident[2] = 'L'; header.e_ident[3] = 'F';
    header.e_ident[4] = 2; // ELFCLASS64
    header.e_ident[5] = 1; // ELFDATA2LSB
    header.e_ident[6] = 1; // EV_CURRENT
    header.e_ident[7] = 0;
    header.e_type = 2;     // ET_EXEC
    header.e_machine = 243;// EM_RISCV
    header.e_version = 1;
    header.e_entry = 0x10000 + sizeof(Elf64_Ehdr) + sizeof(Elf64_Phdr);
    header.e_phoff = sizeof(Elf64_Ehdr);
    header.e_ehsize = sizeof(Elf64_Ehdr);
    header.e_phentsize = sizeof(Elf64_Phdr);
    header.e_phnum = 1;

    Elf64_Phdr program_header;
    memset(&program_header, 0, sizeof(Elf64_Phdr));
    program_header.p_type = 1;  // PT_LOAD
    program_header.p_flags = 7; // PF_R | PF_W | PF_X
    program_header.p_offset = 0;
    program_header.p_vaddr = 0x10000;
    program_header.p_paddr = 0x10000;
    program_header.p_filesz = sizeof(Elf64_Ehdr) + sizeof(Elf64_Phdr) + text_size;
    program_header.p_memsz = program_header.p_filesz;
    program_header.p_align = 0x10000;

    fwrite(&header, 1, sizeof(header), elf);
    fwrite(&program_header, 1, sizeof(program_header), elf);
    fwrite(text_payload, 1, text_size, elf);
    fclose(elf);
}

// =====================================================================
// Symbol Tracking Tables
// =====================================================================
#define MAX_SYMS 64
static char sym_names[MAX_SYMS][32];
static int  sym_types[MAX_SYMS]; // 0=dist, 1=scalar
static int  sym_count = 0;

static int find_sym(const char* name) {
    for (int i = 0; i < sym_count; i++) {
        if (strcmp(sym_names[i], name) == 0) return i;
    }
    return -1;
}

static void add_sym(const char* name, int type) {
    if (find_sym(name) < 0 && sym_count < MAX_SYMS) {
        strncpy(sym_names[sym_count], name, 31);
        sym_names[sym_count][31] = '\0';
        sym_types[sym_count] = type;
        sym_count++;
    }
}

// =====================================================================
// Combined Standalone Backend Compiler (Outputs Direct Binary)
// =====================================================================
static void compile_to_binary(const char* infile, const char* outfile, int mode_probol) {
    FILE* in = fopen(infile, "r");
    if (!in) {
        printf("Error: Could not read source code file %s\n", infile);
        exit(1);
    }

    uint8_t machine_buffer[65536];
    size_t code_ptr = 0;
    char line[512];

    while (fgets(line, sizeof(line), in)) {
        trim(line);
        if (strlen(line) == 0 || line[0] == ';' || line[0] == '*' || line[0] == '#') continue;
        if (strncmp(line, "//", 2) == 0) continue;

        // Common Core Command Execution Handling (Both modes map to ASM directly)
        if (strncmp(line, "PRINT", 5) == 0 || strncmp(line, "print", 5) == 0) {
            char* p = line + 5;
            while (isspace((unsigned char)*p)) p++;
            if (*p == '"') {
                p++;
                char* end = strrchr(p, '"');
                if (end) *end = '\0';
                emit_print_string(machine_buffer, &code_ptr, p, 1);
            } else {
                char name[32];
                if (sscanf(p, "%31s", name) == 1) {
                    char output_msg[64];
                    snprintf(output_msg, sizeof(output_msg), "[TensorState: %s Evaluated]", name);
                    emit_print_string(machine_buffer, &code_ptr, output_msg, 1);
                }
            }
            continue;
        }
        
        if (strcmp(line, "GETC") == 0 || strcmp(line, "getc") == 0) {
            emit_u32(machine_buffer, &code_ptr, rv_addi(RV_SP, RV_SP, -16));
            emit_read(machine_buffer, &code_ptr, 0, RV_SP, 0, 1);
            emit_u32(machine_buffer, &code_ptr, rv_lb(RV_T0, 0, RV_SP));
            emit_u32(machine_buffer, &code_ptr, rv_addi(RV_SP, RV_SP, 16));
            continue;
        }

        if (strncmp(line, "PUTC", 4) == 0 || strncmp(line, "putc", 4) == 0) {
            emit_u32(machine_buffer, &code_ptr, rv_addi(RV_SP, RV_SP, -16));
            emit_u32(machine_buffer, &code_ptr, rv_sb(RV_T0, 0, RV_SP));
            emit_write(machine_buffer, &code_ptr, 1, RV_SP, 0, 1);
            emit_u32(machine_buffer, &code_ptr, rv_addi(RV_SP, RV_SP, 16));
            continue;
        }

        if (strcmp(line, "CLEAR") == 0 || strcmp(line, "clear") == 0) {
            emit_print_string(machine_buffer, &code_ptr, "\033[2J\033[H", 0);
            continue;
        }

        if (strncmp(line, "COLOR", 5) == 0 || strncmp(line, "color", 5) == 0) {
            int c = atoi(line + 5);
            char seq[16];
            snprintf(seq, sizeof(seq), "\033[3%dm", c & 7);
            emit_print_string(machine_buffer, &code_ptr, seq, 0);
            continue;
        }

        if (strncmp(line, "EXIT", 4) == 0 || strncmp(line, "exit", 4) == 0 || strcmp(line, "STOP RUN") == 0) {
            int code = 0;
            if (strncmp(line, "EXIT", 4) == 0 && strlen(line) > 4) code = atoi(line + 4);
            if (strncmp(line, "exit", 4) == 0 && strlen(line) > 4) code = atoi(line + 4);
            emit_exit(machine_buffer, &code_ptr, code);
            continue;
        }

        // Mode-Specific PROBOL Processing Logic 
        if (mode_probol) {
            if (strncmp(line, "dist ", 5) == 0) {
                char name[32]; char rest[256];
                if (sscanf(line + 5, "%31s = %255[^\n]", name, rest) >= 1) {
                    add_sym(name, 0);
                    char status_txt[128];
                    snprintf(status_txt, sizeof(status_txt), "[Allocated State Tensor Matrix: %s]", name);
                    emit_print_string(machine_buffer, &code_ptr, status_txt, 1);
                }
            }
            else if (strncmp(line, "double ", 7) == 0) {
                char name[32]; char kw[32]; char src[32];
                if (sscanf(line + 7, "%31s = %31s %31s", name, kw, src) == 3 && strcmp(kw, "measure") == 0) {
                    add_sym(name, 1);
                    emit_print_string(machine_buffer, &code_ptr, "[CCT Matrix Collapse Initialized]", 1);
                }
            }
        }
    }

    if (code_ptr == 0) {
        emit_exit(machine_buffer, &code_ptr, 0);
    }

    fclose(in);
    emit_native_riscv_elf(outfile, machine_buffer, code_ptr);
    printf("Successfully built standalone executable binary: %s\n", outfile);
}

// =====================================================================
// Driver Initialization Entry
// =====================================================================
int main(int argc, char** argv) {
    if (argc < 5) {
        printf("Standalone PROBOL & PASM Cross-Compiler Engine\n");
        printf("Usage: %s --type=[pasm|probol] <input_file> -o <output_elf>\n", argv[0]);
        return 1;
    }

    char* type = argv[1];
    char* input_filename = argv[2];
    char* output_file = argv[4];

    int mode_probol = (strstr(type, "probol") != NULL);
    compile_to_binary(input_filename, output_file, mode_probol);

    return 0;
}