/******************************************************************************
*
* Copyright (C) 2026  Brian Berlin
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 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 General Public License for more details.
* 
* You should have received a copy of the GNU General Public License
* along with this program.  If not, see <http://www.gnu.org/licenses/>
*
******************************************************************************/
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "shader_bundle.h"
#include "toml.h"
#include "util.h"

static char * shader_bundle_strdup ( const char *value ) {

    char *copy;

    if ( !value )
        return NULL;

    copy = malloc ( strlen(value) + 1 );
    if ( !copy )
        return NULL;

    strcpy ( copy, value );
    return copy;

} /* end shader_bundle_strdup */

static char * shader_bundle_dup_key ( t_toml_node *tbl, const char *key ) {

    return shader_bundle_strdup ( toml_get_key_value(tbl, key) );

} /* end shader_bundle_dup_key */

static void shader_bundle_read_int_key ( t_toml_node *tbl, const char *key,
        int *value ) {

    char *text;
    char *endptr;
    long result;

    text = toml_get_key_value ( tbl, key );
    if ( !text )
        return;

    endptr = text;
    result = strtol ( text, &endptr, 0 );
    if ( endptr != text )
        *value = (int) result;

} /* end shader_bundle_read_int_key */

static void shader_bundle_read_float_key ( t_toml_node *tbl, const char *key,
        float *value ) {

    char *text;
    char *endptr;
    float result;

    text = toml_get_key_value ( tbl, key );
    if ( !text )
        return;

    endptr = text;
    result = strtof ( text, &endptr );
    if ( endptr != text )
        *value = result;

} /* end shader_bundle_read_float_key */

static int shader_bundle_table_matches ( t_toml_node *node, const char *name ) {

    return node && ( node->type == TOML_TABLE || node->type == TOML_TABLE_ARY ) &&
           node->value && !strcmp ( node->value, name );

} /* end shader_bundle_table_matches */

static int shader_bundle_count_child_tables ( t_toml_node *parent,
        const char *name ) {

    int count = 0;
    t_toml_node *node;

    if ( !parent )
        return 0;

    node = parent->child;
    while ( node ) {
        if ( shader_bundle_table_matches(node, name) )
            count++;
        node = node->next;
    }

    return count;

} /* end shader_bundle_count_child_tables */

static void shader_bundle_parse_uniform ( t_toml_node *tbl,
        t_shader_bundle_uniform *uniform ) {

    uniform->name = shader_bundle_dup_key ( tbl, "name" );
    uniform->target = shader_bundle_dup_key ( tbl, "target" );
    uniform->resource = shader_bundle_dup_key ( tbl, "resource" );
    uniform->type = shader_bundle_dup_key ( tbl, "type" );
    uniform->offset = -1;

    shader_bundle_read_int_key ( tbl, "offset", &uniform->offset );

} /* end shader_bundle_parse_uniform */

