// luma.gl
// SPDX-License-Identifier: MIT
// Copyright (c) vis.gl contributors
import { glsl } from "../glsl-utils/highlight.js";
import { resolveModules } from "./resolve-modules.js";
import { getPlatformShaderDefines } from "./platform-defines.js";
import { injectShader, DECLARATION_INJECT_MARKER } from "./shader-injections.js";
import { transpileGLSLShader } from "../shader-transpiler/transpile-glsl-shader.js";
import { normalizeShaderHooks, getShaderHooks } from "./shader-hooks.js";
import { assert } from "../utils/assert.js";
import { getShaderInfo } from "../glsl-utils/get-shader-info.js";
const INJECT_SHADER_DECLARATIONS = `\n\n${DECLARATION_INJECT_MARKER}\n`;
/**
 * Precision prologue to inject before functions are injected in shader
 * TODO - extract any existing prologue in the fragment source and move it up...
 */
const FRAGMENT_SHADER_PROLOGUE = `\
precision highp float;
`;
/**
 * Inject a list of shader modules into a single shader source for WGSL
 */
export function assembleShaderWGSL(options) {
    const modules = resolveModules(options.modules || []);
    return {
        source: assembleWGSLShader(options.platformInfo, {
            ...options,
            source: options.source,
            stage: 'vertex',
            modules
        }),
        getUniforms: assembleGetUniforms(modules)
    };
}
/**
 * Injects dependent shader module sources into pair of main vertex/fragment shader sources for WGSL
 */
export function assembleShaderPairWGSL(options) {
    const modules = resolveModules(options.modules || []);
    return {
        vs: assembleWGSLShader(options.platformInfo, {
            ...options,
            source: options.vs,
            stage: 'vertex',
            modules
        }),
        fs: assembleWGSLShader(options.platformInfo, {
            ...options,
            source: options.fs,
            stage: 'fragment',
            modules
        }),
        getUniforms: assembleGetUniforms(modules)
    };
}
/**
 * Injects dependent shader module sources into pair of main vertex/fragment shader sources for GLSL
 */
export function assembleShaderPairGLSL(options) {
    const { vs, fs } = options;
    const modules = resolveModules(options.modules || []);
    return {
        vs: assembleGLSLShader(options.platformInfo, {
            ...options,
            source: vs,
            stage: 'vertex',
            modules
        }),
        fs: assembleGLSLShader(options.platformInfo, {
            ...options,
            source: fs,
            stage: 'fragment',
            modules
        }),
        getUniforms: assembleGetUniforms(modules)
    };
}
/**
 * Pulls together complete source code for either a vertex or a fragment shader
 * adding prologues, requested module chunks, and any final injections.
 * @param gl
 * @param options
 * @returns
 */
