#include "pdf_bridge.h"

#include <mupdf/fitz.h>

#include <math.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#define PARDES_PDF_INITIAL_RESULT_QUADS 64

struct pardes_pdf_document {
    fz_context *ctx;
    fz_document *doc;
    int page_count;

    /* One UI thread owns a document. Page switches invalidate this pair. */
    int cached_page_number;
    fz_page *cached_page;
    fz_stext_page *cached_text;
    fz_rect cached_bounds;
};

static void
pardes_pdf_drop_cached_page(pardes_pdf_document *document)
{
    fz_drop_stext_page(document->ctx, document->cached_text);
    fz_drop_page(document->ctx, document->cached_page);
    document->cached_text = NULL;
    document->cached_page = NULL;
    document->cached_page_number = -1;
}

static int
pardes_pdf_valid_bounds(fz_rect bounds)
{
    return isfinite(bounds.x0) && isfinite(bounds.y0) &&
        isfinite(bounds.x1) && isfinite(bounds.y1) &&
        bounds.x1 > bounds.x0 && bounds.y1 > bounds.y0;
}

/*
 * Load into temporaries first, then replace the cache in one non-throwing
 * commit. Thus a failed page/text load leaves the previous cached page valid.
 */
static void
pardes_pdf_cache_page(
    pardes_pdf_document *document,
    int page_number,
    int require_text)
{
    fz_context *ctx = document->ctx;
    fz_page *page = NULL;
    fz_stext_page *text = NULL;
    fz_rect bounds;
    fz_stext_options options = {0};
    options.flags = FZ_STEXT_DEHYPHENATE;

    if (document->cached_page_number == page_number) {
        if (require_text && document->cached_text == NULL)
            document->cached_text = fz_new_stext_page_from_page(
                ctx, document->cached_page, &options);
        return;
    }

    fz_var(page);
    fz_var(text);
    fz_try(ctx)
    {
        page = fz_load_page(ctx, document->doc, page_number);
        bounds = fz_bound_page(ctx, page);
        if (!pardes_pdf_valid_bounds(bounds))
            fz_throw(ctx, FZ_ERROR_FORMAT, "PDF page has empty bounds");
        if (require_text)
            text = fz_new_stext_page_from_page(ctx, page, &options);
    }
    fz_catch(ctx)
    {
        fz_drop_stext_page(ctx, text);
        fz_drop_page(ctx, page);
        fz_rethrow(ctx);
    }

    pardes_pdf_drop_cached_page(document);
    document->cached_page_number = page_number;
    document->cached_page = page;
    document->cached_text = text;
    document->cached_bounds = bounds;
}

static float
pardes_pdf_clamp_unit(float value)
{
    if (value < 0.0f)
        return 0.0f;
    if (value > 1.0f)
        return 1.0f;
    return value;
}

static int
pardes_pdf_valid_point(pardes_pdf_point point)
{
    return isfinite(point.x) && isfinite(point.y);
}

static fz_point
pardes_pdf_point_to_page(pardes_pdf_document *document, pardes_pdf_point point)
{
    fz_rect bounds = document->cached_bounds;
    float x = pardes_pdf_clamp_unit(point.x);
    float y = pardes_pdf_clamp_unit(point.y);
    return fz_make_point(
        bounds.x0 + x * (bounds.x1 - bounds.x0),
        bounds.y0 + y * (bounds.y1 - bounds.y0));
}

static pardes_pdf_point
pardes_pdf_point_from_page(pardes_pdf_document *document, fz_point point)
{
    fz_rect bounds = document->cached_bounds;
    pardes_pdf_point result;
    result.x = pardes_pdf_clamp_unit(
        (point.x - bounds.x0) / (bounds.x1 - bounds.x0));
    result.y = pardes_pdf_clamp_unit(
        (point.y - bounds.y0) / (bounds.y1 - bounds.y0));
    return result;
}

static pardes_pdf_quad
pardes_pdf_quad_from_page(pardes_pdf_document *document, fz_quad quad)
{
    pardes_pdf_quad result;
    result.ul = pardes_pdf_point_from_page(document, quad.ul);
    result.ur = pardes_pdf_point_from_page(document, quad.ur);
    result.ll = pardes_pdf_point_from_page(document, quad.ll);
    result.lr = pardes_pdf_point_from_page(document, quad.lr);
    return result;
}

static int
pardes_pdf_valid_quad(pardes_pdf_quad quad)
{
    return pardes_pdf_valid_point(quad.ul) &&
        pardes_pdf_valid_point(quad.ur) &&
        pardes_pdf_valid_point(quad.ll) &&
        pardes_pdf_valid_point(quad.lr);
}

