type ShaderType =
    | WebGLRenderingContext["VERTEX_SHADER"]
    | WebGLRenderingContext["FRAGMENT_SHADER"];

interface ShaderSource {
    source: string;
    kind: ShaderType;
}

type ShaderAttributes = Map<string, number>;
type ShaderUniforms = Map<string, WebGLUniformLocation>;

class ShaderError extends Error {}

class Shader {
    private program_: WebGLProgram;
    private attributes_: ShaderAttributes;
    private uniforms_: ShaderUniforms;

    constructor(
        program: WebGLProgram,
        attributes: ShaderAttributes,
        uniforms: ShaderUniforms
    ) {
        this.program_ = program;
        this.attributes_ = attributes;
        this.uniforms_ = uniforms;
    }

    static builder(gl: WebGLRenderingContext): ShaderBuilder {
        return new ShaderBuilder(gl);
    }

    get program(): WebGLProgram {
        return this.program_;
    }

    public getAttribute(name: string): number {
        const attribute = this.attributes_.get(name);

        if (attribute === undefined) {
            throw new ShaderError(`undefined shader attribute: ${name}`);
        }

        return attribute;
    }

    public getUniform(name: string): WebGLUniformLocation {
        const uniform = this.uniforms_.get(name);

        if (uniform === undefined) {
            throw new ShaderError(`undefined shader uniform: ${name}`);
        }

        return uniform;
    }

    get uniforms(): ShaderUniforms {
        return this.uniforms_;
    }
}

class ShaderBuilder {
    private gl: WebGLRenderingContext;
    private sources: Array<ShaderSource>;
    private attributes: Array<string>;
    private uniforms: Array<string>;

    public constructor(gl: WebGLRenderingContext) {
        this.gl = gl;
        this.sources = new Array<ShaderSource>();
        this.attributes = new Array<string>();
        this.uniforms = new Array<string>();
    }

    public addShader(source: string, kind: ShaderType): ShaderBuilder {
        this.sources.push({ source, kind });
        return this;
    }

    public addAttribute(name: string): ShaderBuilder {
        this.attributes.push(name);
        return this;
    }

    public addUniforms(name: string): ShaderBuilder {
        this.uniforms.push(name);
        return this;
    }

    public build(): Shader {
        // Load, compile and link shader sources
        const shaders = this.sources.map(({ source, kind }) => {
            return this.loadShader(source, kind);
        });

        const shaderProgram = this.gl.createProgram();
        if (shaderProgram === null) {
            throw new ShaderError("failed to create shader program");
        }

        for (const shader of shaders) {
            this.gl.attachShader(shaderProgram, shader);
        }

        this.gl.linkProgram(shaderProgram);
        if (!this.gl.getProgramParameter(shaderProgram, this.gl.LINK_STATUS)) {
            let message = "failed to link shader program";
            const log = this.gl.getProgramInfoLog(shaderProgram);
            if (log !== null) {
                message = `failed to link shader program: ${log}`;
            }

            throw new ShaderError(message);
        }

        // Find attribute and uniform locations
        const attributes = this.attributes.reduce((acc, attribute) => {
            const attributeLocation = this.gl.getAttribLocation(
                shaderProgram,
                attribute
            );

            if (attributeLocation === -1) {
                throw new ShaderError(
                    `shader attribute '${attribute}' could not be found`
                );
            }

            return new Map<string, number>([
                ...acc,
                [attribute, attributeLocation],
            ]);
        }, new Map<string, number>());

        const uniforms = this.uniforms.reduce((acc, uniform) => {
            const uniformLocation = this.gl.getUniformLocation(
                shaderProgram,
                uniform
            );

            if (uniformLocation === null) {
                throw new ShaderError(
                    `shader uniform '${uniform}' could not be found`
                );
            }

            return new Map<string, WebGLUniformLocation>([
                ...acc,
                [uniform, uniformLocation],
            ]);
        }, new Map<string, WebGLUniformLocation>());

        // Build actual shader object
        return new Shader(shaderProgram, attributes, uniforms);
    }

    private loadShader(source: string, kind: ShaderType): WebGLShader {
        const shader = this.gl.createShader(kind);
        if (shader === null) {
            throw new ShaderError(`failed to initialize shader "${source}"`);
        }

        this.gl.shaderSource(shader, source);
        this.gl.compileShader(shader);

        if (!this.gl.getShaderParameter(shader, this.gl.COMPILE_STATUS)) {
            let message = `failed to compile shader "${source}"`;
            const log = this.gl.getShaderInfoLog(shader);
            if (log !== null) {
                message = `failed to compile shader "${source}": ${log}`;
            }

            this.gl.deleteShader(shader);

            throw new ShaderError(message);
        }

        return shader;
    }
}

export { Shader, ShaderError };
