package org.dbsp.sqlCompiler.compiler.frontend.calciteCompiler;

import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.sql.SqlFunction;
import org.apache.calcite.sql.SqlFunctionCategory;
import org.apache.calcite.sql.SqlIdentifier;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.SqlOperandCountRange;
import org.apache.calcite.sql.type.InferTypes;
import org.apache.calcite.sql.type.OperandTypes;
import org.apache.calcite.sql.type.ReturnTypes;
import org.apache.calcite.sql.type.SqlOperandCountRanges;
import org.apache.calcite.sql.type.SqlOperandTypeChecker;
import org.apache.calcite.sql.type.SqlOperandTypeInference;
import org.apache.calcite.sql.type.SqlSingleOperandTypeChecker;
import org.apache.calcite.sql.type.SqlTypeName;
import org.dbsp.sqlCompiler.compiler.DBSPCompiler;
import org.dbsp.sqlCompiler.compiler.errors.CompilationError;
import org.dbsp.sqlCompiler.compiler.errors.InternalCompilerError;
import org.dbsp.sqlCompiler.compiler.errors.SourcePositionRange;
import org.dbsp.sqlCompiler.compiler.frontend.ExpressionCompiler;
import org.dbsp.sqlCompiler.compiler.frontend.TypeCompiler;
import org.dbsp.sqlCompiler.compiler.frontend.calciteObject.CalciteObject;
import org.dbsp.sqlCompiler.ir.DBSPFunction;
import org.dbsp.sqlCompiler.ir.DBSPParameter;
import org.dbsp.sqlCompiler.ir.expression.DBSPApplyExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPApplyMethodExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPBlockExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPCastExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPVariablePath;
import org.dbsp.sqlCompiler.ir.statement.DBSPLetStatement;
import org.dbsp.sqlCompiler.ir.statement.DBSPStatement;
import org.dbsp.sqlCompiler.ir.type.DBSPType;
import org.dbsp.sqlCompiler.ir.type.primitive.DBSPTypeAny;
import org.dbsp.sqlCompiler.ir.type.primitive.DBSPTypeString;
import org.dbsp.sqlCompiler.ir.type.user.DBSPTypeResult;
import org.dbsp.util.Linq;
import org.dbsp.util.Utilities;

import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.List;

/** Represents a function that is implemented by the user in Rust */
public class ExternalFunction extends SqlFunction {
    public final CalciteObject node;
    public final RelDataType returnType;
    public final List<RelDataTypeField> parameterList;
    /** True if the function is generated by the compiler */
    public final boolean generated;
    public final SourcePositionRange position;
    /** Cached return type */
    @Nullable
    public DBSPType structReturnType = null;
    @Nullable
    public final RexNode body;

    static final String JSONSTRING_AS = "jsonstring_as_";

    public static SqlOperandTypeChecker createTypeChecker(String function, List<RelDataTypeField> parameters) {
        if (parameters.isEmpty())
            return OperandTypes.NILADIC;
        SqlSingleOperandTypeChecker[] checkers = new SqlSingleOperandTypeChecker[parameters.size()];
        StringBuilder builder = new StringBuilder();
        builder.append(function).append("(");
        for (int i = 0; i < parameters.size(); i++) {
            RelDataTypeField field = parameters.get(i);
            SqlTypeName typeName = field.getType().getSqlTypeName();
            // This type checker is a bit too strict.
            checkers[i] = new ExactTypeNameChecker(typeName);
            // This type checker allows implicit casts from any type in the type family,
            // which is not great.
            // checkers[i] = OperandTypes.typeName(typeName);
            if (i > 0)
                builder.append(", ");
            builder.append("<")
                    .append(field.getType().toString())
                    .append(">");
        }
        builder.append(")");
        if (checkers.length > 1)
            return OperandTypes.sequence(builder.toString(), checkers);
        return checkers[0];
    }

    /** True if this function's body is provided in SQL */
    public boolean isSqlFunction() {
        return this.body != null;
    }

    static SqlOperandTypeInference createTypeInference(List<RelDataTypeField> parameters) {
        return InferTypes.explicit(Linq.map(parameters, RelDataTypeField::getType));
    }

    public ExternalFunction(SqlIdentifier name, RelDataType returnType,
                            List<RelDataTypeField> parameters, @Nullable RexNode body, boolean generated) {
        super(name.getSimple(), SqlKind.OTHER_FUNCTION, ReturnTypes.explicit(returnType),
                createTypeInference(parameters), createTypeChecker(name.getSimple(), parameters),
                SqlFunctionCategory.USER_DEFINED_FUNCTION);
        this.node = CalciteObject.create(name);
        this.position = new SourcePositionRange(name.getParserPosition());
        this.body = body;
        this.returnType = returnType;
        this.parameterList = parameters;
        this.generated = generated;
    }

    @Override
    public boolean isDeterministic() {
        return false;
    }

    @Override
    public SqlOperandCountRange getOperandCountRange() {
        return SqlOperandCountRanges.of(this.parameterList.size());
    }

    /** Return type as declared in the SQL program, not as implemented.
     * I.e., may be a struct. */
    public DBSPType getFunctionReturnType(TypeCompiler compiler) {
        if (this.structReturnType == null)
            this.structReturnType = compiler.convertType(this.position, this.returnType, true);
        return this.structReturnType;
    }

