/*****************************************************************************

 * chroma_probe.c: chroma conversion probing

 *****************************************************************************

 * Copyright (C) 2025 VLC authors and VideoLAN

 *

 * This program is free software; you can redistribute it and/or modify it

 * under the terms of the GNU Lesser General Public License as published by

 * the Free Software Foundation; either version 2.1 of the License, or

 * (at your option) any later version.

 *

 * This program is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

 * GNU Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with this program; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.

 *****************************************************************************/

#ifdef HAVE_CONFIG_H

# include "config.h"

#endif

#include <vlc_chroma_probe.h>

#include <vlc_fourcc.h>

#include <vlc_threads.h>

#include <vlc_modules.h>

#include <vlc_sort.h>

#include <vlc_memstream.h>

static int

modules_Probe(vlc_chroma_conv_vec *chroma_table)

{

    module_t **mods;

    ssize_t total = vlc_module_match("chroma probe", NULL, false, &mods, NULL);

    if (total == -1)

        return -ENOENT;

    for (ssize_t i = 0; i < total; ++i)

    {

        vlc_chroma_conv_probe fn = vlc_module_map(NULL, mods[i]);

        if (fn == NULL)

            continue;

        fn(chroma_table);

    }

    free(mods);

    return 0;

}

/* Breadth First Search (BFS) node */

struct bfs_node

{

    vlc_fourcc_t chain[VLC_CHROMA_CONV_CHAIN_COUNT_MAX];

    unsigned depth; /* Max deep is VLC_CHROMA_CONV_CHAIN_COUNT_MAX -1 */

    float cost_factor;

};

typedef struct VLC_VECTOR(struct bfs_node) bfs_queue_vec;

static int

bfs_Run(vlc_fourcc_t chroma_from, vlc_fourcc_t chroma_to, unsigned max_depth,

        const vlc_chroma_conv_vec *chroma_table, int flags,

        bfs_queue_vec *queue)

{

    struct bfs_node start = {

        .chain[0] = chroma_from,

        .cost_factor = 1,

        .depth = 0,

    };

    bool success = vlc_vector_push(queue, start);

    if (!success)

        return -ENOMEM;

    for (size_t queue_idx = 0; queue_idx < queue->size; queue_idx++)

    {

        const struct bfs_node current = queue->data[queue_idx];

        vlc_fourcc_t current_chroma = current.chain[current.depth];

        if (chroma_to != 0 && current_chroma == chroma_to)

            continue; /* Found a path to 'chroma_to' */

        if (current.depth == max_depth)

            continue;

        /* Enqueue neighbors */

        for (size_t chroma_idx = 0; chroma_idx < chroma_table->size; chroma_idx++)

        {

            struct vlc_chroma_conv_entry *entry = &chroma_table->data[chroma_idx];

            vlc_fourcc_t from = entry->in;

            vlc_fourcc_t to = entry->out;

            float cost_factor = entry->cost_factor;

            if (from == current_chroma)

            {

                vlc_fourcc_t next_chroma = to;

                /* Apply filters from flags */

                if (flags & VLC_CHROMA_CONV_FLAG_ONLY_YUV)

                {

                    if (!vlc_fourcc_IsYUV(next_chroma))

                        continue;

                }

                else if (flags & VLC_CHROMA_CONV_FLAG_ONLY_RGB)

                {

                    const vlc_chroma_description_t *desc =

                        vlc_fourcc_GetChromaDescription(next_chroma);

                    if (desc == NULL || desc->subtype != VLC_CHROMA_SUBTYPE_RGB)

                        continue;

                }

                /* If next_chroma is already in the chain at any previous step,

                 * we've encountered a cycle or a duplicate. */

                bool already_visited = false;

                for (size_t i = 0; i < current.depth; ++i)

                    if (current.chain[i] == next_chroma)

                    {

                        already_visited = true;

                        break;

                    }

                if (already_visited)

                    continue;

                struct bfs_node next = current;

                next.depth = current.depth + 1;

                next.cost_factor = current.cost_factor * cost_factor;

                next.chain[next.depth] = next_chroma;

                success = vlc_vector_push(queue, next);

                if (!success)

                    return -ENOMEM;

            }

        }

    }

    return 0;

}

