// Copyright 2021 The Tint Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "src/tint/transform/multiplanar_external_texture.h"

#include <string>
#include <vector>

#include "src/tint/ast/function.h"
#include "src/tint/program_builder.h"
#include "src/tint/sem/call.h"
#include "src/tint/sem/function.h"
#include "src/tint/sem/variable.h"

TINT_INSTANTIATE_TYPEINFO(tint::transform::MultiplanarExternalTexture);
TINT_INSTANTIATE_TYPEINFO(tint::transform::MultiplanarExternalTexture::NewBindingPoints);

namespace tint::transform {
namespace {

/// This struct stores symbols for new bindings created as a result of
/// transforming a texture_external instance.
struct NewBindingSymbols {
    Symbol params;
    Symbol plane_0;
    Symbol plane_1;
};
}  // namespace

/// State holds the current transform state
struct MultiplanarExternalTexture::State {
    /// The clone context.
    CloneContext& ctx;

    /// ProgramBuilder for the context
    ProgramBuilder& b;

    /// Destination binding locations for the expanded texture_external provided
    /// as input into the transform.
    const NewBindingPoints* new_binding_points;

    /// Symbol for the GammaTransferParams
    Symbol gamma_transfer_struct_sym;

    /// Symbol for the ExternalTextureParams struct
    Symbol params_struct_sym;

    /// Symbol for the textureLoadExternal function
    Symbol texture_load_external_sym;

    /// Symbol for the textureSampleExternal function
    Symbol texture_sample_external_sym;

    /// Symbol for the gammaCorrection function
    Symbol gamma_correction_sym;

    /// Storage for new bindings that have been created corresponding to an
    /// original texture_external binding.
    std::unordered_map<const sem::Variable*, NewBindingSymbols> new_binding_symbols;

    /// Constructor
    /// @param context the clone
    /// @param newBindingPoints the input destination binding locations for the
    /// expanded texture_external
    State(CloneContext& context, const NewBindingPoints* newBindingPoints)
        : ctx(context), b(*context.dst), new_binding_points(newBindingPoints) {}