static int shader_bundle_parse_stage ( t_toml_node *tbl,
        t_shader_bundle_stage *stage ) {

    int uniform_index = 0;
    int sampler_index = 0;
    t_toml_node *node;

    stage->stage = shader_bundle_dup_key ( tbl, "stage" );
    stage->entrypoint = shader_bundle_dup_key ( tbl, "entrypoint" );
    stage->spirv = shader_bundle_dup_key ( tbl, "spirv" );
    stage->msl = shader_bundle_dup_key ( tbl, "msl" );
    stage->dxil = shader_bundle_dup_key ( tbl, "dxil" );
    stage->common_uniform_slot = -1;
    stage->param_uniform_slot = -1;

    shader_bundle_read_int_key ( tbl, "samplers", &stage->samplers );
    shader_bundle_read_int_key ( tbl, "storage_textures",
            &stage->storage_textures );
    shader_bundle_read_int_key ( tbl, "storage_buffers",
            &stage->storage_buffers );
    shader_bundle_read_int_key ( tbl, "uniform_buffers",
            &stage->uniform_buffers );
    shader_bundle_read_int_key ( tbl, "common_uniform_slot",
            &stage->common_uniform_slot );
    shader_bundle_read_int_key ( tbl, "common_uniform_size",
            &stage->common_uniform_size );
    shader_bundle_read_int_key ( tbl, "param_uniform_slot",
            &stage->param_uniform_slot );
    shader_bundle_read_int_key ( tbl, "param_uniform_size",
            &stage->param_uniform_size );

    stage->num_uniforms = shader_bundle_count_child_tables ( tbl,
            "pass.stage.uniform" );
    if ( stage->num_uniforms ) {
        stage->uniforms = calloc ( (size_t) stage->num_uniforms,
                sizeof(t_shader_bundle_uniform) );
        if ( !stage->uniforms )
            return -1;
    }

    stage->num_sampler_bindings = shader_bundle_count_child_tables ( tbl,
            "pass.stage.sampler" );
    if ( stage->num_sampler_bindings ) {
        stage->sampler_bindings = calloc (
                (size_t) stage->num_sampler_bindings,
                sizeof(t_shader_bundle_sampler) );
        if ( !stage->sampler_bindings )
            return -1;
    }

    node = tbl->child;
    while ( node ) {
        if ( shader_bundle_table_matches(node, "pass.stage.uniform") &&
                uniform_index < stage->num_uniforms ) {
            shader_bundle_parse_uniform ( node,
                    &stage->uniforms[uniform_index] );
            if ( !stage->uniforms[uniform_index].name ||
                    !stage->uniforms[uniform_index].resource ||
                    stage->uniforms[uniform_index].offset < 0 )
                return -1;
            uniform_index++;
        }
        else if ( shader_bundle_table_matches(node, "pass.stage.sampler") &&
                sampler_index < stage->num_sampler_bindings ) {
            stage->sampler_bindings[sampler_index].name =
                    shader_bundle_dup_key ( node, "name" );
            stage->sampler_bindings[sampler_index].slot = -1;
            shader_bundle_read_int_key ( node, "slot",
                    &stage->sampler_bindings[sampler_index].slot );
            if ( !stage->sampler_bindings[sampler_index].name ||
                    stage->sampler_bindings[sampler_index].slot < 0 )
                return -1;
            sampler_index++;
        }

        node = node->next;
    }

    if ( stage->common_uniform_slot >= stage->uniform_buffers ||
            stage->param_uniform_slot >= stage->uniform_buffers )
        return -1;
    if ( ( stage->common_uniform_slot >= 0 && stage->common_uniform_size <= 0 ) ||
            ( stage->param_uniform_slot >= 0 && stage->param_uniform_size <= 0 ) )
        return -1;

    return 0;

} /* end shader_bundle_parse_stage */

static void shader_bundle_parse_param ( t_toml_node *tbl,
        t_shader_bundle_param *param ) {

    param->name = shader_bundle_dup_key ( tbl, "name" );
    param->label = shader_bundle_dup_key ( tbl, "label" );
    param->group = shader_bundle_dup_key ( tbl, "group" );
    param->stage = shader_bundle_dup_key ( tbl, "stage" );
    param->target = shader_bundle_dup_key ( tbl, "target" );
    param->resource = shader_bundle_dup_key ( tbl, "resource" );
    param->type = shader_bundle_dup_key ( tbl, "type" );

    shader_bundle_read_float_key ( tbl, "default", &param->def );
    shader_bundle_read_float_key ( tbl, "min", &param->min );
    shader_bundle_read_float_key ( tbl, "max", &param->max );
    shader_bundle_read_float_key ( tbl, "step", &param->step );
    shader_bundle_read_int_key ( tbl, "divider", &param->divider );
    shader_bundle_read_int_key ( tbl, "offset", &param->offset );

    param->val = param->def;

} /* end shader_bundle_parse_param */

static int shader_bundle_parse_pass ( t_toml_node *tbl,
        t_shader_bundle_pass *pass ) {

    int stage_index = 0;
    int param_index = 0;
    t_toml_node *node;

    shader_bundle_read_int_key ( tbl, "index", &pass->index );
    pass->source = shader_bundle_dup_key ( tbl, "source" );
    pass->filter_linear = shader_bundle_dup_key ( tbl, "filter_linear" );
    pass->wrap_mode = shader_bundle_dup_key ( tbl, "wrap_mode" );
    pass->mipmap_input = shader_bundle_dup_key ( tbl, "mipmap_input" );
    pass->scale_type = shader_bundle_dup_key ( tbl, "scale_type" );
    pass->scale_type_x = shader_bundle_dup_key ( tbl, "scale_type_x" );
    pass->scale_type_y = shader_bundle_dup_key ( tbl, "scale_type_y" );
    pass->scale = shader_bundle_dup_key ( tbl, "scale" );
    pass->scale_x = shader_bundle_dup_key ( tbl, "scale_x" );
    pass->scale_y = shader_bundle_dup_key ( tbl, "scale_y" );
    pass->alias = shader_bundle_dup_key ( tbl, "alias" );
    pass->float_framebuffer = shader_bundle_dup_key ( tbl, "float_framebuffer" );
    pass->srgb_framebuffer = shader_bundle_dup_key ( tbl, "srgb_framebuffer" );

    pass->num_stages = shader_bundle_count_child_tables ( tbl, "pass.stage" );
    if ( pass->num_stages ) {
        pass->stages = calloc ( (size_t) pass->num_stages,
                sizeof(t_shader_bundle_stage) );
        if ( !pass->stages )
            return -1;
    }

    pass->num_params = shader_bundle_count_child_tables ( tbl, "pass.param" );
    if ( pass->num_params ) {
        pass->params = calloc ( (size_t) pass->num_params,
                sizeof(t_shader_bundle_param) );
        if ( !pass->params )
            return -1;
    }

    node = tbl->child;
    while ( node ) {
        if ( shader_bundle_table_matches(node, "pass.stage") &&
                stage_index < pass->num_stages ) {
            if ( shader_bundle_parse_stage ( node,
                        &pass->stages[stage_index] ) )
                return -1;
            stage_index++;
        }
        else if ( shader_bundle_table_matches(node, "pass.param") &&
                param_index < pass->num_params ) {
            shader_bundle_parse_param ( node, &pass->params[param_index] );
            param_index++;
        }

        node = node->next;
    }

    return 0;

} /* end shader_bundle_parse_pass */