static fz_quad
pardes_pdf_quad_to_page(pardes_pdf_document *document, pardes_pdf_quad quad)
{
    fz_quad result;
    result.ul = pardes_pdf_point_to_page(document, quad.ul);
    result.ur = pardes_pdf_point_to_page(document, quad.ur);
    result.ll = pardes_pdf_point_to_page(document, quad.ll);
    result.lr = pardes_pdf_point_to_page(document, quad.lr);
    return result;
}

static size_t
pardes_pdf_grown_capacity(size_t current, size_t needed, size_t maximum)
{
    size_t capacity = current == 0 ? PARDES_PDF_INITIAL_RESULT_QUADS : current;
    while (capacity < needed) {
        if (capacity > maximum / 2) {
            capacity = maximum;
            break;
        }
        capacity *= 2;
    }
    return capacity;
}

pardes_pdf_document *
pardes_pdf_open(const char *path, int *page_count)
{
    pardes_pdf_document *state;
    fz_context *ctx;
    fz_document *doc = NULL;
    int pages = 0;

    if (path == NULL || page_count == NULL)
        return NULL;

    state = (pardes_pdf_document *)calloc(1, sizeof(*state));
    if (state == NULL)
        return NULL;
    state->cached_page_number = -1;

    ctx = fz_new_context(NULL, NULL, PARDES_PDF_STORE_LIMIT_BYTES);
    if (ctx == NULL) {
        free(state);
        return NULL;
    }

    fz_var(doc);
    fz_try(ctx)
    {
        fz_register_document_handlers(ctx);
        doc = fz_open_document(ctx, path);
        pages = fz_count_pages(ctx, doc);
        if (pages < 1)
            fz_throw(ctx, FZ_ERROR_FORMAT, "PDF has no pages");
    }
    fz_catch(ctx)
    {
        fz_report_error(ctx);
        fz_drop_document(ctx, doc);
        fz_drop_context(ctx);
        free(state);
        return NULL;
    }

    state->ctx = ctx;
    state->doc = doc;
    state->page_count = pages;
    *page_count = pages;
    return state;
}

void
pardes_pdf_close(pardes_pdf_document *document)
{
    if (document == NULL)
        return;
    pardes_pdf_drop_cached_page(document);
    fz_drop_document(document->ctx, document->doc);
    fz_drop_context(document->ctx);
    free(document);
}

int
pardes_pdf_render(
    pardes_pdf_document *document,
    int page_number,
    int max_dimension,
    pardes_pdf_pixmap *out)
{
    return pardes_pdf_render_with_highlights(
        document, page_number, max_dimension, NULL, 0, out);
}

