/*
 * DWN - Desktop Window Manager
 * retoor <retoor@molodetz.nl>
 * OCR text extraction using Tesseract - Production Hardened
 */

#include "ocr.h"
#include "util.h"

#include <stdlib.h>
#include <string.h>
#include <stdatomic.h>
#include <tesseract/capi.h>
#include <leptonica/allheaders.h>

/* Thread state tracking */
typedef enum {
    OCR_THREAD_RUNNING,
    OCR_THREAD_COMPLETED,
    OCR_THREAD_JOINED
} OcrThreadState;

typedef struct OcrRequestInternal {
    OcrRequest base;
    OcrThreadState thread_state;
    atomic_int cancelled;
    struct OcrRequestInternal *next;
} OcrRequestInternal;

static TessBaseAPI *tess_api = NULL;
static bool initialized = false;
static pthread_mutex_t ocr_mutex = PTHREAD_MUTEX_INITIALIZER;
static OcrRequestInternal *ocr_queue = NULL;
static atomic_int ocr_shutting_down = 0;

bool ocr_init(void)
{
    if (initialized) {
        return true;
    }
    
    atomic_store(&ocr_shutting_down, 0);
    
    tess_api = TessBaseAPICreate();
    if (tess_api == NULL) {
        LOG_ERROR("Failed to create Tesseract API");
        return false;
    }
    
    if (TessBaseAPIInit3(tess_api, NULL, "eng") != 0) {
        LOG_WARN("Tesseract initialization failed, OCR unavailable");
        TessBaseAPIDelete(tess_api);
        tess_api = NULL;
        return false;
    }
    
    initialized = true;
    LOG_INFO("OCR module initialized with Tesseract");
    return true;
}

void ocr_cleanup(void)
{
    atomic_store(&ocr_shutting_down, 1);
    
    pthread_mutex_lock(&ocr_mutex);
    
    OcrRequestInternal *req = ocr_queue;
    while (req != NULL) {
        OcrRequestInternal *next = req->next;
        
        /* Signal cancellation */
        atomic_store(&req->cancelled, 1);
        
        /* Remove from queue */
        req->next = NULL;
        
        pthread_mutex_unlock(&ocr_mutex);
        
        /* Join thread if not already joined */
        if (req->thread_state != OCR_THREAD_JOINED) {
            pthread_join(req->base.thread, NULL);
            req->thread_state = OCR_THREAD_JOINED;
        }
        
        /* Free resources */
        if (req->base.png_data != NULL) {
            dwn_free(req->base.png_data);
            req->base.png_data = NULL;
        }
        ocr_result_free(&req->base.result);
        dwn_free(req);
        
        pthread_mutex_lock(&ocr_mutex);
        req = next;
    }
    ocr_queue = NULL;
    
    pthread_mutex_unlock(&ocr_mutex);
    
    if (tess_api != NULL) {
        TessBaseAPIEnd(tess_api);
        TessBaseAPIDelete(tess_api);
        tess_api = NULL;
    }
    initialized = false;
}

bool ocr_is_available(void)
{
    return initialized && tess_api != NULL && !atomic_load(&ocr_shutting_down);
}

const char *ocr_status_string(OcrStatus status)
{
    switch (status) {
        case OCR_OK:
            return "OK";
        case OCR_ERROR_INVALID_ARG:
            return "Invalid argument";
        case OCR_ERROR_INIT:
            return "Initialization error";
        case OCR_ERROR_IMAGE:
            return "Image processing error";
        case OCR_ERROR_NO_MEMORY:
            return "Out of memory";
        case OCR_ERROR_NOT_AVAILABLE:
            return "OCR not available";
        default:
            return "Unknown error";
    }
}

void ocr_result_free(OcrResult *result)
{
    if (result == NULL) {
        return;
    }
    if (result->text != NULL) {
        dwn_free(result->text);
        result->text = NULL;
    }
    result->length = 0;
    result->confidence = 0.0f;
}

