package io.nosqlbench.nb.api.expr.providers;

/*
 * Copyright (c) nosqlbench
 *
 * 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.
 */


import io.nosqlbench.nb.annotations.Service;
import io.nosqlbench.nb.api.system.NBEnvironment;
import io.nosqlbench.nb.api.expr.ExprFunctionProvider;
import io.nosqlbench.nb.api.expr.ExprRuntimeContext;
import io.nosqlbench.nb.api.expr.TemplateContext;
import io.nosqlbench.nb.api.expr.annotations.ExprExample;
import io.nosqlbench.nb.api.expr.annotations.ExprFunctionSpec;

import java.util.Map;

/**
 * Unified parameter provider that consolidates all parameter and template variable access functions.
 *
 * <p>This provider offers two complementary APIs:</p>
 * <ul>
 *   <li><b>User-facing API:</b> {@code param()}, {@code paramOr()}, {@code hasParam()} - clean, simple parameter access</li>
 *   <li><b>Internal API:</b> {@code _template*()} functions - generated by TemplateRewriter for legacy template syntax support</li>
 * </ul>
 *
 * <h3>Unified Resolution Strategy</h3>
 * <p>All functions use consistent parameter resolution order:</p>
 * <ol>
 *   <li>TemplateContext overrides (if active)</li>
 *   <li>Explicit parameters</li>
 *   <li>Environment variables (for template functions)</li>
 *   <li>Default value (if provided)</li>
 * </ol>
 */
@Service(value = ExprFunctionProvider.class, selector = "unified-parameters")
public class UnifiedParameterProvider implements ExprFunctionProvider {

    /**
     * Resolve a parameter value using the unified resolution strategy.
     *
     * @param context the expression runtime context
     * @param name the parameter name
     * @param checkEnvironment whether to check environment variables
     * @return the resolved value, or null if not found
     */
    private Object resolveParameter(ExprRuntimeContext context, String name, boolean checkEnvironment) {
        Map<String, ?> params = context.parameters();

        // Check TemplateContext overrides first (if context is active)
        if (TemplateContext.isActive()) {
            TemplateContext ctx = TemplateContext.current();
            if (ctx.hasOverride(name)) {
                return ctx.getOverride(name);
            }
        }

        // Check explicit parameters
        if (params.containsKey(name)) {
            return params.get(name);
        }

        // Check environment variables (only for template functions)
        if (checkEnvironment && NBEnvironment.INSTANCE.containsKey(name)) {
            return NBEnvironment.INSTANCE.get(name);
        }

        return null;
    }

    /**
     * Track a parameter access in the TemplateContext (if active).
     *
     * @param name the parameter name
     * @param value the resolved value
     */
    private void trackAccess(String name, String value) {
        if (TemplateContext.isActive()) {
            TemplateContext.current().trackAccess(name, value);
        }
    }

    // ============================================================
    // User-Facing Parameter API
    // ============================================================

    @ExprExample(args = {"\"mode\""}, expect = "\"strict\"")
    @ExprExample(args = {"\"threshold\""}, expect = "42")
    @ExprFunctionSpec(
        name = "param",
        synopsis = "param(name)",
        description = "Return the value of a required workload parameter, raising an error when missing."
    )
    private Object param(ExprRuntimeContext context, Object[] args) {
        if (args.length != 1) {
            throw new IllegalArgumentException("param(name) requires exactly one argument");
        }
        String name = String.valueOf(args[0]);
        Object value = resolveParameter(context, name, false);
        if (value == null) {
            throw new IllegalArgumentException("Parameter '" + name + "' was not provided");
        }
        trackAccess(name, String.valueOf(value));
        return value;
    }

    @ExprExample(args = {"\"threshold\"", "\"fallback\""}, expect = "42")
    @ExprExample(args = {"\"missing\"", "\"fallback\""}, expect = "\"fallback\"")
    @ExprFunctionSpec(
        name = "paramOr",
        synopsis = "paramOr(name, default)",
        description = "Return the value of a workload parameter or the provided default when absent or null. " +
            "Supports lexical scoping: once a parameter is resolved (from params or default), subsequent " +
            "calls without a default will reuse that value."
    )
    private Object paramOr(ExprRuntimeContext context, Object[] args) {
        if (args.length < 1 || args.length > 2) {
            throw new IllegalArgumentException("paramOr(name, default?) expects one or two arguments");
        }
        String name = String.valueOf(args[0]);

        // Lexical scoping: check if this parameter was already resolved in the workload
        String bindingKey = "__param_" + name;
        if (context.hasVariable(bindingKey)) {
            Object cached = context.getVariable(bindingKey);
            if (cached != null) {
                trackAccess(name, String.valueOf(cached));
            }
            return cached;
        }

        // Not yet resolved - determine the value
        Object value = resolveParameter(context, name, false);
        Object result;
        if (value == null && args.length == 2) {
            result = args[1];
        } else {
            result = value;
        }

        // Store in binding for lexical scoping
        context.setVariable(bindingKey, result);

        // Track template variable access for workload validation
        if (result != null) {
            trackAccess(name, String.valueOf(result));
        }

        return result;
    }