    /// Processes the module
    void Process() {
        auto& sem = ctx.src->Sem();

        // For each texture_external binding, we replace it with a texture_2d<f32>
        // binding and create two additional bindings (one texture_2d<f32> to
        // represent the secondary plane and one uniform buffer for the
        // ExternalTextureParams struct).
        for (auto* var : ctx.src->AST().GlobalVariables()) {
            auto* sem_var = sem.Get(var);
            if (!sem_var->Type()->UnwrapRef()->Is<sem::ExternalTexture>()) {
                continue;
            }

            // If the attributes are empty, then this must be a texture_external
            // passed as a function parameter. These variables are transformed
            // elsewhere.
            if (var->attributes.empty()) {
                continue;
            }

            // If we find a texture_external binding, we know we must emit the
            // ExternalTextureParams struct.
            if (!params_struct_sym.IsValid()) {
                createExtTexParamsStructs();
            }

            // The binding points for the newly introduced bindings must have been
            // provided to this transform. We fetch the new binding points by
            // providing the original texture_external binding points into the
            // passed map.
            BindingPoint bp = {var->BindingPoint().group->value,
                               var->BindingPoint().binding->value};

            BindingsMap::const_iterator it = new_binding_points->bindings_map.find(bp);
            if (it == new_binding_points->bindings_map.end()) {
                b.Diagnostics().add_error(
                    diag::System::Transform,
                    "missing new binding points for texture_external at binding {" +
                        std::to_string(bp.group) + "," + std::to_string(bp.binding) + "}");
                continue;
            }

            BindingPoints bps = it->second;

            // Symbols for the newly created bindings must be saved so they can be
            // passed as parameters later. These are placed in a map and keyed by
            // the source symbol associated with the texture_external binding that
            // corresponds with the new destination bindings.
            // NewBindingSymbols new_binding_syms;
            auto& syms = new_binding_symbols[sem_var];
            syms.plane_0 = ctx.Clone(var->symbol);
            syms.plane_1 = b.Symbols().New("ext_tex_plane_1");
            b.Global(syms.plane_1, b.ty.sampled_texture(ast::TextureDimension::k2d, b.ty.f32()),
                     b.GroupAndBinding(bps.plane_1.group, bps.plane_1.binding));
            syms.params = b.Symbols().New("ext_tex_params");
            b.Global(syms.params, b.ty.type_name("ExternalTextureParams"),
                     ast::StorageClass::kUniform,
                     b.GroupAndBinding(bps.params.group, bps.params.binding));

            // Replace the original texture_external binding with a texture_2d<f32>
            // binding.
            ast::AttributeList cloned_attributes = ctx.Clone(var->attributes);
            const ast::Expression* cloned_constructor = ctx.Clone(var->constructor);

            auto* replacement =
                b.Var(syms.plane_0, b.ty.sampled_texture(ast::TextureDimension::k2d, b.ty.f32()),
                      cloned_constructor, cloned_attributes);
            ctx.Replace(var, replacement);
        }

        // We must update all the texture_external parameters for user declared
        // functions.
        for (auto* fn : ctx.src->AST().Functions()) {
            for (const ast::Variable* param : fn->params) {
                if (auto* sem_var = sem.Get(param)) {
                    if (!sem_var->Type()->UnwrapRef()->Is<sem::ExternalTexture>()) {
                        continue;
                    }
                    // If we find a texture_external, we must ensure the
                    // ExternalTextureParams struct exists.
                    if (!params_struct_sym.IsValid()) {
                        createExtTexParamsStructs();
                    }
                    // When a texture_external is found, we insert all components
                    // the texture_external into the parameter list. We must also place
                    // the new symbols into the transform state so they can be used when
                    // transforming function calls.
                    auto& syms = new_binding_symbols[sem_var];
                    syms.plane_0 = ctx.Clone(param->symbol);
                    syms.plane_1 = b.Symbols().New("ext_tex_plane_1");
                    syms.params = b.Symbols().New("ext_tex_params");
                    auto tex2d_f32 = [&] {
                        return b.ty.sampled_texture(ast::TextureDimension::k2d, b.ty.f32());
                    };
                    ctx.Replace(param, b.Param(syms.plane_0, tex2d_f32()));
                    ctx.InsertAfter(fn->params, param, b.Param(syms.plane_1, tex2d_f32()));
                    ctx.InsertAfter(fn->params, param,
                                    b.Param(syms.params, b.ty.type_name(params_struct_sym)));
                }
            }
        }

        // Transform the original textureLoad and textureSampleLevel calls into
        // textureLoadExternal and textureSampleExternal calls.
        ctx.ReplaceAll([&](const ast::CallExpression* expr) -> const ast::CallExpression* {
            auto* builtin = sem.Get(expr)->Target()->As<sem::Builtin>();

            if (builtin && !builtin->Parameters().empty() &&
                builtin->Parameters()[0]->Type()->Is<sem::ExternalTexture>() &&
                builtin->Type() != sem::BuiltinType::kTextureDimensions) {
                if (auto* var_user = sem.Get<sem::VariableUser>(expr->args[0])) {
                    auto it = new_binding_symbols.find(var_user->Variable());
                    if (it == new_binding_symbols.end()) {
                        // If valid new binding locations were not provided earlier, we
                        // would have been unable to create these symbols. An error
                        // message was emitted earlier, so just return early to avoid
                        // internal compiler errors and retain a clean error message.
                        return nullptr;
                    }
                    auto& syms = it->second;

                    if (builtin->Type() == sem::BuiltinType::kTextureLoad) {
                        return createTexLdExt(expr, syms);
                    }

                    if (builtin->Type() == sem::BuiltinType::kTextureSampleLevel) {
                        return createTexSmpExt(expr, syms);
                    }
                }

            } else if (sem.Get(expr)->Target()->Is<sem::Function>()) {
                // The call expression may be to a user-defined function that
                // contains a texture_external parameter. These need to be expanded
                // out to multiple plane textures and the texture parameters
                // structure.
                for (auto* arg : expr->args) {
                    if (auto* var_user = sem.Get<sem::VariableUser>(arg)) {
                        // Check if a parameter is a texture_external by trying to find
                        // it in the transform state.
                        auto it = new_binding_symbols.find(var_user->Variable());
                        if (it != new_binding_symbols.end()) {
                            auto& syms = it->second;
                            // When we find a texture_external, we must unpack it into its
                            // components.
                            ctx.Replace(arg, b.Expr(syms.plane_0));
                            ctx.InsertAfter(expr->args, arg, b.Expr(syms.plane_1));
                            ctx.InsertAfter(expr->args, arg, b.Expr(syms.params));
                        }
                    }
                }
            }

            return nullptr;
        });
    }