OcrStatus ocr_extract_from_png(const unsigned char *png_data, size_t size,
                               const char *language, OcrResult *result)
{
    if (result == NULL) {
        return OCR_ERROR_INVALID_ARG;
    }
    
    memset(result, 0, sizeof(OcrResult));
    
    if (png_data == NULL || size == 0) {
        return OCR_ERROR_INVALID_ARG;
    }
    
    if (!ocr_is_available()) {
        return OCR_ERROR_NOT_AVAILABLE;
    }
    
    const char *lang = (language != NULL && language[0] != '\0') ? language : "eng";
    
    if (TessBaseAPIInit3(tess_api, NULL, lang) != 0) {
        LOG_WARN("Failed to set language '%s', falling back to 'eng'", lang);
        if (TessBaseAPIInit3(tess_api, NULL, "eng") != 0) {
            return OCR_ERROR_INIT;
        }
    }
    
    PIX *image = pixReadMem(png_data, size);
    if (image == NULL) {
        LOG_ERROR("Failed to read PNG image for OCR");
        return OCR_ERROR_IMAGE;
    }
    
    TessBaseAPISetImage2(tess_api, image);
    
    char *text = TessBaseAPIGetUTF8Text(tess_api);
    if (text == NULL) {
        pixDestroy(&image);
        return OCR_ERROR_IMAGE;
    }
    
    int conf = TessBaseAPIMeanTextConf(tess_api);
    float confidence = (float)conf / 100.0f;
    
    size_t text_len = strlen(text);
    result->text = dwn_malloc(text_len + 1);
    if (result->text == NULL) {
        TessDeleteText(text);
        pixDestroy(&image);
        return OCR_ERROR_NO_MEMORY;
    }
    
    memcpy(result->text, text, text_len + 1);
    result->length = text_len;
    result->confidence = confidence;
    
    TessDeleteText(text);
    pixDestroy(&image);
    
    return OCR_OK;
}

static void *ocr_worker(void *arg)
{
    OcrRequestInternal *req = (OcrRequestInternal *)arg;
    
    if (atomic_load(&ocr_shutting_down) || atomic_load(&req->cancelled)) {
        req->base.state = OCR_STATE_ERROR;
        req->base.status = OCR_ERROR_NOT_AVAILABLE;
        return NULL;
    }
    
    /* Create local Tesseract instance for thread safety */
    TessBaseAPI *local_tess = TessBaseAPICreate();
    if (local_tess == NULL) {
        req->base.status = OCR_ERROR_INIT;
        req->base.state = OCR_STATE_ERROR;
        return NULL;
    }
    
    const char *lang = req->base.language[0] ? req->base.language : "eng";
    if (TessBaseAPIInit3(local_tess, NULL, lang) != 0) {
        if (TessBaseAPIInit3(local_tess, NULL, "eng") != 0) {
            TessBaseAPIDelete(local_tess);
            req->base.status = OCR_ERROR_INIT;
            req->base.state = OCR_STATE_ERROR;
            return NULL;
        }
    }
    
    /* Check for cancellation */
    if (atomic_load(&req->cancelled)) {
        TessBaseAPIEnd(local_tess);
        TessBaseAPIDelete(local_tess);
        req->base.status = OCR_ERROR_NOT_AVAILABLE;
        req->base.state = OCR_STATE_ERROR;
        return NULL;
    }
    
    PIX *image = pixReadMem(req->base.png_data, req->base.png_size);
    if (image == NULL) {
        TessBaseAPIEnd(local_tess);
        TessBaseAPIDelete(local_tess);
        req->base.status = OCR_ERROR_IMAGE;
        req->base.state = OCR_STATE_ERROR;
        return NULL;
    }
    
    TessBaseAPISetImage2(local_tess, image);
    
    /* Check for cancellation again before expensive OCR operation */
    if (atomic_load(&req->cancelled)) {
        pixDestroy(&image);
        TessBaseAPIEnd(local_tess);
        TessBaseAPIDelete(local_tess);
        req->base.status = OCR_ERROR_NOT_AVAILABLE;
        req->base.state = OCR_STATE_ERROR;
        return NULL;
    }
    
    char *text = TessBaseAPIGetUTF8Text(local_tess);
    int conf = TessBaseAPIMeanTextConf(local_tess);
    
    if (text != NULL) {
        size_t len = strlen(text);
        req->base.result.text = dwn_malloc(len + 1);
        if (req->base.result.text != NULL) {
            memcpy(req->base.result.text, text, len + 1);
            req->base.result.length = len;
            req->base.result.confidence = (float)conf / 100.0f;
            req->base.status = OCR_OK;
            req->base.state = OCR_STATE_COMPLETED;
        } else {
            req->base.status = OCR_ERROR_NO_MEMORY;
            req->base.state = OCR_STATE_ERROR;
        }
        TessDeleteText(text);
    } else {
        req->base.status = OCR_ERROR_IMAGE;
        req->base.state = OCR_STATE_ERROR;
    }
    
    pixDestroy(&image);
    TessBaseAPIEnd(local_tess);
    TessBaseAPIDelete(local_tess);
    
    return NULL;
}