void shader_bundle_init ( t_shader_bundle *bundle ) {

    if ( bundle )
        memset ( bundle, 0, sizeof(t_shader_bundle) );

} /* end shader_bundle_init */

void shader_bundle_free ( t_shader_bundle *bundle ) {

    if ( !bundle )
        return;

    for ( int i = 0; i < bundle->num_passes; ++i ) {
        t_shader_bundle_pass *pass = &bundle->passes[i];

        free ( pass->source );
        free ( pass->filter_linear );
        free ( pass->wrap_mode );
        free ( pass->mipmap_input );
        free ( pass->scale_type );
        free ( pass->scale_type_x );
        free ( pass->scale_type_y );
        free ( pass->scale );
        free ( pass->scale_x );
        free ( pass->scale_y );
        free ( pass->alias );
        free ( pass->float_framebuffer );
        free ( pass->srgb_framebuffer );

        for ( int j = 0; j < pass->num_stages; ++j ) {
            free ( pass->stages[j].stage );
            free ( pass->stages[j].entrypoint );
            free ( pass->stages[j].spirv );
            free ( pass->stages[j].msl );
            free ( pass->stages[j].dxil );
            for ( int k = 0; k < pass->stages[j].num_uniforms; ++k ) {
                free ( pass->stages[j].uniforms[k].name );
                free ( pass->stages[j].uniforms[k].target );
                free ( pass->stages[j].uniforms[k].resource );
                free ( pass->stages[j].uniforms[k].type );
            }
            free ( pass->stages[j].uniforms );
            for ( int k = 0; k < pass->stages[j].num_sampler_bindings; ++k )
                free ( pass->stages[j].sampler_bindings[k].name );
            free ( pass->stages[j].sampler_bindings );
        }
        free ( pass->stages );

        for ( int j = 0; j < pass->num_params; ++j ) {
            free ( pass->params[j].name );
            free ( pass->params[j].label );
            free ( pass->params[j].group );
            free ( pass->params[j].stage );
            free ( pass->params[j].target );
            free ( pass->params[j].resource );
            free ( pass->params[j].type );
        }
        free ( pass->params );
    }

    for ( int i = 0; i < bundle->num_textures; ++i ) {
        free ( bundle->textures[i].name );
        free ( bundle->textures[i].file );
        free ( bundle->textures[i].filter_linear );
        free ( bundle->textures[i].wrap_mode );
        free ( bundle->textures[i].mipmap );
    }
    free ( bundle->textures );

    free ( bundle->dir );
    free ( bundle->id );
    free ( bundle->source_hash );
    free ( bundle->glsl_to_spirv );
    free ( bundle->shadercross );
    free ( bundle->spirv_cross );
    free ( bundle->source_preset );
    free ( bundle->passes );

    shader_bundle_init ( bundle );

} /* end shader_bundle_free */