    @ExprExample(args = {"\"mode\""}, expect = "true")
    @ExprExample(args = {"\"missing\""}, expect = "false")
    @ExprFunctionSpec(
        name = "hasParam",
        synopsis = "hasParam(name)",
        description = "Check if a workload parameter was provided."
    )
    private Object hasParam(ExprRuntimeContext context, Object[] args) {
        if (args.length != 1) {
            throw new IllegalArgumentException("hasParam(name) requires exactly one argument");
        }
        String name = String.valueOf(args[0]);
        return resolveParameter(context, name, false) != null;
    }

    // ============================================================
    // Internal Template API (generated by TemplateRewriter)
    // ============================================================

    @ExprExample(args = {"\"retries\"", "\"3\""}, expect = "\"3\"")
    @ExprFunctionSpec(
        name = "_templateSet",
        synopsis = "_templateSet(name, default)",
        description = "Set parameter if unset, then return it (implements the := operator). " +
            "This mimics the <<key:=default>> behavior where the default is set and used on first access."
    )
    private Object _templateSet(ExprRuntimeContext context, Object[] args) {
        if (args.length != 2) {
            throw new IllegalArgumentException("_templateSet(name, default) requires exactly two arguments");
        }

        String name = String.valueOf(args[0]);
        String defaultValue = String.valueOf(args[1]);
        TemplateContext ctx = TemplateContext.current();

        // Check if already resolved (override → params → environment)
        Object resolved = resolveParameter(context, name, true);
        if (resolved != null) {
            String value = String.valueOf(resolved);
            ctx.trackAccess(name, value);
            return value;
        }

        // Set override and use default
        ctx.setOverride(name, defaultValue);
        ctx.trackAccess(name, defaultValue);
        return defaultValue;
    }

    @ExprExample(args = {"\"debug\"", "\"verbose\""}, expect = "\"verbose\"")
    @ExprExample(args = {"\"missing\"", "\"alternate\""}, expect = "\"\"")
    @ExprFunctionSpec(
        name = "_templateAlt",
        synopsis = "_templateAlt(name, alternate)",
        description = "Return alternate value if parameter IS set, empty string otherwise " +
            "(implements the :+ operator). This mimics the <<key:+alternate>> behavior."
    )
    private Object _templateAlt(ExprRuntimeContext context, Object[] args) {
        if (args.length != 2) {
            throw new IllegalArgumentException("_templateAlt(name, alternate) requires exactly two arguments");
        }

        String name = String.valueOf(args[0]);
        String alternate = String.valueOf(args[1]);
        TemplateContext ctx = TemplateContext.current();

        // Check if parameter is set (in any location)
        boolean isSet = resolveParameter(context, name, true) != null;

        if (isSet) {
            ctx.trackAccess(name, alternate);
            return alternate;
        }

        return "";
    }

    @ExprFunctionSpec(
        name = "_templateTrack",
        synopsis = "_templateTrack(name, value)",
        description = "Track template variable access for validation. This is used internally " +
            "to record which template variables were accessed during workload processing."
    )
    private Object _templateTrack(ExprRuntimeContext context, Object[] args) {
        if (args.length != 2) {
            throw new IllegalArgumentException("_templateTrack(name, value) requires exactly two arguments");
        }

        String name = String.valueOf(args[0]);
        String value = String.valueOf(args[1]);

        TemplateContext.current().trackAccess(name, value);
        return value;
    }

    @ExprExample(args = {"\"myvar\"", "\"default\""}, expect = "\"default\"")
    @ExprFunctionSpec(
        name = "_templateGet",
        synopsis = "_templateGet(name, default)",
        description = "Get a template variable value if already set, otherwise use default. " +
            "This implements lexical scoping where template variables set in one place are " +
            "accessible in later references without re-specifying the default."
    )
    private Object _templateGet(ExprRuntimeContext context, Object[] args) {
        if (args.length != 2) {
            throw new IllegalArgumentException("_templateGet(name, default) requires exactly two arguments");
        }

        String name = String.valueOf(args[0]);
        String defaultValue = String.valueOf(args[1]);
        TemplateContext ctx = TemplateContext.current();

        // Check if already set via _templateSet (override → params → environment)
        Object resolved = resolveParameter(context, name, true);
        if (resolved != null) {
            String value = String.valueOf(resolved);
            ctx.trackAccess(name, value);
            return value;
        }

        // Use default without setting
        ctx.trackAccess(name, defaultValue);
        return defaultValue;
    }
}