static uint64_t

GetChromaBits(const vlc_chroma_description_t *desc,

              unsigned width, unsigned height)

{

    if (desc->plane_count == 0)

    {

        /* Fallback to the size of the subtype */

        switch (desc->subtype)

        {

            case VLC_CHROMA_SUBTYPE_OTHER:

                return 0;

            case VLC_CHROMA_SUBTYPE_YUV444:

                return width * height * 3 * desc->color_bits;

            case VLC_CHROMA_SUBTYPE_YUV440:

            case VLC_CHROMA_SUBTYPE_YUV422:

                return width * height * 2 * desc->color_bits;

            case VLC_CHROMA_SUBTYPE_YUV420:

            case VLC_CHROMA_SUBTYPE_YUV411:

                return width * height * 1.5 * desc->color_bits;

            case VLC_CHROMA_SUBTYPE_YUV410:

                return width * height * 1.125 * desc->color_bits;

            case VLC_CHROMA_SUBTYPE_YUV211:

            case VLC_CHROMA_SUBTYPE_GREY:

                return width * height * desc->color_bits;

            case VLC_CHROMA_SUBTYPE_RGB:

                return width * height * 4 * desc->color_bits;

            default:

                vlc_assert_unreachable();

        }

    }

    uint64_t total_bits = 0;

    for (unsigned i = 0; i < desc->plane_count; i++)

    {

        const vlc_rational_t rw = desc->p[i].w;

        const vlc_rational_t rh = desc->p[i].h;

        unsigned plane_width = width * rw.num / rw.den;

        unsigned plane_height = height * rh.num / rh.den;

        uint64_t plane_pixels = plane_width * plane_height;

        uint64_t plane_bits = plane_pixels * desc->pixel_bits;

        total_bits += plane_bits;

    }

    return total_bits;

}

static float

GetColorRatio(enum vlc_chroma_subtype subtype)

{

    switch (subtype)

    {

        case VLC_CHROMA_SUBTYPE_YUV444:

        case VLC_CHROMA_SUBTYPE_RGB:

            return 1.f;

        case VLC_CHROMA_SUBTYPE_YUV422:

            return 0.67;

        case VLC_CHROMA_SUBTYPE_YUV440:

            return 0.5; /* should be like YUV422, but it is less common */

        case VLC_CHROMA_SUBTYPE_YUV420:

            return 0.5;

        case VLC_CHROMA_SUBTYPE_YUV411:

            return 0.33;

        case VLC_CHROMA_SUBTYPE_YUV410:

            return 0.25;

        case VLC_CHROMA_SUBTYPE_YUV211:

        case VLC_CHROMA_SUBTYPE_OTHER:

            return 0.2;

        case VLC_CHROMA_SUBTYPE_GREY:

            return 0.1;

        default:

            vlc_assert_unreachable();

    }

}

static float

CompareDescs(const vlc_chroma_description_t *in_desc,

             const vlc_chroma_description_t *out_desc)

{

    /* Compare color bits */

    float bits_ratio;

    if (in_desc->color_bits == 0 || out_desc->color_bits == 0)

        bits_ratio = 1.f;

    else

    {

        bits_ratio = out_desc->color_bits / in_desc->color_bits;

        if (bits_ratio > 1.f)

            bits_ratio = 1.f;

    }

    /* Compare color ratios, favor same or near subtype */

    if (in_desc->subtype == out_desc->subtype)

        return bits_ratio;

    float color_ratio = GetColorRatio(out_desc->subtype)

                      / GetColorRatio(in_desc->subtype);

    if (color_ratio > 1.f)

        color_ratio = 1.f;

    /* Malus for CPU YUV <-> Other. Favor staying in the same color model. */

    bool in_is_yuv = vlc_chroma_description_IsYUV(in_desc);

    bool out_is_yuv = vlc_chroma_description_IsYUV(out_desc);

    if ((in_desc->plane_count != 0 && out_desc->plane_count != 0)

     && (in_is_yuv || out_is_yuv) && (in_is_yuv != out_is_yuv))

        color_ratio *= 0.9;

    return color_ratio * bits_ratio;

}

