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

int main(int ac, char** av) {
    printf("Copyright (C) 2026 Creeper Lv.\n");
    if (ac < 2) {
        printf("Usage: SagittariusAssembler <input_file_0> <input_file_1> ... [-o <output_file>]\n");
        return 1;
    }
    
    char* output_file = NULL;
    char** input_files = malloc(ac * sizeof(char*));
    int input_count = 0;
    
    for (int i = 1; i < ac; ++i) {
        if (strcmp(av[i], "-o") == 0) {
            if (i + 1 < ac) {
                output_file = av[i+1];
                i++;
            } else {
                printf("Error: Missing output file after -o\n");
                free(input_files);
                return 1;
            }
        } else {
            input_files[input_count++] = av[i];
        }
    }
    
    if (input_count == 0) {
        printf("Error: No input files specified\n");
        free(input_files);
        return 1;
    }
    
    char* dynamic_out = NULL;
    if (!output_file) {
        size_t len = strlen(input_files[0]);
        dynamic_out = malloc(len + 5);
        sprintf(dynamic_out, "%s.out", input_files[0]);
        output_file = dynamic_out;
    }
    
    printf("Assembling %d file(s) into '%s'...\n", input_count, output_file);
    
    Sag_IntermediateProgram combined;
    memset(&combined, 0, sizeof(combined));
    
    for (int i = 0; i < input_count; ++i) {
        FILE* f = fopen(input_files[i], "r");
        if (!f) {
            printf("Error: Cannot open input file '%s'\n", input_files[i]);
            free(input_files);
            if (dynamic_out) free(dynamic_out);
            return 1;
        }
        
        Sag_IntermediateProgram ip;
        memset(&ip, 0, sizeof(ip));
        if (!Sag_Scan(f, &ip)) {
            printf("Error scanning '%s'\n", input_files[i]);
            fclose(f);
            free(input_files);
            if (dynamic_out) free(dynamic_out);
            return 1;
        }
        fclose(f);
        
        if (i == 0) {
            combined = ip;
        } else {
            Sag_IntermediateProgram temp;
            memset(&temp, 0, sizeof(temp));
            if (!Sag_Combine(&combined, &ip, &temp)) {
                printf("Error combining intermediate programs\n");
                free(input_files);
                if (dynamic_out) free(dynamic_out);
                return 1;
            }
            free(combined.insts);
            free(combined.data);
            free(combined.Consts);
            free(ip.insts);
            free(ip.data);
            free(ip.Consts);
            combined = temp;
        }
    }
    
    // Finalize
    SagittariusProgram prog;
    memset(&prog, 0, sizeof(prog));
    if (!Sag_Finalize(&combined, &prog)) {
        printf("Error: Finalization failed!\n");
        free(combined.insts);
        free(combined.data);
        free(combined.Consts);
        free(input_files);
        if (dynamic_out) free(dynamic_out);
        return 1;
    }
    
    // Write output
    FILE* out_f = fopen(output_file, "wb");
    if (!out_f) {
        printf("Error: Cannot open output file '%s' for writing\n", output_file);
        free(combined.insts);
        free(combined.data);
        free(combined.Consts);
        free(prog.instructions);
        if (prog.data) free(prog.data);
        free(input_files);
        if (dynamic_out) free(dynamic_out);
        return 1;
    }
    
    if (!Sag_WriteProgram(out_f, &prog)) {
        printf("Error writing to output file '%s'\n", output_file);
        fclose(out_f);
        free(combined.insts);
        free(combined.data);
        free(combined.Consts);
        free(prog.instructions);
        if (prog.data) free(prog.data);
        free(input_files);
        if (dynamic_out) free(dynamic_out);
        return 1;
    }
    fclose(out_f);
    
    printf("Successfully assembled! Instruction count: %llu, Data size: %llu bytes\n", 
           (unsigned long long)prog.instCount, (unsigned long long)prog.data_size);
           
    // Cleanup
    free(combined.insts);
    free(combined.data);
    free(combined.Consts);
    free(prog.instructions);
    if (prog.data) free(prog.data);
    free(input_files);
    if (dynamic_out) free(dynamic_out);
    
    return 0;
}