    /// Creates the parameter structs associated with the transform.
    void createExtTexParamsStructs() {
        // Create GammaTransferParams struct.
        ast::StructMemberList gamma_transfer_member_list = {
            b.Member("G", b.ty.f32()), b.Member("A", b.ty.f32()),      b.Member("B", b.ty.f32()),
            b.Member("C", b.ty.f32()), b.Member("D", b.ty.f32()),      b.Member("E", b.ty.f32()),
            b.Member("F", b.ty.f32()), b.Member("padding", b.ty.u32())};

        gamma_transfer_struct_sym = b.Symbols().New("GammaTransferParams");

        b.Structure(gamma_transfer_struct_sym, gamma_transfer_member_list);

        // Create ExternalTextureParams struct.
        ast::StructMemberList ext_tex_params_member_list = {
            b.Member("numPlanes", b.ty.u32()),
            b.Member("yuvToRgbConversionMatrix", b.ty.mat3x4(b.ty.f32())),
            b.Member("gammaDecodeParams", b.ty.type_name("GammaTransferParams")),
            b.Member("gammaEncodeParams", b.ty.type_name("GammaTransferParams")),
            b.Member("gamutConversionMatrix", b.ty.mat3x3(b.ty.f32()))};

        params_struct_sym = b.Symbols().New("ExternalTextureParams");

        b.Structure(params_struct_sym, ext_tex_params_member_list);
    }

    /// Creates the gammaCorrection function if needed and returns a call
    /// expression to it.
    void createGammaCorrectionFn() {
        using f32 = ProgramBuilder::f32;
        ast::VariableList varList = {b.Param("v", b.ty.vec3<f32>()),
                                     b.Param("params", b.ty.type_name(gamma_transfer_struct_sym))};

        ast::StatementList statementList = {
            // let cond = abs(v) < vec3(params.D);
            b.Decl(b.Let(
                "cond", nullptr,
                b.LessThan(b.Call("abs", "v"), b.vec3<f32>(b.MemberAccessor("params", "D"))))),
            // let t = sign(v) * ((params.C * abs(v)) + params.F);
            b.Decl(b.Let("t", nullptr,
                         b.Mul(b.Call("sign", "v"),
                               b.Add(b.Mul(b.MemberAccessor("params", "C"), b.Call("abs", "v")),
                                     b.MemberAccessor("params", "F"))))),
            // let f = (sign(v) * pow(((params.A * abs(v)) + params.B),
            // vec3(params.G))) + params.E;
            b.Decl(b.Let(
                "f", nullptr,
                b.Mul(b.Call("sign", "v"),
                      b.Add(b.Call("pow",
                                   b.Add(b.Mul(b.MemberAccessor("params", "A"), b.Call("abs", "v")),
                                         b.MemberAccessor("params", "B")),
                                   b.vec3<f32>(b.MemberAccessor("params", "G"))),
                            b.MemberAccessor("params", "E"))))),
            // return select(f, t, cond);
            b.Return(b.Call("select", "f", "t", "cond"))};

        gamma_correction_sym = b.Symbols().New("gammaCorrection");

        b.Func(gamma_correction_sym, varList, b.ty.vec3<f32>(), statementList, {});
    }