int
pardes_pdf_render_with_highlights(
    pardes_pdf_document *document,
    int page_number,
    int max_dimension,
    const pardes_pdf_highlight *highlights,
    size_t highlight_count,
    pardes_pdf_pixmap *out)
{
    fz_context *ctx;
    fz_pixmap *pixmap = NULL;
    fz_device *device = NULL;
    fz_path *path = NULL;
    fz_matrix ctm;
    fz_rect bounds;
    float longest;
    float scale = 2.0f; /* 144 dpi; PDF user space is 72 dpi. */
    size_t i;

    if (document == NULL || out == NULL || page_number < 0 ||
        page_number >= document->page_count || max_dimension < 1 ||
        (highlight_count != 0 && highlights == NULL) ||
        highlight_count > PARDES_PDF_MAX_RESULT_QUADS)
        return PARDES_PDF_ERROR;

    memset(out, 0, sizeof(*out));
    ctx = document->ctx;
    fz_var(pixmap);
    fz_var(device);
    fz_var(path);
    fz_try(ctx)
    {
        pardes_pdf_cache_page(document, page_number, 0);
        bounds = document->cached_bounds;
        longest = fmaxf(bounds.x1 - bounds.x0, bounds.y1 - bounds.y0);
        if (!(longest > 0.0f))
            fz_throw(ctx, FZ_ERROR_FORMAT, "PDF page has empty bounds");
        if (longest * scale > (float)max_dimension)
            scale = (float)max_dimension / longest;
        /*
         * Unlike a bare scale, fz_transform_page translates non-zero crop
         * boxes to the pixmap origin and rounds both dimensions coherently.
         * The overlay draw device receives this exact CTM, so normalized page
         * geometry and page pixels stay aligned for offset and rotated pages.
         */
        ctm = fz_transform_page(bounds, scale * 72.0f, 0.0f);

        /*
         * alpha=0 is intentional. MuPDF clears this RGB pixmap to opaque
         * white before drawing. Asking for alpha=1 instead produces
         * premultiplied transparent samples, while Pardes consumes straight
         * RGBA after the Zig-side RGB expansion.
         */
        pixmap = fz_new_pixmap_from_page(
            ctx, document->cached_page, ctm, fz_device_rgb(ctx), 0);
        if (fz_pixmap_width(ctx, pixmap) < 1 ||
            fz_pixmap_height(ctx, pixmap) < 1)
            fz_throw(ctx, FZ_ERROR_FORMAT, "PDF page rendered empty");

        if (highlight_count != 0) {
            device = fz_new_draw_device(ctx, ctm, pixmap);
            for (i = 0; i < highlight_count; ++i) {
                fz_quad quad;
                float color[3];
                float alpha;

                if (!pardes_pdf_valid_quad(highlights[i].quad))
                    fz_throw(ctx, FZ_ERROR_ARGUMENT,
                        "PDF highlight contains a non-finite point");
                if (highlights[i].rgba[3] == 0)
                    continue;

                /* Kind is semantic metadata; caller-supplied RGBA is final. */
                (void)highlights[i].kind;
                quad = pardes_pdf_quad_to_page(document, highlights[i].quad);
                color[0] = highlights[i].rgba[0] / 255.0f;
                color[1] = highlights[i].rgba[1] / 255.0f;
                color[2] = highlights[i].rgba[2] / 255.0f;
                alpha = highlights[i].rgba[3] / 255.0f;

                path = fz_new_path(ctx);
                fz_moveto(ctx, path, quad.ul.x, quad.ul.y);
                fz_lineto(ctx, path, quad.ur.x, quad.ur.y);
                fz_lineto(ctx, path, quad.lr.x, quad.lr.y);
                fz_lineto(ctx, path, quad.ll.x, quad.ll.y);
                fz_closepath(ctx, path);
                fz_fill_path(ctx, device, path, 0, fz_identity,
                    fz_device_rgb(ctx), color, alpha,
                    fz_default_color_params);
                fz_drop_path(ctx, path);
                path = NULL;
            }
            fz_close_device(ctx, device);
        }
    }
    fz_always(ctx)
    {
        fz_drop_path(ctx, path);
        fz_drop_device(ctx, device);
    }
    fz_catch(ctx)
    {
        fz_report_error(ctx);
        fz_drop_pixmap(ctx, pixmap);
        return PARDES_PDF_ERROR;
    }

    out->handle = pixmap;
    out->samples = fz_pixmap_samples(ctx, pixmap);
    out->width = fz_pixmap_width(ctx, pixmap);
    out->height = fz_pixmap_height(ctx, pixmap);
    out->stride = fz_pixmap_stride(ctx, pixmap);
    out->components = fz_pixmap_components(ctx, pixmap);
    return PARDES_PDF_OK;
}

void
pardes_pdf_drop_pixmap(pardes_pdf_document *document, void *pixmap)
{
    if (document != NULL && pixmap != NULL)
        fz_drop_pixmap(document->ctx, (fz_pixmap *)pixmap);
}

int
pardes_pdf_page_text(
    pardes_pdf_document *document,
    int page_number,
    pardes_pdf_text *out)
{
    fz_context *ctx;
    fz_buffer *buffer = NULL;
    unsigned char *data = NULL;
    size_t len = 0;

    if (document == NULL || out == NULL || page_number < 0 ||
        page_number >= document->page_count)
        return PARDES_PDF_ERROR;

    memset(out, 0, sizeof(*out));
    ctx = document->ctx;
    fz_var(buffer);
    fz_try(ctx)
    {
        pardes_pdf_cache_page(document, page_number, 1);
        buffer = fz_new_buffer_from_stext_page(ctx, document->cached_text);
        len = fz_buffer_storage(ctx, buffer, &data);
    }
    fz_catch(ctx)
    {
        fz_report_error(ctx);
        fz_drop_buffer(ctx, buffer);
        return PARDES_PDF_ERROR;
    }

    out->handle = buffer;
    out->data = data;
    out->len = len;
    return PARDES_PDF_OK;
}

void
pardes_pdf_drop_text(pardes_pdf_document *document, void *buffer)
{
    if (document != NULL && buffer != NULL)
        fz_drop_buffer(document->ctx, (fz_buffer *)buffer);
}

