/**
 * Copyright © 2012-2025 by Data Virtuality GmbH
 */
package org.teiid.query.resolver.util;

import org.teiid.core.types.DataTypeManager;
import org.teiid.core.types.NullType;
import org.teiid.metadata.Column;
import org.teiid.metadata.Procedure;
import org.teiid.query.metadata.QueryMetadataInterface;
import org.teiid.query.sql.lang.Command;
import org.teiid.query.sql.lang.SetQuery;
import org.teiid.query.sql.proc.Block;
import org.teiid.query.sql.proc.CommandStatement;
import org.teiid.query.sql.proc.CreateProcedureCommand;
import org.teiid.query.sql.proc.Statement;
import org.teiid.query.sql.symbol.Expression;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * Utilities to extract expected types from a procedure RETURNS clause
 * and apply them to returnable statements (e.g., SetQuery). This ensures
 * NULL-only queries (like "SELECT NULL UNION SELECT NULL") are typed to the
 * declared RETURNS.
 *
 * @author Alexander Fischer
 * @date 2025.10.15
 */
public final class ProcedureReturnsTypeApplier {

    private ProcedureReturnsTypeApplier() {}

    /** Build expected Java classes from procedure RETURNS. */
    public static List<Class<?>> expectedFromReturns(Procedure procMeta) {
        if (procMeta == null || procMeta.getResultSet() == null || procMeta.getResultSet().getColumns() == null) {
            return Collections.emptyList();
        }
        List<Class<?>> expected = new ArrayList<>(procMeta.getResultSet().getColumns().size());
        for (Column c : procMeta.getResultSet().getColumns()) {
            String rt = c.getRuntimeType();
            if (rt == null && c.getDatatype() != null) {
                rt = c.getDatatype().getRuntimeTypeName();
            }
            Class<?> clazz = (rt != null) ? DataTypeManager.getDataTypeClass(rt) : NullType.class;
            if (clazz == null) {
                clazz = NullType.class;
            }
            expected.add(clazz);
        }
        return expected;
    }

    /** Build expected Java classes from already resolved expressions (e.g., command.getResultSetColumns()). */
    public static List<Class<?>> expectedFromExpressions(List<? extends Expression> cols) {
        if (cols == null || cols.isEmpty()) {
            return Collections.emptyList();
        }
        List<Class<?>> expected = new ArrayList<>(cols.size());
        for (Expression e : cols) {
            Class<?> t = (e != null) ? e.getType() : null;
            expected.add(t != null ? t : NullType.class);
        }
        return expected;
    }

    /**
     * Apply expected types to the first returnable statement inside a procedure block
     * (typically a SetQuery), if present.
     */
    public static void applyToFirstReturnable(CreateProcedureCommand cpc,
                                              List<Class<?>> expected,
                                              QueryMetadataInterface qmi) {
        if (cpc == null || cpc.getBlock() == null || expected == null || expected.isEmpty()) return;

        Block block = cpc.getBlock();
        List<Statement> stmts = block.getStatements();
        if (stmts == null) return;

        for (Statement st : stmts) {
            if (st instanceof CommandStatement) {
                CommandStatement cs = (CommandStatement) st;
                if (!cs.isReturnable()) {
                    continue;
                }
                Command inner = cs.getCommand();
                if (inner instanceof SetQuery) {
                    ((SetQuery) inner).setProjectedTypes(expected, qmi.getDesignTimeMetadata());
                }
                break;
            }
        }
    }

    /**
     * Apply expected types to a command if it is a SetQuery (e.g., inside UpdateProcedureResolver
     * for a returnable CommandStatement).
     */
    public static void applyToCommandIfSetQuery(Command cmd,
                                                List<Class<?>> expected,
                                                QueryMetadataInterface qmi) {
        if (cmd instanceof SetQuery && expected != null && !expected.isEmpty()) {
            ((SetQuery) cmd).setProjectedTypes(expected, qmi);
        }
    }
}
