/* oneocr_loader.c -- Minimal OneOCR DLL loader for Wine
 * Compile: x86_64-w64-mingw32-gcc -O2 -o oneocr_loader.exe oneocr_loader.c
 * Usage:   wine oneocr_loader.exe <dll_dir> <image_bmp> <model_key>
 * Output:  JSON to stdout
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <windows.h>

/* DLL function types */
typedef long long (*fn_CreateOcrInitOptions)(long long*);
typedef long long (*fn_OcrInitOptionsSetUseModelDelayLoad)(long long, char);
typedef long long (*fn_CreateOcrPipeline)(const char*, const char*, long long, long long*);
typedef long long (*fn_CreateOcrProcessOptions)(long long*);
typedef long long (*fn_OcrProcessOptionsSetMaxRecognitionLineCount)(long long, long long);
typedef long long (*fn_RunOcrPipeline)(long long, void*, long long, long long*);
typedef long long (*fn_GetImageAngle)(long long, float*);
typedef long long (*fn_GetOcrLineCount)(long long, long long*);
typedef long long (*fn_GetOcrLine)(long long, long long, long long*);
typedef long long (*fn_GetOcrLineContent)(long long, const char**);
typedef long long (*fn_GetOcrLineBoundingBox)(long long, void**);
typedef long long (*fn_GetOcrLineWordCount)(long long, long long*);
typedef long long (*fn_GetOcrWord)(long long, long long, long long*);
typedef long long (*fn_GetOcrWordContent)(long long, const char**);
typedef long long (*fn_GetOcrWordBoundingBox)(long long, void**);
typedef long long (*fn_GetOcrWordConfidence)(long long, float*);
typedef void (*fn_ReleaseOcrResult)(long long);
typedef void (*fn_ReleaseOcrInitOptions)(long long);
typedef void (*fn_ReleaseOcrPipeline)(long long);
typedef void (*fn_ReleaseOcrProcessOptions)(long long);

#pragma pack(push, 1)
typedef struct {
    int type;       /* 3 = BGRA 4-channel (matches engine.py) */
    int width;
    int height;
    int reserved;
    long long step;
    unsigned char *data;
} ImageStruct;

typedef struct {
    float x1, y1, x2, y2, x3, y3, x4, y4;
} BBox;
#pragma pack(pop)

/* Simple BMP loader (32-bit BGRA) */
static unsigned char* load_bmp_bgra(const char* path, int* w, int* h) {
    FILE* f = fopen(path, "rb");
    if (!f) return NULL;
    
    unsigned char header[54];
    fread(header, 1, 54, f);
    
    *w = *(int*)(header + 18);
    *h = *(int*)(header + 22);
    int bpp = *(short*)(header + 28);
    int offset = *(int*)(header + 10);
    int abs_h = *h < 0 ? -*h : *h;
    
    fseek(f, offset, SEEK_SET);
    
    /* Allocate BGRA buffer */
    unsigned char* bgra = (unsigned char*)malloc((*w) * abs_h * 4);
    
    if (bpp == 24) {
        int row_size = ((*w * 3 + 3) & ~3);
        unsigned char* row = (unsigned char*)malloc(row_size);
        for (int y = 0; y < abs_h; y++) {
            int dest_y = (*h > 0) ? (abs_h - 1 - y) : y;
            fread(row, 1, row_size, f);
            for (int x = 0; x < *w; x++) {
                bgra[(dest_y * *w + x) * 4 + 0] = row[x * 3 + 0]; /* B */
                bgra[(dest_y * *w + x) * 4 + 1] = row[x * 3 + 1]; /* G */
                bgra[(dest_y * *w + x) * 4 + 2] = row[x * 3 + 2]; /* R */
                bgra[(dest_y * *w + x) * 4 + 3] = 255;              /* A */
            }
        }
        free(row);
    } else if (bpp == 32) {
        for (int y = 0; y < abs_h; y++) {
            int dest_y = (*h > 0) ? (abs_h - 1 - y) : y;
            fread(bgra + dest_y * *w * 4, 1, *w * 4, f);
        }
    }
    
    *h = abs_h;
    fclose(f);
    return bgra;
}