export function assembleWGSLShader(platformInfo, options) {
    const { 
    // id,
    source, stage, modules, 
    // defines = {},
    hookFunctions = [], inject = {}, log } = options;
    assert(typeof source === 'string', 'shader source must be a string');
    // const isVertex = type === 'vs';
    // const sourceLines = source.split('\n');
    const coreSource = source;
    // Combine Module and Application Defines
    // const allDefines = {};
    // modules.forEach(module => {
    //   Object.assign(allDefines, module.getDefines());
    // });
    // Object.assign(allDefines, defines);
    // Add platform defines (use these to work around platform-specific bugs and limitations)
    // Add common defines (GLSL version compatibility, feature detection)
    // Add precision declaration for fragment shaders
    let assembledSource = '';
    //   prologue
    //     ? `\
    // ${getShaderNameDefine({id, source, type})}
    // ${getShaderType(type)}
    // ${getPlatformShaderDefines(platformInfo)}
    // ${getApplicationDefines(allDefines)}
    // ${isVertex ? '' : FRAGMENT_SHADER_PROLOGUE}
    // `
    // `;
    const hookFunctionMap = normalizeShaderHooks(hookFunctions);
    // Add source of dependent modules in resolved order
    const hookInjections = {};
    const declInjections = {};
    const mainInjections = {};
    for (const key in inject) {
        const injection = typeof inject[key] === 'string' ? { injection: inject[key], order: 0 } : inject[key];
        const match = /^(v|f)s:(#)?([\w-]+)$/.exec(key);
        if (match) {
            const hash = match[2];
            const name = match[3];
            if (hash) {
                if (name === 'decl') {
                    declInjections[key] = [injection];
                }
                else {
                    mainInjections[key] = [injection];
                }
            }
            else {
                hookInjections[key] = [injection];
            }
        }
        else {
            // Regex injection
            mainInjections[key] = [injection];
        }
    }
    // TODO - hack until shadertool modules support WebGPU
    const modulesToInject = platformInfo.type !== 'webgpu' ? modules : [];
    for (const module of modulesToInject) {
        if (log) {
            module.checkDeprecations(coreSource, log);
        }
        const moduleSource = module.getModuleSource(stage, 'wgsl');
        // Add the module source, and a #define that declares it presence
        assembledSource += moduleSource;
        const injections = module.injections[stage];
        for (const key in injections) {
            const match = /^(v|f)s:#([\w-]+)$/.exec(key);
            if (match) {
                const name = match[2];
                const injectionType = name === 'decl' ? declInjections : mainInjections;
                injectionType[key] = injectionType[key] || [];
                injectionType[key].push(injections[key]);
            }
            else {
                hookInjections[key] = hookInjections[key] || [];
                hookInjections[key].push(injections[key]);
            }
        }
    }
    // For injectShader
    assembledSource += INJECT_SHADER_DECLARATIONS;
    assembledSource = injectShader(assembledSource, stage, declInjections);
    assembledSource += getShaderHooks(hookFunctionMap[stage], hookInjections);
    // Add the version directive and actual source of this shader
    assembledSource += coreSource;
    // Apply any requested shader injections
    assembledSource = injectShader(assembledSource, stage, mainInjections);
    return assembledSource;
}
/**
 * Pulls together complete source code for either a vertex or a fragment shader
 * adding prologues, requested module chunks, and any final injections.
 * @param gl
 * @param options
 * @returns
 */
function assembleGLSLShader(platformInfo, options) {
    const { id, source, stage, language = 'glsl', modules, defines = {}, hookFunctions = [], inject = {}, prologue = true, log } = options;
    assert(typeof source === 'string', 'shader source must be a string');
    const sourceVersion = language === 'glsl' ? getShaderInfo(source).version : -1;
    const targetVersion = platformInfo.shaderLanguageVersion;
    const sourceVersionDirective = sourceVersion === 100 ? '#version 100' : '#version 300 es';
    const sourceLines = source.split('\n');
    // TODO : keep all pre-processor statements at the beginning of the shader.
    const coreSource = sourceLines.slice(1).join('\n');
    // Combine Module and Application Defines
    const allDefines = {};
    modules.forEach(module => {
        Object.assign(allDefines, module.getDefines());
    });
    Object.assign(allDefines, defines);
    // Add platform defines (use these to work around platform-specific bugs and limitations)
    // Add common defines (GLSL version compatibility, feature detection)
    // Add precision declaration for fragment shaders
    let assembledSource = '';
    switch (language) {
        case 'wgsl':
            break;
        case 'glsl':
            assembledSource = prologue
                ? `\
${sourceVersionDirective}

// ----- PROLOGUE -------------------------
${getShaderNameDefine({ id, source, stage })}
${`#define SHADER_TYPE_${stage.toUpperCase()}`}
${getPlatformShaderDefines(platformInfo)}
${stage === 'fragment' ? FRAGMENT_SHADER_PROLOGUE : ''}

// ----- APPLICATION DEFINES -------------------------

${getApplicationDefines(allDefines)}

`
                : `${sourceVersionDirective}
`;
            break;
    }
    const hookFunctionMap = normalizeShaderHooks(hookFunctions);
    // Add source of dependent modules in resolved order
    const hookInjections = {};
    const declInjections = {};
    const mainInjections = {};
    for (const key in inject) {
        // @ts-expect-error
        const injection = typeof inject[key] === 'string' ? { injection: inject[key], order: 0 } : inject[key];
        const match = /^(v|f)s:(#)?([\w-]+)$/.exec(key);
        if (match) {
            const hash = match[2];
            const name = match[3];
            if (hash) {
                if (name === 'decl') {
                    declInjections[key] = [injection];
                }
                else {
                    mainInjections[key] = [injection];
                }
            }
            else {
                hookInjections[key] = [injection];
            }
        }
        else {
            // Regex injection
            mainInjections[key] = [injection];
        }
    }
    for (const module of modules) {
        if (log) {
            module.checkDeprecations(coreSource, log);
        }
        const moduleSource = module.getModuleSource(stage);
        // Add the module source, and a #define that declares it presence
        assembledSource += moduleSource;
        const injections = module.injections[stage];
        for (const key in injections) {
            const match = /^(v|f)s:#([\w-]+)$/.exec(key);
            if (match) {
                const name = match[2];
                const injectionType = name === 'decl' ? declInjections : mainInjections;
                injectionType[key] = injectionType[key] || [];
                injectionType[key].push(injections[key]);
            }
            else {
                hookInjections[key] = hookInjections[key] || [];
                hookInjections[key].push(injections[key]);
            }
        }
    }
    assembledSource += '// ----- MAIN SHADER SOURCE -------------------------';
    // For injectShader
    assembledSource += INJECT_SHADER_DECLARATIONS;
    assembledSource = injectShader(assembledSource, stage, declInjections);
    assembledSource += getShaderHooks(hookFunctionMap[stage], hookInjections);
    // Add the version directive and actual source of this shader
    assembledSource += coreSource;
    // Apply any requested shader injections
    assembledSource = injectShader(assembledSource, stage, mainInjections);
    if (language === 'glsl' && sourceVersion !== targetVersion) {
        assembledSource = transpileGLSLShader(assembledSource, stage);
    }
    return assembledSource.trim();
}
/**
 * Returns a combined `getUniforms` covering the options for all the modules,
 * the created function will pass on options to the inidividual `getUniforms`
 * function of each shader module and combine the results into one object that
 * can be passed to setUniforms.
 * @param modules
 * @returns
 */
export function assembleGetUniforms(modules) {
    return function getUniforms(opts) {
        const uniforms = {};
        for (const module of modules) {
            // `modules` is already sorted by dependency level. This guarantees that
            // modules have access to the uniforms that are generated by their dependencies.
            const moduleUniforms = module.getUniforms(opts, uniforms);
            Object.assign(uniforms, moduleUniforms);
        }
        return uniforms;
    };
}
/**
 * Generate "glslify-compatible" SHADER_NAME defines
 * These are understood by the GLSL error parsing function
 * If id is provided and no SHADER_NAME constant is present in source, create one
 */
function getShaderNameDefine(options) {
    const { id, source, stage } = options;
    const injectShaderName = id && source.indexOf('SHADER_NAME') === -1;
    return injectShaderName
        ? `
#define SHADER_NAME ${id}_${stage}

`
        : '';
}
/** Generates application defines from an object of key value pairs */
function getApplicationDefines(defines = {}) {
    let sourceText = '';
    for (const define in defines) {
        const value = defines[define];
        if (value || Number.isFinite(value)) {
            sourceText += `#define ${define.toUpperCase()} ${defines[define]}\n`;
        }
    }
    return sourceText;
}
/*
function getHookFunctions(
  hookFunctions: Record<string, HookFunction>,
  hookInjections: Record<string, Injection[]>
): string {
  let result = '';
  for (const hookName in hookFunctions) {
    const hookFunction = hookFunctions[hookName];
    result += `void ${hookFunction.signature} {\n`;
    if (hookFunction.header) {
      result += `  ${hookFunction.header}`;
    }
    if (hookInjections[hookName]) {
      const injections = hookInjections[hookName];
      injections.sort((a: {order: number}, b: {order: number}): number => a.order - b.order);
      for (const injection of injections) {
        result += `  ${injection.injection}\n`;
      }
    }
    if (hookFunction.footer) {
      result += `  ${hookFunction.footer}`;
    }
    result += '}\n';
  }

  return result;
}

function normalizeHookFunctions(hookFunctions: (string | HookFunction)[]): {
  vs: Record<string, HookFunction>;
  fs: Record<string, HookFunction>;
} {
  const result: {vs: Record<string, any>; fs: Record<string, any>} = {
    vs: {},
    fs: {}
  };

  hookFunctions.forEach((hookFunction: string | HookFunction) => {
    let opts: HookFunction;
    let hook: string;
    if (typeof hookFunction !== 'string') {
      opts = hookFunction;
      hook = opts.hook;
    } else {
      opts = {} as HookFunction;
      hook = hookFunction;
    }
    hook = hook.trim();
    const [stage, signature] = hook.split(':');
    const name = hook.replace(/\(.+/, '');
    if (stage !== 'vs' && stage !== 'fs') {
      throw new Error(stage);
    }
    result[stage][name] = Object.assign(opts, {signature});
  });

  return result;
}
*/