static void

vlc_chroma_conv_result_FromNode(struct vlc_chroma_conv_result *res,

                                const struct bfs_node *node,

                                unsigned width, unsigned height)

{

    res->chain_count = node->depth + 1;

    res->cost = 0;

    uint64_t total_cost = 0;

    float total_quality = 1.f;

    for (size_t i = 0; i < res->chain_count; ++i)

    {

        res->chain[i] = node->chain[i];

        if (i > 0)

        {

            const vlc_chroma_description_t *from_desc =

                vlc_fourcc_GetChromaDescription(res->chain[i - 1]);

            const vlc_chroma_description_t *to_desc =

                vlc_fourcc_GetChromaDescription(res->chain[i]);

            if (from_desc == NULL || to_desc == NULL)

            {

                /* Unlikely, fallback for a big cost */

                total_cost += width * height * 4 * 8 * node->cost_factor;

                continue;

            }

            uint64_t from_bits = GetChromaBits(from_desc, width, height);

            uint64_t to_bits = GetChromaBits(to_desc, width, height);

            /* Unlikely case */

            if (from_bits == 0) /* OTHER -> ANY */

                from_bits = to_bits;

            else if (to_bits == 0) /* ANY -> OTHER */

                to_bits = from_bits;

            total_cost += (from_bits + to_bits) * node->cost_factor;

            float quality = CompareDescs(from_desc, to_desc);

            assert(quality > 0.f && quality <= 1.f);

            total_quality *= quality;

        }

    }

    res->cost = total_cost / width / height;

    res->quality = 100 * total_quality;

}

static int

SortResults(const void *a, const void *b, void *arg)

{

    const struct vlc_chroma_conv_result *ra = a;

    const struct vlc_chroma_conv_result *rb = b;

    bool *sort_by_quality = arg;

    int cost_score = 0, quality_score = 0;

    /* Lower cost is better */

    if (ra->cost < rb->cost)

        cost_score = -1;

    else if (ra->cost > rb->cost)

        cost_score = 1;

    /* Higher Quality is better */

    if (ra->quality > rb->quality)

        quality_score = -1;

    else if (ra->quality < rb->quality)

        quality_score = 1;

    /* Fallback to secondary score in same score */

    if (*sort_by_quality)

        return quality_score != 0 ? quality_score : cost_score;

    else

        return cost_score != 0 ? cost_score : quality_score;

}

static bool

bfs_node_IsResult(const struct bfs_node *node, vlc_fourcc_t to)

{

    vlc_fourcc_t current_chroma = node->chain[node->depth];

    return to == 0 || current_chroma == to;

}

static bool

vlc_chroma_conv_result_Equals(struct vlc_chroma_conv_result *a,

                              struct vlc_chroma_conv_result *b)

{

    if (a->chain_count != b->chain_count)

        return false;

    if (a->quality != b->quality)

        return false;

    /* Don't check cost since we want to merge results with different costs */

    for (size_t i = 0; i < a->chain_count; ++i)

        if (a->chain[i] != b->chain[i])

            return false;

    return true;

}

struct vlc_chroma_conv_result *

vlc_chroma_conv_Probe(vlc_fourcc_t from, vlc_fourcc_t to,

                      unsigned width, unsigned height,

                      unsigned max_indirect_steps, int flags, size_t *count)