    /// Constructs a StatementList containing all the statements making up the
    /// bodies of the textureSampleExternal and textureLoadExternal functions.
    /// @param call_type determines which function body to generate
    /// @returns a statement list that makes of the body of the chosen function
    ast::StatementList createTexFnExtStatementList(sem::BuiltinType call_type) {
        using f32 = ProgramBuilder::f32;
        const ast::CallExpression* single_plane_call = nullptr;
        const ast::CallExpression* plane_0_call = nullptr;
        const ast::CallExpression* plane_1_call = nullptr;
        if (call_type == sem::BuiltinType::kTextureSampleLevel) {
            // textureSampleLevel(plane0, smp, coord.xy, 0.0);
            single_plane_call = b.Call("textureSampleLevel", "plane0", "smp", "coord", 0.0f);
            // textureSampleLevel(plane0, smp, coord.xy, 0.0);
            plane_0_call = b.Call("textureSampleLevel", "plane0", "smp", "coord", 0.0f);
            // textureSampleLevel(plane1, smp, coord.xy, 0.0);
            plane_1_call = b.Call("textureSampleLevel", "plane1", "smp", "coord", 0.0f);
        } else if (call_type == sem::BuiltinType::kTextureLoad) {
            // textureLoad(plane0, coords.xy, 0);
            single_plane_call = b.Call("textureLoad", "plane0", "coord", 0);
            // textureLoad(plane0, coords.xy, 0);
            plane_0_call = b.Call("textureLoad", "plane0", "coord", 0);
            // textureLoad(plane1, coords.xy, 0);
            plane_1_call = b.Call("textureLoad", "plane1", "coord", 0);
        } else {
            TINT_ICE(Transform, b.Diagnostics()) << "unhandled builtin: " << call_type;
        }

        return {
            // var color: vec3<f32>;
            b.Decl(b.Var("color", b.ty.vec3(b.ty.f32()))),
            // if ((params.numPlanes == 1u))
            b.If(b.create<ast::BinaryExpression>(
                     ast::BinaryOp::kEqual, b.MemberAccessor("params", "numPlanes"), b.Expr(1u)),
                 b.Block(
                     // color = textureLoad(plane0, coord, 0).rgb;
                     b.Assign("color", b.MemberAccessor(single_plane_call, "rgb"))),
                 b.Block(
                     // color = vec4<f32>(plane_0_call.r, plane_1_call.rg, 1.0) *
                     //         params.yuvToRgbConversionMatrix;
                     b.Assign("color",
                              b.Mul(b.vec4<f32>(b.MemberAccessor(plane_0_call, "r"),
                                                b.MemberAccessor(plane_1_call, "rg"), 1.0f),
                                    b.MemberAccessor("params", "yuvToRgbConversionMatrix"))))),
            // color = gammaConversion(color, gammaDecodeParams);
            b.Assign("color", b.Call("gammaCorrection", "color",
                                     b.MemberAccessor("params", "gammaDecodeParams"))),
            // color = (params.gamutConversionMatrix * color);
            b.Assign("color", b.Mul(b.MemberAccessor("params", "gamutConversionMatrix"), "color")),
            // color = gammaConversion(color, gammaEncodeParams);
            b.Assign("color", b.Call("gammaCorrection", "color",
                                     b.MemberAccessor("params", "gammaEncodeParams"))),
            // return vec4<f32>(color, 1.0f);
            b.Return(b.vec4<f32>("color", 1.0f))};
    }

    /// Creates the textureSampleExternal function if needed and returns a call
    /// expression to it.
    /// @param expr the call expression being transformed
    /// @param syms the expanded symbols to be used in the new call
    /// @returns a call expression to textureSampleExternal
    const ast::CallExpression* createTexSmpExt(const ast::CallExpression* expr,
                                               NewBindingSymbols syms) {
        ast::ExpressionList params;
        const ast::Expression* plane_0_binding_param = ctx.Clone(expr->args[0]);

        if (expr->args.size() != 3) {
            TINT_ICE(Transform, b.Diagnostics()) << "expected textureSampleLevel call with a "
                                                    "texture_external to have 3 parameters, found "
                                                 << expr->args.size() << " parameters";
        }

        // TextureSampleExternal calls the gammaCorrection function, so ensure it
        // exists.
        if (!gamma_correction_sym.IsValid()) {
            createGammaCorrectionFn();
        }

        if (!texture_sample_external_sym.IsValid()) {
            texture_sample_external_sym = b.Symbols().New("textureSampleExternal");

            // Emit the textureSampleExternal function.
            ast::VariableList varList = {
                b.Param("plane0", b.ty.sampled_texture(ast::TextureDimension::k2d, b.ty.f32())),
                b.Param("plane1", b.ty.sampled_texture(ast::TextureDimension::k2d, b.ty.f32())),
                b.Param("smp", b.ty.sampler(ast::SamplerKind::kSampler)),
                b.Param("coord", b.ty.vec2(b.ty.f32())),
                b.Param("params", b.ty.type_name(params_struct_sym))};

            ast::StatementList statementList =
                createTexFnExtStatementList(sem::BuiltinType::kTextureSampleLevel);

            b.Func(texture_sample_external_sym, varList, b.ty.vec4(b.ty.f32()), statementList, {});
        }

        const ast::IdentifierExpression* exp = b.Expr(texture_sample_external_sym);
        params = {plane_0_binding_param, b.Expr(syms.plane_1), ctx.Clone(expr->args[1]),
                  ctx.Clone(expr->args[2]), b.Expr(syms.params)};
        return b.Call(exp, params);
    }