    /** Visitor used to generate the function body from a RexNode */
    static class FunctionBodyGenerator extends ExpressionCompiler {
        final List<DBSPParameter> parameters;

        /* Customize the input ref to refer to parameters instead */
        @Override
        public DBSPExpression visitInputRef(RexInputRef inputRef) {
            CalciteObject node = CalciteObject.create(null, inputRef);
            int index = inputRef.getIndex();
            if (index < this.parameters.size()) {
                return this.parameters.get(index).asVariable().applyCloneIfNeeded();
            }
            throw new InternalCompilerError("Parameter index out of bounds ", node);
        }

        public FunctionBodyGenerator(DBSPCompiler compiler, List<DBSPParameter> parameters) {
            super(null, null, compiler);
            this.parameters = parameters;
        }
    }

    /** Get the declaration of an external function */
    public DBSPFunction getDeclaration(TypeCompiler typeCompiler) {
        List<DBSPParameter> parameters = new ArrayList<>();
        for (RelDataTypeField param: this.parameterList) {
            DBSPType type = typeCompiler.convertType(this.position, param.getType(), false);
            DBSPParameter p = new DBSPParameter(param.getName(), type);
            parameters.add(p);
        }
        DBSPType returnType = typeCompiler.convertType(SourcePositionRange.INVALID, this.returnType, false);
        // Must wrap the return type in a Result
        returnType = new DBSPTypeResult(returnType);
        return new DBSPFunction(this.node, this.getName(), parameters, returnType, null, Linq.list());
    }

    @Nullable
    public DBSPFunction getImplementation(TypeCompiler typeCompiler, DBSPCompiler compiler) {
        if (this.body != null) {
            List<DBSPParameter> parameters = new ArrayList<>();
            for (RelDataTypeField param: this.parameterList) {
                DBSPType type = typeCompiler.convertType(this.position, param.getType(), false);
                DBSPParameter p = new DBSPParameter(param.getName(), type);
                parameters.add(p);
            }
            FunctionBodyGenerator generator = new FunctionBodyGenerator(compiler, parameters);
            DBSPType returnType = typeCompiler.convertType(this.position, this.returnType, false);
            DBSPExpression functionBody = generator.compile(this.body);

            DBSPType functionType = functionBody.getType();
            if (!returnType.sameType(functionType)) {
                if (returnType.is(DBSPTypeString.class)) {
                    functionBody = functionBody.cast(
                            functionBody.getNode(), returnType, DBSPCastExpression.CastType.SqlUnsafe);
                } else if (returnType.sameType(functionType.withMayBeNull(true))) {
                    functionBody = functionBody.cast(
                            functionBody.getNode(), returnType, DBSPCastExpression.CastType.SqlUnsafe);
                } else {
                    throw new CompilationError("Function " + Utilities.singleQuote(this.getName()) +
                            " should return " + Utilities.singleQuote(this.returnType.getFullTypeString()) +
                            " but returns " + Utilities.singleQuote(this.body.getType().getFullTypeString()), this.node);
                }
            }
            return new DBSPFunction(this.node, this.getName(), parameters, returnType, functionBody, Linq.list());
        }
        if (!this.generated)
            return null;
        // Some classes of functions are automatically generated by the compiler.
        if (this.getName().toLowerCase().startsWith(JSONSTRING_AS)) {
            /*
            TYPE ~ struct
            pub fn JSONSTRING_AS_<TYPE>(_pos: &SourcePositionRange, s: Option<SqlString>) -> Option<Tup1<TYPE>> {
                let s = s?.str();
                let strct: Option<struct_0> = from_json_string(s);
                strct.map(move |x: _, | -> _ { x.into() })
            }
             */
            DBSPType returnType = typeCompiler.convertType(this.position, this.returnType, false);
            DBSPType structType = this.getFunctionReturnType(typeCompiler);
            if (this.parameterList.size() != 1) {
                compiler.reportError(this.position, "Incorrect signature",
                        "Function " + Utilities.singleQuote(this.getName()) +
                                " must have a single parameter");
                return null;
            }
            DBSPType parameterType = typeCompiler.convertType(this.position, this.parameterList.get(0).getType(), false);
            DBSPParameter param = new DBSPParameter("s", parameterType);
            List<DBSPStatement> statements = new ArrayList<>();
            if (parameterType.mayBeNull)
                statements.add(new DBSPLetStatement("s", param.asVariable().question()));
            DBSPExpression toStruct = new DBSPApplyExpression("from_json_string",
                    structType.withMayBeNull(true),
                    new DBSPApplyMethodExpression("str",
                            DBSPTypeAny.getDefault(), param.asVariable()));
            DBSPLetStatement strct = new DBSPLetStatement("strct", toStruct);
            statements.add(strct);
            DBSPVariablePath var = DBSPTypeAny.getDefault().var();
            DBSPExpression into = new DBSPApplyMethodExpression("into", DBSPTypeAny.getDefault(), var).closure(var);
            DBSPExpression retval = new DBSPApplyMethodExpression("map", returnType, strct.getVarReference(), into);
            DBSPExpression body = new DBSPBlockExpression(statements, retval);
            // DBSPExpression ok = new DBSPApplyMethodExpression("Ok", new DBSPTypeResult(returnType), retval);
            // TODO: the function should return Result
            return new DBSPFunction(this.node, this.getName(), Linq.list(param), returnType, body, Linq.list());
        }
        throw new InternalCompilerError("Cannot generate implementation for " + this, this.node);
    }
}