/* Escape JSON string */
static void json_escape(const char* s, char* out, int max) {
    int j = 0;
    out[j++] = '"';
    for (int i = 0; s[i] && j < max - 3; i++) {
        if (s[i] == '"') { out[j++] = '\\'; out[j++] = '"'; }
        else if (s[i] == '\\') { out[j++] = '\\'; out[j++] = '\\'; }
        else if (s[i] == '\n') { out[j++] = '\\'; out[j++] = 'n'; }
        else if (s[i] == '\r') { out[j++] = '\\'; out[j++] = 'r'; }
        else if (s[i] == '\t') { out[j++] = '\\'; out[j++] = 't'; }
        else out[j++] = s[i];
    }
    out[j++] = '"';
    out[j] = 0;
}

int main(int argc, char** argv) {
    if (argc < 4) {
        fprintf(stderr, "Usage: %s <dll_dir> <image.bmp> <model_key_hex>\n", argv[0]);
        return 1;
    }
    
    const char* dll_dir = argv[1];
    const char* img_path = argv[2];
    const char* key_hex = argv[3];
    
    /* Set DLL search path */
    SetDllDirectoryA(dll_dir);
    char old_path[32768];
    GetEnvironmentVariableA("PATH", old_path, sizeof(old_path));
    char new_path[32768];
    snprintf(new_path, sizeof(new_path), "%s;%s", dll_dir, old_path);
    SetEnvironmentVariableA("PATH", new_path);
    
    /* Load DLL */
    char dll_path[MAX_PATH];
    snprintf(dll_path, sizeof(dll_path), "%s\\oneocr.dll", dll_dir);
    
    HMODULE hmod = LoadLibraryA(dll_path);
    if (!hmod) {
        fprintf(stderr, "{\"error\": \"LoadLibrary failed: %lu\"}\n", GetLastError());
        return 1;
    }
    
    /* Get function pointers */
    #define GETFN(name) fn_##name p##name = (fn_##name)GetProcAddress(hmod, #name); \
        if (!p##name) { fprintf(stderr, "{\"error\": \"GetProcAddress(%s) failed\"}\n", #name); return 1; }
    
    GETFN(CreateOcrInitOptions)
    GETFN(OcrInitOptionsSetUseModelDelayLoad) 
    GETFN(CreateOcrPipeline)
    GETFN(CreateOcrProcessOptions)
    GETFN(OcrProcessOptionsSetMaxRecognitionLineCount)
    GETFN(RunOcrPipeline)
    GETFN(GetImageAngle)
    GETFN(GetOcrLineCount)
    GETFN(GetOcrLine)
    GETFN(GetOcrLineContent)
    GETFN(GetOcrLineBoundingBox)
    GETFN(GetOcrLineWordCount)
    GETFN(GetOcrWord)
    GETFN(GetOcrWordContent)
    GETFN(GetOcrWordBoundingBox)
    GETFN(GetOcrWordConfidence)
    GETFN(ReleaseOcrResult)
    GETFN(ReleaseOcrInitOptions)
    GETFN(ReleaseOcrPipeline)
    GETFN(ReleaseOcrProcessOptions)
    
    /* Model path and key */
    char model_path[MAX_PATH];
    snprintf(model_path, sizeof(model_path), "%s\\oneocr.onemodel", dll_dir);
    
    /* Decode hex key */
    int key_len = strlen(key_hex) / 2;
    char key[64];
    for (int i = 0; i < key_len && i < 63; i++) {
        sscanf(key_hex + i*2, "%2hhx", &key[i]);
    }
    key[key_len] = 0;
    
    /* Initialize pipeline */
    long long init_opts = 0;
    pCreateOcrInitOptions(&init_opts);
    
    long long pipeline = 0;
    long long res = pCreateOcrPipeline(model_path, key, init_opts, &pipeline);
    if (res != 0) {
        fprintf(stderr, "{\"error\": \"CreateOcrPipeline failed: %lld\"}\n", res);
        return 1;
    }
    
    long long proc_opts = 0;
    pCreateOcrProcessOptions(&proc_opts);
    pOcrProcessOptionsSetMaxRecognitionLineCount(proc_opts, 200);
    
    /* Load image */
    int w = 0, h = 0;
    unsigned char* data = load_bmp_bgra(img_path, &w, &h);
    if (!data) {
        fprintf(stderr, "{\"error\": \"Failed to load image\"}\n");
        return 1;
    }
    
    ImageStruct img = {3, w, h, 0, (long long)(w * 4), data};
    
    /* Run OCR */
    long long result = 0;
    res = pRunOcrPipeline(pipeline, &img, proc_opts, &result);
    if (res != 0) {
        fprintf(stderr, "{\"error\": \"RunOcrPipeline failed: %lld\"}\n", res);
        return 1;
    }
    
    /* Extract results */
    float angle = 0;
    pGetImageAngle(result, &angle);
    
    long long line_count = 0;
    pGetOcrLineCount(result, &line_count);
    
    /* Output JSON */
    char buf[65536];
    int pos = 0;
    pos += snprintf(buf + pos, sizeof(buf) - pos, 
        "{\"text_angle\": %.4f, \"lines\": [", angle);
    
    for (long long i = 0; i < line_count; i++) {
        long long line = 0;
        pGetOcrLine(result, i, &line);
        
        const char* line_text = NULL;
        pGetOcrLineContent(line, &line_text);
        
        BBox* line_bbox = NULL;
        pGetOcrLineBoundingBox(line, (void**)&line_bbox);
        
        long long word_count = 0;
        pGetOcrLineWordCount(line, &word_count);
        
        if (i > 0) pos += snprintf(buf + pos, sizeof(buf) - pos, ", ");
        
        char esc_line[4096];
        json_escape(line_text ? line_text : "", esc_line, sizeof(esc_line));
        
        pos += snprintf(buf + pos, sizeof(buf) - pos,
            "{\"text\": %s, \"bbox\": [%.1f,%.1f,%.1f,%.1f,%.1f,%.1f,%.1f,%.1f], \"words\": [",
            esc_line,
            line_bbox ? line_bbox->x1 : 0, line_bbox ? line_bbox->y1 : 0,
            line_bbox ? line_bbox->x2 : 0, line_bbox ? line_bbox->y2 : 0,
            line_bbox ? line_bbox->x3 : 0, line_bbox ? line_bbox->y3 : 0,
            line_bbox ? line_bbox->x4 : 0, line_bbox ? line_bbox->y4 : 0);
        
        for (long long j = 0; j < word_count; j++) {
            long long word = 0;
            pGetOcrWord(line, j, &word);
            
            const char* word_text = NULL;
            pGetOcrWordContent(word, &word_text);
            
            BBox* word_bbox = NULL;
            pGetOcrWordBoundingBox(word, (void**)&word_bbox);
            
            float word_conf = 0;
            pGetOcrWordConfidence(word, &word_conf);
            
            if (j > 0) pos += snprintf(buf + pos, sizeof(buf) - pos, ", ");
            
            char esc_word[2048];
            json_escape(word_text ? word_text : "", esc_word, sizeof(esc_word));
            
            pos += snprintf(buf + pos, sizeof(buf) - pos,
                "{\"text\": %s, \"confidence\": %.4f, \"bbox\": [%.1f,%.1f,%.1f,%.1f,%.1f,%.1f,%.1f,%.1f]}",
                esc_word, word_conf,
                word_bbox ? word_bbox->x1 : 0, word_bbox ? word_bbox->y1 : 0,
                word_bbox ? word_bbox->x2 : 0, word_bbox ? word_bbox->y2 : 0,
                word_bbox ? word_bbox->x3 : 0, word_bbox ? word_bbox->y3 : 0,
                word_bbox ? word_bbox->x4 : 0, word_bbox ? word_bbox->y4 : 0);
        }
        
        pos += snprintf(buf + pos, sizeof(buf) - pos, "]}");
    }
    
    pos += snprintf(buf + pos, sizeof(buf) - pos, "]}");
    
    /* Write JSON to stdout */
    printf("%s\n", buf);
    fflush(stdout);
    
    /* Cleanup */
    pReleaseOcrResult(result);
    free(data);
    pReleaseOcrProcessOptions(proc_opts);
    pReleaseOcrPipeline(pipeline);
    pReleaseOcrInitOptions(init_opts);
    FreeLibrary(hmod);
    
    return 0;
}