{

    assert(from != 0);

    assert(max_indirect_steps <= VLC_CHROMA_CONV_MAX_INDIRECT_STEPS);

    vlc_chroma_conv_vec chroma_table;

    vlc_vector_init(&chroma_table);

    if (width == 0 || height == 0)

    {

        width = 3840;

        height = 2160;

    }

    if (max_indirect_steps > 0)

    {

        /* Allow indirect steps only when converting from/to a GPU chroma */

        bool from_cpu = vlc_fourcc_GetChromaBPP(from) != 0;

        bool to_cpu = to == 0 ? true : vlc_fourcc_GetChromaBPP(to) != 0;

        if (from_cpu && to_cpu)

            max_indirect_steps--;

    }

    /* Probe modules */

    int ret = modules_Probe(&chroma_table);

    if (ret != 0 || chroma_table.size == 0)

    {

        vlc_vector_destroy(&chroma_table);

        return NULL;

    }

    /* Run tree search */

    bfs_queue_vec bfs_queue;

    vlc_vector_init(&bfs_queue);

    ret = bfs_Run(from, to, max_indirect_steps + 1 , &chroma_table, flags,

                  &bfs_queue);

    vlc_vector_destroy(&chroma_table);

    size_t result_count = 0;

    for (size_t i = 1 /* skip start node */; i < bfs_queue.size; ++i)

        if (bfs_node_IsResult(&bfs_queue.data[i], to))

            result_count++;

    if (unlikely(ret != 0) || result_count == 0)

    {

        vlc_vector_destroy(&bfs_queue);

        return NULL;

    }

    /* Allocate the result array */

    struct VLC_VECTOR(struct vlc_chroma_conv_result) result_vec =

        VLC_VECTOR_INITIALIZER;

    bool success = vlc_vector_push_hole(&result_vec, result_count);

    if (!success)

    {

        vlc_vector_destroy(&bfs_queue);

        return NULL;

    }

    /* Fill the result from the tree search */

    size_t res_idx = 0;

    for (size_t i = 1 /* skip start node */; i < bfs_queue.size; ++i)

    {

        const struct bfs_node *node = &bfs_queue.data[i];

        if (!bfs_node_IsResult(node, to))

            continue;

        assert(res_idx < result_count);

        struct vlc_chroma_conv_result *res = &result_vec.data[res_idx++];

        vlc_chroma_conv_result_FromNode(res, node, width, height);

    }

    assert(res_idx == result_count);

    vlc_vector_destroy(&bfs_queue);

    /* Sort */

    bool sort_by_quality = (flags & VLC_CHROMA_CONV_FLAG_SORT_COST) == 0;

    vlc_qsort(result_vec.data, result_count,

              sizeof(struct vlc_chroma_conv_result), SortResults,

              &sort_by_quality);

    /* Remove duplicate entries, it can happen when more the 2 modules probe

     * the same conversion. They are not necessarily one after the other as

     * they might have different quality. */

    for (size_t i = 0; i < result_vec.size - 1; ++ i)

    {

        struct vlc_chroma_conv_result *cur = &result_vec.data[i];

        size_t j = i + 1;

        while (j < result_vec.size)

        {

            struct vlc_chroma_conv_result *next = &result_vec.data[j];

            if (vlc_chroma_conv_result_Equals(cur, next))

            {

                /* Keep the lowest cost */

                if (next->cost < cur->cost)

                    cur->cost = next->cost;

                vlc_vector_remove(&result_vec, j);

                /* update pointer after possible realloc */

                cur = &result_vec.data[i];

            }

            else

                j++;

        }

    }

    *count = result_vec.size;

    return result_vec.data;

}

char *

vlc_chroma_conv_result_ToString(const struct vlc_chroma_conv_result *res)

{

    struct vlc_memstream ms;

    int ret = vlc_memstream_open(&ms);

    if (ret != 0)

        return NULL;

    vlc_memstream_printf(&ms, "[c=%u|q=%u] ", res->cost, res->quality);

    for (size_t i = 0; i < res->chain_count; ++i)

    {

        vlc_memstream_printf(&ms, "%4.4s", (const char *) &res->chain[i]);

        if (i != res->chain_count - 1)

            vlc_memstream_puts(&ms, " -> ");

    }

    ret = vlc_memstream_close(&ms);

    return ret == 0 ? ms.ptr : NULL;

}