    /// Creates the textureLoadExternal function if needed and returns a call
    /// expression to it.
    /// @param expr the call expression being transformed
    /// @param syms the expanded symbols to be used in the new call
    /// @returns a call expression to textureLoadExternal
    const ast::CallExpression* createTexLdExt(const ast::CallExpression* expr,
                                              NewBindingSymbols syms) {
        ast::ExpressionList params;
        const ast::Expression* plane_0_binding_param = ctx.Clone(expr->args[0]);

        if (expr->args.size() != 2) {
            TINT_ICE(Transform, b.Diagnostics())
                << "expected textureLoad call with a texture_external "
                   "to have 2 parameters, found "
                << expr->args.size() << " parameters";
        }

        // TextureLoadExternal calls the gammaCorrection function, so ensure it
        // exists.
        if (!gamma_correction_sym.IsValid()) {
            createGammaCorrectionFn();
        }

        if (!texture_load_external_sym.IsValid()) {
            texture_load_external_sym = b.Symbols().New("textureLoadExternal");

            // Emit the textureLoadExternal function.
            ast::VariableList var_list = {
                b.Param("plane0", b.ty.sampled_texture(ast::TextureDimension::k2d, b.ty.f32())),
                b.Param("plane1", b.ty.sampled_texture(ast::TextureDimension::k2d, b.ty.f32())),
                b.Param("coord", b.ty.vec2(b.ty.i32())),
                b.Param("params", b.ty.type_name(params_struct_sym))};

            ast::StatementList statement_list =
                createTexFnExtStatementList(sem::BuiltinType::kTextureLoad);

            b.Func(texture_load_external_sym, var_list, b.ty.vec4(b.ty.f32()), statement_list, {});
        }

        const ast::IdentifierExpression* exp = b.Expr(texture_load_external_sym);
        params = {plane_0_binding_param, b.Expr(syms.plane_1), ctx.Clone(expr->args[1]),
                  b.Expr(syms.params)};
        return b.Call(exp, params);
    }
};

MultiplanarExternalTexture::NewBindingPoints::NewBindingPoints(BindingsMap inputBindingsMap)
    : bindings_map(std::move(inputBindingsMap)) {}
MultiplanarExternalTexture::NewBindingPoints::~NewBindingPoints() = default;

MultiplanarExternalTexture::MultiplanarExternalTexture() = default;
MultiplanarExternalTexture::~MultiplanarExternalTexture() = default;

bool MultiplanarExternalTexture::ShouldRun(const Program* program, const DataMap&) const {
    for (auto* node : program->ASTNodes().Objects()) {
        if (auto* ty = node->As<ast::Type>()) {
            if (program->Sem().Get<sem::ExternalTexture>(ty)) {
                return true;
            }
        }
    }
    return false;
}

// Within this transform, an instance of a texture_external binding is unpacked
// into two texture_2d<f32> bindings representing two possible planes of a
// single texture and a uniform buffer binding representing a struct of
// parameters. Calls to textureLoad or textureSampleLevel that contain a
// texture_external parameter will be transformed into a newly generated version
// of the function, which can perform the desired operation on a single RGBA
// plane or on separate Y and UV planes.
void MultiplanarExternalTexture::Run(CloneContext& ctx, const DataMap& inputs, DataMap&) const {
    auto* new_binding_points = inputs.Get<NewBindingPoints>();

    if (!new_binding_points) {
        ctx.dst->Diagnostics().add_error(
            diag::System::Transform,
            "missing new binding point data for " + std::string(TypeInfo().name));
        return;
    }

    State state(ctx, new_binding_points);

    state.Process();

    ctx.Clone();
}

}  // namespace tint::transform