int
pardes_pdf_search_page(
    pardes_pdf_document *document,
    int page_number,
    const char *needle,
    pardes_pdf_search_result *out)
{
    fz_context *ctx;
    fz_search *search = NULL;
    pardes_pdf_search_quad *items = NULL;
    size_t capacity = 0;
    size_t quad_count = 0;
    size_t hit_count = 0;
    int status = PARDES_PDF_ERROR;

    if (document == NULL || needle == NULL || out == NULL || page_number < 0 ||
        page_number >= document->page_count)
        return PARDES_PDF_ERROR;

    memset(out, 0, sizeof(*out));
    ctx = document->ctx;
    fz_var(search);
    fz_var(items);
    fz_var(status);
    fz_try(ctx)
    {
        pardes_pdf_cache_page(document, page_number, 1);
        search = fz_new_search(ctx);
        fz_search_set_options(
            ctx, search, FZ_SEARCH_IGNORE_CASE, needle);

        /* fz_feed_search takes ownership; retain the document's cache ref. */
        fz_feed_search(ctx, search,
            fz_keep_stext_page(ctx, document->cached_text), page_number);

        for (;;) {
            fz_search_result result = fz_search_forwards(ctx, search);
            if (result.reason == FZ_SEARCH_COMPLETE)
                break;
            if (result.reason == FZ_SEARCH_MORE_INPUT) {
                /* This is deliberately a one-page search. */
                fz_feed_search(ctx, search, NULL,
                    result.u.more_input.seq_needed);
                continue;
            }
            if (result.reason == FZ_SEARCH_MATCH) {
                fz_search_result_details *details = result.u.match.result;
                size_t match_quads;
                size_t needed;
                size_t next_capacity;
                size_t i;

                if (details == NULL || details->num_quads < 0 ||
                    (details->num_quads > 0 && details->quads == NULL))
                    fz_throw(ctx, FZ_ERROR_FORMAT,
                        "MuPDF returned malformed search geometry");
                if (hit_count >= PARDES_PDF_MAX_RESULT_QUADS) {
                    status = PARDES_PDF_LIMIT_EXCEEDED;
                    fz_throw(ctx, FZ_ERROR_LIMIT,
                        "PDF search result limit exceeded");
                }

                match_quads = (size_t)details->num_quads;
                if (match_quads > PARDES_PDF_MAX_RESULT_QUADS - quad_count) {
                    status = PARDES_PDF_LIMIT_EXCEEDED;
                    fz_throw(ctx, FZ_ERROR_LIMIT,
                        "PDF search quad limit exceeded");
                }
                needed = quad_count + match_quads;
                if (needed > capacity) {
                    next_capacity = pardes_pdf_grown_capacity(
                        capacity, needed, PARDES_PDF_MAX_RESULT_QUADS);
                    items = fz_realloc_array(
                        ctx, items, next_capacity, pardes_pdf_search_quad);
                    capacity = next_capacity;
                }
                for (i = 0; i < match_quads; ++i) {
                    items[quad_count + i].quad = pardes_pdf_quad_from_page(
                        document, details->quads[i].quad);
                    items[quad_count + i].hit = hit_count;
                }
                quad_count = needed;
                ++hit_count;
                continue;
            }
            fz_throw(ctx, FZ_ERROR_FORMAT,
                "MuPDF returned an unknown search state");
        }
        status = PARDES_PDF_OK;
    }
    fz_always(ctx)
    {
        fz_drop_search(ctx, search);
    }
    fz_catch(ctx)
    {
        if (status != PARDES_PDF_LIMIT_EXCEEDED)
            fz_report_error(ctx);
        fz_free(ctx, items);
        return status;
    }

    out->handle = items;
    out->quads = items;
    out->quad_count = quad_count;
    out->hit_count = hit_count;
    return PARDES_PDF_OK;
}

void
pardes_pdf_drop_search_result(pardes_pdf_document *document, void *result)
{
    if (document != NULL && result != NULL)
        fz_free(document->ctx, result);
}