OcrRequest *ocr_extract_async(const unsigned char *png_data, size_t size,
                              const char *language,
                              void (*callback)(OcrRequest *))
{
    if (!initialized || atomic_load(&ocr_shutting_down)) {
        return NULL;
    }
    
    if (png_data == NULL || size == 0) {
        return NULL;
    }
    
    /* Check for reasonable size to prevent memory exhaustion */
    if (size > 100 * 1024 * 1024) {  /* 100MB limit */
        LOG_ERROR("OCR image too large: %zu bytes", size);
        return NULL;
    }
    
    OcrRequestInternal *req = dwn_calloc(1, sizeof(OcrRequestInternal));
    if (req == NULL) {
        return NULL;
    }
    
    req->base.png_data = dwn_malloc(size);
    if (req->base.png_data == NULL) {
        dwn_free(req);
        return NULL;
    }
    
    memcpy(req->base.png_data, png_data, size);
    req->base.png_size = size;
    
    if (language != NULL && language[0] != '\0') {
        strncpy(req->base.language, language, sizeof(req->base.language) - 1);
        req->base.language[sizeof(req->base.language) - 1] = '\0';
    }
    
    req->base.state = OCR_STATE_PENDING;
    req->base.callback = callback;
    atomic_init(&req->cancelled, 0);
    req->thread_state = OCR_THREAD_RUNNING;
    
    pthread_mutex_lock(&ocr_mutex);
    req->next = ocr_queue;
    ocr_queue = req;
    pthread_mutex_unlock(&ocr_mutex);
    
    if (pthread_create(&req->base.thread, NULL, ocr_worker, req) != 0) {
        pthread_mutex_lock(&ocr_mutex);
        /* Remove from queue */
        OcrRequestInternal **pp = &ocr_queue;
        while (*pp != NULL) {
            if (*pp == req) {
                *pp = req->next;
                break;
            }
            pp = &(*pp)->next;
        }
        pthread_mutex_unlock(&ocr_mutex);
        
        dwn_free(req->base.png_data);
        dwn_free(req);
        return NULL;
    }
    
    LOG_DEBUG("OCR async request started");
    return (OcrRequest *)req;
}

void ocr_process_pending(void)
{
    if (atomic_load(&ocr_shutting_down)) {
        return;
    }
    
    pthread_mutex_lock(&ocr_mutex);
    OcrRequestInternal **pp = &ocr_queue;
    
    while (*pp != NULL) {
        OcrRequestInternal *req = *pp;
        
        if (req->base.state == OCR_STATE_COMPLETED ||
            req->base.state == OCR_STATE_ERROR) {
            
            /* Remove from queue */
            *pp = req->next;
            pthread_mutex_unlock(&ocr_mutex);
            
            /* Join thread if not already joined */
            if (req->thread_state != OCR_THREAD_JOINED) {
                pthread_join(req->base.thread, NULL);
                req->thread_state = OCR_THREAD_JOINED;
            }
            
            /* Invoke callback */
            if (req->base.callback != NULL) {
                req->base.callback((OcrRequest *)req);
            }
            
            pthread_mutex_lock(&ocr_mutex);
        } else {
            pp = &(*pp)->next;
        }
    }
    
    pthread_mutex_unlock(&ocr_mutex);
}

void ocr_cancel(OcrRequest *base_req)
{
    if (base_req == NULL) {
        return;
    }
    
    OcrRequestInternal *req = (OcrRequestInternal *)base_req;
    
    /* Signal cancellation */
    atomic_store(&req->cancelled, 1);
    
    /* Remove from queue */
    pthread_mutex_lock(&ocr_mutex);
    OcrRequestInternal **pp = &ocr_queue;
    while (*pp != NULL) {
        if (*pp == req) {
            *pp = req->next;
            break;
        }
        pp = &(*pp)->next;
    }
    pthread_mutex_unlock(&ocr_mutex);
    
    /* Wait for thread to complete (it will check cancelled flag) */
    if (req->thread_state != OCR_THREAD_JOINED) {
        pthread_join(req->base.thread, NULL);
        req->thread_state = OCR_THREAD_JOINED;
    }
    
    /* Free resources */
    if (req->base.png_data != NULL) {
        dwn_free(req->base.png_data);
        req->base.png_data = NULL;
    }
    ocr_result_free(&req->base.result);
    dwn_free(req);
}