int shader_bundle_load ( const char *bundle_dir, t_shader_bundle *bundle ) {

    char filename[KAT_PATH_MAX];
    int pass_index = 0;
    int texture_index = 0;
    int result = -1;
    t_toml_node *tbl = NULL;
    t_toml_node *node;

    if ( !bundle_dir || !bundle )
        return -1;

    shader_bundle_init ( bundle );

    snprintf ( filename, sizeof(filename), "%s/bundle.toml", bundle_dir );
    tbl = toml_parse_file ( filename );
    if ( !tbl )
        goto done;

    bundle->dir = shader_bundle_strdup ( bundle_dir );
    bundle->id = shader_bundle_dup_key ( tbl, "id" );
    bundle->source_hash = shader_bundle_dup_key ( tbl, "source_hash" );
    bundle->glsl_to_spirv = shader_bundle_dup_key ( tbl, "glsl_to_spirv" );
    bundle->shadercross = shader_bundle_dup_key ( tbl, "shadercross" );
    bundle->spirv_cross = shader_bundle_dup_key ( tbl, "spirv_cross" );
    bundle->source_preset = shader_bundle_dup_key ( tbl, "source_preset" );

    if ( !bundle->id )
        goto done;

    bundle->num_passes = shader_bundle_count_child_tables ( tbl, "pass" );
    if ( bundle->num_passes ) {
        bundle->passes = calloc ( (size_t) bundle->num_passes,
                sizeof(t_shader_bundle_pass) );
        if ( !bundle->passes )
            goto done;
    }

    bundle->num_textures = shader_bundle_count_child_tables ( tbl, "texture" );
    if ( bundle->num_textures ) {
        bundle->textures = calloc ( (size_t) bundle->num_textures,
                sizeof(t_shader_bundle_texture) );
        if ( !bundle->textures )
            goto done;
    }

    node = tbl->child;
    while ( node ) {
        if ( shader_bundle_table_matches(node, "pass") &&
                pass_index < bundle->num_passes ) {
            if ( shader_bundle_parse_pass(node, &bundle->passes[pass_index]) )
                goto done;
            pass_index++;
        }
        else if ( shader_bundle_table_matches(node, "texture") &&
                texture_index < bundle->num_textures ) {
            t_shader_bundle_texture *texture =
                    &bundle->textures[texture_index];
            texture->name = shader_bundle_dup_key ( node, "name" );
            texture->file = shader_bundle_dup_key ( node, "file" );
            texture->filter_linear = shader_bundle_dup_key ( node, "linear" );
            texture->wrap_mode = shader_bundle_dup_key ( node, "wrap_mode" );
            texture->mipmap = shader_bundle_dup_key ( node, "mipmap" );
            if ( !texture->name || !texture->file )
                goto done;
            texture_index++;
        }

        node = node->next;
    }

    result = 0;

done:
    if ( tbl )
        toml_free_node_tree ( tbl );
    if ( result )
        shader_bundle_free ( bundle );

    return result;

} /* end shader_bundle_load */

t_shader_bundle_param * shader_bundle_find_param ( t_shader_bundle *bundle,
        int pass_index, const char *name ) {

    if ( !bundle || !name )
        return NULL;

    for ( int i = 0; i < bundle->num_passes; ++i ) {
        t_shader_bundle_pass *pass = &bundle->passes[i];

        if ( pass->index != pass_index )
            continue;

        for ( int j = 0; j < pass->num_params; ++j ) {
            if ( pass->params[j].name && !strcmp(pass->params[j].name, name) )
                return &pass->params[j];
        }
    }

    return NULL;

} /* end shader_bundle_find_param */

void shader_bundle_reset_params ( t_shader_bundle *bundle ) {

    if ( !bundle )
        return;

    for ( int i = 0; i < bundle->num_passes; ++i ) {
        t_shader_bundle_pass *pass = &bundle->passes[i];
        for ( int j = 0; j < pass->num_params; ++j )
            pass->params[j].val = pass->params[j].def;
    }

} /* end shader_bundle_reset_params */

void shader_bundle_clamp_params ( t_shader_bundle *bundle ) {

    if ( !bundle )
        return;

    for ( int i = 0; i < bundle->num_passes; ++i ) {
        t_shader_bundle_pass *pass = &bundle->passes[i];
        for ( int j = 0; j < pass->num_params; ++j ) {
            if ( pass->params[j].val < pass->params[j].min )
                pass->params[j].val = pass->params[j].min;
            if ( pass->params[j].val > pass->params[j].max )
                pass->params[j].val = pass->params[j].max;
        }
    }

} /* end shader_bundle_clamp_params */

int shader_bundle_param_is_default ( const t_shader_bundle_param *param ) {

    if ( !param )
        return 1;

    return fabsf ( param->val - param->def ) <= 0.000001f;

} /* end shader_bundle_param_is_default */

const char * shader_get_float_format ( float step ) {

    if ( step < 0.0099 )
        return "%+#05.3lf";
    else if ( step < 0.099 )
        return "%+#04.2lf";
    else if ( step < 0.999 )
        return "%+#03.1lf";
    else
        return "%+0.0lf";

} /* end shader_get_float_format */