int
pardes_pdf_select(
    pardes_pdf_document *document,
    int page_number,
    pardes_pdf_point start,
    pardes_pdf_point end,
    pardes_pdf_selection *out)
{
    fz_context *ctx;
    fz_quad *raw = NULL;
    pardes_pdf_quad *items = NULL;
    fz_point a;
    fz_point b;
    int capacity = PARDES_PDF_INITIAL_RESULT_QUADS;
    int count = 0;
    int status = PARDES_PDF_ERROR;
    int i;

    if (document == NULL || out == NULL || page_number < 0 ||
        page_number >= document->page_count ||
        !pardes_pdf_valid_point(start) || !pardes_pdf_valid_point(end))
        return PARDES_PDF_ERROR;

    memset(out, 0, sizeof(*out));
    ctx = document->ctx;
    fz_var(raw);
    fz_var(items);
    fz_var(status);
    fz_try(ctx)
    {
        pardes_pdf_cache_page(document, page_number, 1);
        a = pardes_pdf_point_to_page(document, start);
        b = pardes_pdf_point_to_page(document, end);
        (void)fz_snap_selection(
            ctx, document->cached_text, &a, &b, FZ_SELECT_WORDS);

        raw = fz_malloc_array(ctx, capacity, fz_quad);
        for (;;) {
            int next_capacity;
            count = fz_highlight_selection(
                ctx, document->cached_text, a, b, raw, capacity);
            if (count < 0)
                fz_throw(ctx, FZ_ERROR_FORMAT,
                    "MuPDF returned an invalid selection count");
            if (count < capacity)
                break;

            /* One extra slot distinguishes exactly-at-limit from truncation. */
            if (capacity >= PARDES_PDF_MAX_RESULT_QUADS + 1) {
                status = PARDES_PDF_LIMIT_EXCEEDED;
                fz_throw(ctx, FZ_ERROR_LIMIT,
                    "PDF selection quad limit exceeded");
            }
            next_capacity = capacity * 2;
            if (next_capacity > PARDES_PDF_MAX_RESULT_QUADS + 1)
                next_capacity = PARDES_PDF_MAX_RESULT_QUADS + 1;
            raw = fz_realloc_array(ctx, raw, next_capacity, fz_quad);
            capacity = next_capacity;
        }
        if (count > PARDES_PDF_MAX_RESULT_QUADS) {
            status = PARDES_PDF_LIMIT_EXCEEDED;
            fz_throw(ctx, FZ_ERROR_LIMIT,
                "PDF selection quad limit exceeded");
        }

        if (count != 0)
            items = fz_malloc_array(ctx, (size_t)count, pardes_pdf_quad);
        for (i = 0; i < count; ++i)
            items[i] = pardes_pdf_quad_from_page(document, raw[i]);
        status = PARDES_PDF_OK;
    }
    fz_always(ctx)
    {
        fz_free(ctx, raw);
    }
    fz_catch(ctx)
    {
        if (status != PARDES_PDF_LIMIT_EXCEEDED)
            fz_report_error(ctx);
        fz_free(ctx, items);
        return status;
    }

    out->handle = items;
    out->quads = items;
    out->quad_count = (size_t)count;
    out->start = pardes_pdf_point_from_page(document, a);
    out->end = pardes_pdf_point_from_page(document, b);
    return PARDES_PDF_OK;
}

void
pardes_pdf_drop_selection(pardes_pdf_document *document, void *selection)
{
    if (document != NULL && selection != NULL)
        fz_free(document->ctx, selection);
}

int
pardes_pdf_point_inside_quad(pardes_pdf_point point, pardes_pdf_quad quad)
{
    fz_quad native;

    if (!pardes_pdf_valid_point(point) || !pardes_pdf_valid_quad(quad))
        return 0;
    native.ul = fz_make_point(quad.ul.x, quad.ul.y);
    native.ur = fz_make_point(quad.ur.x, quad.ur.y);
    native.ll = fz_make_point(quad.ll.x, quad.ll.y);
    native.lr = fz_make_point(quad.lr.x, quad.lr.y);
    return fz_is_point_inside_quad(
        fz_make_point(point.x, point.y), native) != 0;
}

int
pardes_pdf_copy_selection(
    pardes_pdf_document *document,
    int page_number,
    pardes_pdf_point start,
    pardes_pdf_point end,
    pardes_pdf_owned_text *out)
{
    fz_context *ctx;
    char *text = NULL;
    fz_point a;
    fz_point b;

    if (document == NULL || out == NULL || page_number < 0 ||
        page_number >= document->page_count ||
        !pardes_pdf_valid_point(start) || !pardes_pdf_valid_point(end))
        return PARDES_PDF_ERROR;

    memset(out, 0, sizeof(*out));
    ctx = document->ctx;
    fz_var(text);
    fz_try(ctx)
    {
        pardes_pdf_cache_page(document, page_number, 1);
        a = pardes_pdf_point_to_page(document, start);
        b = pardes_pdf_point_to_page(document, end);
        text = fz_copy_selection(ctx, document->cached_text, a, b, 0);
    }
    fz_catch(ctx)
    {
        fz_report_error(ctx);
        fz_free(ctx, text);
        return PARDES_PDF_ERROR;
    }

    out->handle = text;
    out->data = (const unsigned char *)text;
    out->len = text == NULL ? 0 : strlen(text);
    return PARDES_PDF_OK;
}

void
pardes_pdf_drop_owned_text(pardes_pdf_document *document, void *text)
{
    if (document != NULL && text != NULL)
        fz_free(document->ctx, text);
}